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