Always resolve builtin types in dataclass field type inference

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

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__':