From d077e74fad46fcd995300c8864b7be9c4f4c3781 Mon Sep 17 00:00:00 2001 From: Shunping Huang Date: Tue, 17 Dec 2024 14:34:04 -0500 Subject: [PATCH 1/2] Add missing to_type_hint to WindowedValueCoder --- sdks/python/apache_beam/coders/coders.py | 3 +++ sdks/python/apache_beam/coders/coders_test.py | 6 ++++++ 2 files changed, 9 insertions(+) diff --git a/sdks/python/apache_beam/coders/coders.py b/sdks/python/apache_beam/coders/coders.py index e52c6048a15c..22d041f34f8b 100644 --- a/sdks/python/apache_beam/coders/coders.py +++ b/sdks/python/apache_beam/coders/coders.py @@ -1446,6 +1446,9 @@ def from_type_hint(cls, typehint, registry): # pickle coders. return cls(registry.get_coder(typehint.inner_type)) + def to_type_hint(self): + return typehints.WindowedValue[self.wrapped_value_coder.to_type_hint()] + Coder.register_structured_urn( common_urns.coders.WINDOWED_VALUE.urn, WindowedValueCoder) diff --git a/sdks/python/apache_beam/coders/coders_test.py b/sdks/python/apache_beam/coders/coders_test.py index dc9780e36be3..c63281f6d5c4 100644 --- a/sdks/python/apache_beam/coders/coders_test.py +++ b/sdks/python/apache_beam/coders/coders_test.py @@ -258,6 +258,12 @@ def test_numpy_int(self): _ = indata | "CombinePerKey" >> beam.CombinePerKey(sum) +class WindowedValueCoderTest(unittest.TestCase): + def test_to_type_hint(self): + coder = coders.WindowedValueCoder(coders.VarIntCoder()) + self.assertEqual(coder.to_type_hint(), typehints.WindowedValue[int]) + + if __name__ == '__main__': logging.getLogger().setLevel(logging.INFO) unittest.main() From 409599d36f128a874fcaf990def0eeac06b403c4 Mon Sep 17 00:00:00 2001 From: Shunping Huang Date: Tue, 17 Dec 2024 15:11:31 -0500 Subject: [PATCH 2/2] Add type ignore to make mypy happy. --- sdks/python/apache_beam/coders/coders_test.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/sdks/python/apache_beam/coders/coders_test.py b/sdks/python/apache_beam/coders/coders_test.py index c63281f6d5c4..bddd2cb57e06 100644 --- a/sdks/python/apache_beam/coders/coders_test.py +++ b/sdks/python/apache_beam/coders/coders_test.py @@ -261,7 +261,7 @@ def test_numpy_int(self): class WindowedValueCoderTest(unittest.TestCase): def test_to_type_hint(self): coder = coders.WindowedValueCoder(coders.VarIntCoder()) - self.assertEqual(coder.to_type_hint(), typehints.WindowedValue[int]) + self.assertEqual(coder.to_type_hint(), typehints.WindowedValue[int]) # type: ignore[misc] if __name__ == '__main__':