diff --git a/CHANGES.md b/CHANGES.md index cc1ec48ba188..5ea1c49c09e2 100644 --- a/CHANGES.md +++ b/CHANGES.md @@ -78,6 +78,7 @@ ## Breaking Changes +* (Python) Made Beartype the default fallback type checking tool. This can be disabled with the `--disable_beartype` pipeline option. ([#38275](https://github.com/apache/beam/issues/38275)) * X behavior was changed ([#X](https://github.com/apache/beam/issues/X)). ## Deprecations diff --git a/sdks/python/apache_beam/options/pipeline_options.py b/sdks/python/apache_beam/options/pipeline_options.py index d60d75283eab..163729d24dd7 100644 --- a/sdks/python/apache_beam/options/pipeline_options.py +++ b/sdks/python/apache_beam/options/pipeline_options.py @@ -886,6 +886,11 @@ def _add_argparse_args(cls, parser): action='store_false', help='Disable type checking at pipeline construction ' 'time') + parser.add_argument( + '--disable_beartype', + default=False, + action='store_true', + help='Disable the use of beartype for type checking.') parser.add_argument( '--runtime_type_check', default=False, diff --git a/sdks/python/apache_beam/typehints/typehints.py b/sdks/python/apache_beam/typehints/typehints.py index f429935c3a0e..bfd74ef7cd1a 100644 --- a/sdks/python/apache_beam/typehints/typehints.py +++ b/sdks/python/apache_beam/typehints/typehints.py @@ -1485,7 +1485,8 @@ def normalize(x, none_as_type=False): }) -def is_consistent_with(sub, base, use_beartype: bool = False) -> bool: +def is_consistent_with( + sub, base, use_beartype: typing.Optional[bool] = None) -> bool: """Checks whether sub a is consistent with base. This is according to the terminology of PEP 483/484. This relationship is @@ -1494,6 +1495,15 @@ def is_consistent_with(sub, base, use_beartype: bool = False) -> bool: relation, but also handles the special Any type as well as type parameterization. """ + if use_beartype is None: + from apache_beam.options.pipeline_options_context import get_pipeline_options + options = get_pipeline_options() + if options: + from apache_beam.options.pipeline_options import TypeOptions + use_beartype = not options.view_as(TypeOptions).disable_beartype + else: + use_beartype = True + from apache_beam.pvalue import Row from apache_beam.typehints.row_type import RowTypeConstraint if sub == base: diff --git a/sdks/python/apache_beam/typehints/typehints_test.py b/sdks/python/apache_beam/typehints/typehints_test.py index 992c129fd8a5..3806f65ac429 100644 --- a/sdks/python/apache_beam/typehints/typehints_test.py +++ b/sdks/python/apache_beam/typehints/typehints_test.py @@ -1613,6 +1613,33 @@ def test_hint_helper_pipe_union(self): self.assertTrue(is_consistent_with(int, pipe_union_2)) self.assertTrue(is_consistent_with(float, pipe_union_2)) + def test_is_consistent_with_disable_beartype(self): + import unittest.mock + + from apache_beam.options.pipeline_options import PipelineOptions + from apache_beam.options.pipeline_options_context import scoped_pipeline_options + + with unittest.mock.patch( + 'apache_beam.typehints.typehints.is_subhint') as mock_is_subhint: + mock_is_subhint.return_value = True + + class A: + pass + + class B(A): + pass + + options = PipelineOptions([]) + with scoped_pipeline_options(options): + typehints.is_consistent_with(B, A) + self.assertTrue(mock_is_subhint.called) + mock_is_subhint.reset_mock() + + options = PipelineOptions(['--disable_beartype']) + with scoped_pipeline_options(options): + typehints.is_consistent_with(B, A) + self.assertFalse(mock_is_subhint.called) + def test_positional_arg_hints(self): self.assertEqual(typehints.Any, _positional_arg_hints('x', {})) self.assertEqual(int, _positional_arg_hints('x', {'x': int}))