From e96cf241fe12838cce3d8822ff71cee9c0369a12 Mon Sep 17 00:00:00 2001
From: Yi Hu <yathu@google.com>
Date: Tue, 1 Sep 2026 13:05:59 -0400
Subject: [PATCH] Always resolve builtin types in dataclass field type
 inference

exclude_infer_dataclass_field_type was introduced to temporarily keeping old behavior of typehint resolution of dataclass field type.
We now aim to gradualy phase out this restriction. First, expand the types always enabled to all Python built-in types, not limited
to primitive types, and Optional[...] hints as well
---
 .../typehints/trivial_inference.py            | 40 +++++++++++++++++--
 .../typehints/trivial_inference_test.py       | 39 +++++++++++++++---
 2 files changed, 69 insertions(+), 10 deletions(-)

diff --git a/sdks/python/apache_beam/typehints/trivial_inference.py b/sdks/python/apache_beam/typehints/trivial_inference.py
index 69edfc309281..ec6bf6b1c8ef 100644
--- a/sdks/python/apache_beam/typehints/trivial_inference.py
+++ b/sdks/python/apache_beam/typehints/trivial_inference.py
@@ -776,6 +776,34 @@ def infer_return_type_func(f, input_types, debug=False, depth=0):
   return result
 
 
+_BUILTIN_TYPES = (
+    bool,
+    bytearray,
+    bytes,
+    complex,
+    float,
+    int,
+    str,
+    dict,
+    frozenset,
+    list,
+    set,
+    tuple,
+)
+_CONTAINER_CONSTRAINTS = (
+    typehints.ListConstraint,
+    typehints.DictConstraint,
+    typehints.SetTypeConstraint,
+    typehints.FrozenSetTypeConstraint,
+    typehints.TupleConstraint,
+    typehints.TupleSequenceConstraint,
+)
+
+
+def _is_builtin_type(t):
+  return t in _BUILTIN_TYPES or isinstance(t, _CONTAINER_CONSTRAINTS)
+
+
 def resolve_dataclass_field_type(x):
   """
   Resolve a type to Beam typehint under global pipeline option context.
@@ -785,7 +813,7 @@ def resolve_dataclass_field_type(x):
   incorrect typehints; non-deterministic or nullable types disallowed by
   consumer transform but check disabled by Any; tests rely on Any),
   --exclude_infer_dataclass_field_type option to instruct falling back to Any.
-  Fields of builtin primitives are always respected.
+  Fields of builtin types and their Optional are always respected.
   """
   from apache_beam.options.pipeline_options_context import get_pipeline_options
   options = get_pipeline_options()
@@ -795,8 +823,12 @@ def resolve_dataclass_field_type(x):
   else:
     disabled = False
 
+  norm_x = typehints.normalize(x)
   if not disabled:
-    return typehints.normalize(x)
-  if x in (bool, bytes, complex, float, int, str):
-    return x
+    return norm_x
+  if _is_builtin_type(norm_x):
+    return norm_x
+  if (typehints.is_nullable(norm_x) and
+      _is_builtin_type(typehints.get_concrete_type_from_nullable(norm_x))):
+    return norm_x
   return Any
diff --git a/sdks/python/apache_beam/typehints/trivial_inference_test.py b/sdks/python/apache_beam/typehints/trivial_inference_test.py
index dcb0bac97e80..648f073776f8 100644
--- a/sdks/python/apache_beam/typehints/trivial_inference_test.py
+++ b/sdks/python/apache_beam/typehints/trivial_inference_test.py
@@ -21,6 +21,7 @@
 
 import dataclasses
 import types
+import typing
 import unittest
 
 import apache_beam as beam
@@ -501,18 +502,44 @@ class MyDataClass:
       name: str
       tags: list[str]
       custom: BaseClass
-
-    self.assertReturnType(
-        typehints.Tuple[int, str, typehints.List[str], BaseClass],
+      opt_id: typing.Optional[int]
+      opt_custom: typing.Optional[BaseClass]
+      mapping: dict[str, int]
+      coord: tuple[float, float]
+      categories: set[str]
+      immutable: frozenset[int]
+
+    self.assertReturnType(
+        typehints.Tuple[int,
+                        str,
+                        typehints.List[str],
+                        BaseClass,
+                        typehints.Optional[int],
+                        typehints.Optional[BaseClass],
+                        typehints.Dict[str, int],
+                        typehints.Tuple[float, float],
+                        typehints.Set[str],
+                        typehints.FrozenSet[int]],
         python_callable.PythonCallableWithSource(
-            "lambda x: (x.id, x.name, x.tags, x.custom)"), [MyDataClass])
+            "lambda x: (x.id, x.name, x.tags, x.custom, x.opt_id, x.opt_custom, "
+            "x.mapping, x.coord, x.categories, x.immutable)"), [MyDataClass])
 
     options = PipelineOptions(['--exclude_infer_dataclass_field_type'])
     with scoped_pipeline_options(options):
       self.assertReturnType(
-          typehints.Tuple[int, str, typehints.Any, typehints.Any],
+          typehints.Tuple[int,
+                          str,
+                          typehints.List[str],
+                          typehints.Any,
+                          typehints.Optional[int],
+                          typehints.Any,
+                          typehints.Dict[str, int],
+                          typehints.Tuple[float, float],
+                          typehints.Set[str],
+                          typehints.FrozenSet[int]],
           python_callable.PythonCallableWithSource(
-              "lambda x: (x.id, x.name, x.tags, x.custom)"), [MyDataClass])
+              "lambda x: (x.id, x.name, x.tags, x.custom, x.opt_id, x.opt_custom, "
+              "x.mapping, x.coord, x.categories, x.immutable)"), [MyDataClass])
 
 
 if __name__ == '__main__':