(beam) branch master updated: Support infer types involving dataclass fields (#38548)

[yhu]

This is an automated email from the ASF dual-hosted git repository.

Abacn pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/beam.git


The following commit(s) were added to refs/heads/master by this push:
     new 3e5a085bd0d Support infer types involving dataclass fields (#38548)
3e5a085bd0d is described below

commit 3e5a085bd0da62b706726ba812d0530c197ffe8a
Author: Yi Hu <ya...@google.com>
AuthorDate: Thu May 28 15:42:08 2026 -0400

    Support infer types involving dataclass fields (#38548)
---
 sdks/python/apache_beam/typehints/opcodes.py        |  6 ++++++
 sdks/python/apache_beam/typehints/row_type_test.py  | 21 +++++++++++++++++++++
 .../apache_beam/typehints/trivial_inference_test.py | 12 ++++++++++++
 3 files changed, 39 insertions(+)

diff --git a/sdks/python/apache_beam/typehints/opcodes.py 
b/sdks/python/apache_beam/typehints/opcodes.py
index 8e5d7b1e40c..963b5e0850b 100644
--- a/sdks/python/apache_beam/typehints/opcodes.py
+++ b/sdks/python/apache_beam/typehints/opcodes.py
@@ -29,6 +29,7 @@ For internal use only; no backwards-compatibility guarantees.
 """
 # pytype: skip-file
 
+import dataclasses
 import inspect
 import logging
 import sys
@@ -447,6 +448,11 @@ def _getattr(o, name):
     return Const(BoundMethod(func, o))
   elif isinstance(o, row_type.RowTypeConstraint):
     return o.get_type_for(name)
+  elif inspect.isclass(o) and dataclasses.is_dataclass(o):
+    field = o.__dataclass_fields__.get(name)
+    if field is not None:
+      return field.type
+    return Any
   else:
     return Any
 
diff --git a/sdks/python/apache_beam/typehints/row_type_test.py 
b/sdks/python/apache_beam/typehints/row_type_test.py
index 54e64caf6fa..30bda0cd98b 100644
--- a/sdks/python/apache_beam/typehints/row_type_test.py
+++ b/sdks/python/apache_beam/typehints/row_type_test.py
@@ -172,6 +172,27 @@ class RowTypeTest(unittest.TestCase):
         getattr(DerivedDataClass, row_type._BEAM_SCHEMA_ID))
     self.assertNotEqual(schema_for_derived.id, schema_for_base.id)
 
+  def test_dataclass_map_typehints(self):
+    @beam.coders.typecoders.registry.register_row
+    @dataclass(frozen=True)
+    class MyDataClass:
+      id: int
+      name: str
+
+    p = beam.Pipeline()
+    pa = (p | beam.Create([MyDataClass(1, "a"), MyDataClass(2, "b")]))
+    self.assertEqual(pa.element_type, MyDataClass)
+
+    pb = (
+        pa | beam.Map(
+            lambda x: beam.Row(id=x.id, name=x.name, name_hash=hash(x.name))))
+    self.assertTrue(
+        isinstance(pb.element_type, row_type.GeneratedClassRowTypeConstraint))
+    self.assertEqual(
+        pb.element_type,
+        row_type.GeneratedClassRowTypeConstraint(
+            fields=[('id', int), ('name', str), ('name_hash', int)]))
+
 
 if __name__ == '__main__':
   unittest.main()
diff --git a/sdks/python/apache_beam/typehints/trivial_inference_test.py 
b/sdks/python/apache_beam/typehints/trivial_inference_test.py
index fe60974e280..f421819bdca 100644
--- a/sdks/python/apache_beam/typehints/trivial_inference_test.py
+++ b/sdks/python/apache_beam/typehints/trivial_inference_test.py
@@ -19,6 +19,7 @@
 
 # pytype: skip-file
 
+import dataclasses
 import types
 import unittest
 
@@ -487,6 +488,17 @@ class TrivialInferenceTest(unittest.TestCase):
         python_callable.PythonCallableWithSource("lambda x: (x, str(x))"),
         [int])
 
+  def testDataClassFields(self):
+    @dataclasses.dataclass
+    class MyDataClass:
+      id: int
+      name: str
+
+    self.assertReturnType(
+        typehints.Tuple[int, str],
+        python_callable.PythonCallableWithSource("lambda x: (x.id, x.name)"),
+        [MyDataClass])
+
 
 if __name__ == '__main__':
   unittest.main()