diff --git a/sdks/python/apache_beam/typehints/typecheck.py b/sdks/python/apache_beam/typehints/typecheck.py index 7e84779a0f15..46feeaf03b9c 100644 --- a/sdks/python/apache_beam/typehints/typecheck.py +++ b/sdks/python/apache_beam/typehints/typecheck.py @@ -84,25 +84,81 @@ def teardown(self): return self.dofn.teardown() -class OutputCheckWrapperDoFn(AbstractDoFnWrapper): - """A DoFn that verifies against common errors in the output type.""" - def __init__(self, dofn, full_label): +class RuntimeTypeCheckWrapperDoFn(AbstractDoFnWrapper): + """A DoFn wrapper performing runtime type-checking of inputs and outputs. + + This single wrapper performs the work that was previously split between + two nested wrappers (``TypeCheckWrapperDoFn`` wrapped inside + ``OutputCheckWrapperDoFn``): type-checking of inputs and outputs against + declared type hints, verification against common output errors (such as + returning a plain ``str``, ``bytes`` or ``dict`` from a ``ParDo``), and + labeling of raised ``TypeCheckError`` messages with the offending + transform's full label. Merging the wrappers removes one level of Python + call indirection per processed element when ``--runtime_type_check`` is + enabled. + """ + def __init__(self, dofn, type_hints, full_label): super().__init__(dofn) self.full_label = full_label + # Note that a *bound* process method must not be cached on the instance: + # an attribute holding a bound method is visible to the stateful DoFn + # reflection in userstate.get_dofn_specs, and a cached copy can diverge + # from self.dofn.process across (de)serialization, in which case stateful + # DoFn validation would see duplicate StateSpecs/TimerSpecs. Caching the + # underlying (unbound) function is safe: plain functions stored on an + # instance are not bound methods, so DoFn reflection ignores them. + process_fn = dofn._process_argspec_fn() + if hasattr(process_fn, '__func__'): + self._process_fn = process_fn.__func__ + self._process_fn_needs_self = True + else: + self._process_fn = process_fn + self._process_fn_needs_self = False + if type_hints.input_types: + input_args, input_kwargs = type_hints.input_types + self._input_hints = getcallargs_forhints( + process_fn, *input_args, **input_kwargs) + else: + self._input_hints = None + # TODO(robertwb): Multi-output. + self._output_type_hint = type_hints.simple_output_type(full_label) def wrapper(self, method, args, kwargs): try: - result = method(*args, **kwargs) + result = self._type_check_result(method(*args, **kwargs)) except TypeCheckError as e: - # TODO(BEAM-10710): Remove the 'ParDo' prefix for the label name - error_msg = ( - 'Runtime type violation detected within ParDo(%s): ' - '%s' % (self.full_label, e)) - _, _, tb = sys.exc_info() - raise TypeCheckError(error_msg).with_traceback(tb) + raise self._add_label(e) + else: + return self._check_type(result) + + def process(self, *args, **kwargs): + try: + if self._input_hints: + if self._process_fn_needs_self: + actual_inputs = inspect.getcallargs( + self._process_fn, self.dofn, *args, **kwargs) # pylint: disable=deprecated-method + else: + actual_inputs = inspect.getcallargs(self._process_fn, *args, **kwargs) # pylint: disable=deprecated-method + for var, hint in self._input_hints.items(): + if hint is actual_inputs[var]: + # self parameter + continue + _check_instance_type(hint, actual_inputs[var], var, True) + result = self._type_check_result(self.dofn.process(*args, **kwargs)) + except TypeCheckError as e: + raise self._add_label(e) else: return self._check_type(result) + def _add_label(self, e): + """Returns a TypeCheckError labeled with this transform's full label.""" + # TODO(BEAM-10710): Remove the 'ParDo' prefix for the label name + error_msg = ( + 'Runtime type violation detected within ParDo(%s): ' + '%s' % (self.full_label, e)) + _, _, tb = sys.exc_info() + return TypeCheckError(error_msg).with_traceback(tb) + @staticmethod def _check_type(output): if output is None: @@ -120,36 +176,6 @@ def _check_type(output): 'iterable. %s was returned instead.' % type(output)) return output - -class TypeCheckWrapperDoFn(AbstractDoFnWrapper): - """A wrapper around a DoFn which performs type-checking of input and output. - """ - def __init__(self, dofn, type_hints, label=None): - super().__init__(dofn) - self._process_fn = self.dofn._process_argspec_fn() - if type_hints.input_types: - input_args, input_kwargs = type_hints.input_types - self._input_hints = getcallargs_forhints( - self._process_fn, *input_args, **input_kwargs) - else: - self._input_hints = None - # TODO(robertwb): Multi-output. - self._output_type_hint = type_hints.simple_output_type(label) - - def wrapper(self, method, args, kwargs): - result = method(*args, **kwargs) - return self._type_check_result(result) - - def process(self, *args, **kwargs): - if self._input_hints: - actual_inputs = inspect.getcallargs(self._process_fn, *args, **kwargs) # pylint: disable=deprecated-method - for var, hint in self._input_hints.items(): - if hint is actual_inputs[var]: - # self parameter - continue - _check_instance_type(hint, actual_inputs[var], var, True) - return self._type_check_result(self.dofn.process(*args, **kwargs)) - def _type_check_result(self, transform_results): if self._output_type_hint is None or transform_results is None: return transform_results @@ -289,11 +315,9 @@ def visit_transform(self, applied_transform): if isinstance(transform.fn, core.CombineValuesDoFn): transform.fn.combinefn = self._wrapped_fn else: - transform.fn = transform.dofn = OutputCheckWrapperDoFn( - TypeCheckWrapperDoFn( - transform.fn, - transform.get_type_hints(), - applied_transform.full_label), + transform.fn = transform.dofn = RuntimeTypeCheckWrapperDoFn( + transform.fn, + transform.get_type_hints(), applied_transform.full_label) diff --git a/sdks/python/apache_beam/typehints/typecheck_test.py b/sdks/python/apache_beam/typehints/typecheck_test.py index 81046cd3cc0f..ceddc4458a6c 100644 --- a/sdks/python/apache_beam/typehints/typecheck_test.py +++ b/sdks/python/apache_beam/typehints/typecheck_test.py @@ -34,10 +34,14 @@ from apache_beam import Pipeline from apache_beam.options.pipeline_options import PipelineOptions from apache_beam.options.pipeline_options import TypeOptions +from apache_beam.pipeline import PipelineVisitor from apache_beam.testing.test_pipeline import TestPipeline from apache_beam.testing.util import assert_that from apache_beam.testing.util import equal_to +from apache_beam.transforms import combiners +from apache_beam.transforms import userstate from apache_beam.typehints import decorators +from apache_beam.typehints import typecheck from apache_beam.typehints import with_input_types from apache_beam.typehints import with_output_types @@ -105,7 +109,7 @@ def test_wrapper_pass_through(self): # not the same one that actually runs in the pipeline (it is serialized # here and deserialized in the worker). with tempfile.TemporaryDirectory() as tmp_dirname: - path = os.path.join(tmp_dirname + "tmp_filename") + path = os.path.join(tmp_dirname, "tmp_filename") dofn = MyDoFn(path) result = self.p | beam.Create([1, 2, 3]) | beam.ParDo(dofn) assert_that(result, equal_to([1, 2, 3])) @@ -300,5 +304,98 @@ def process(self, element, *args, **kwargs): self.p.run().wait_until_finish() +@with_input_types(int) +@with_output_types(int) +class _AddOneDoFn(beam.DoFn): + def process(self, element): + yield element + 1 + + +class RuntimeTypeCheckWrapperDoFnTest(unittest.TestCase): + """Tests for the merged runtime type-checking wrapper (BEAM-9489).""" + def test_visitor_applies_single_wrapper_layer(self): + # A single RuntimeTypeCheckWrapperDoFn should wrap the user's DoFn + # directly, with no intermediate wrapper layers. + p = beam.Pipeline(options=PipelineOptions(runtime_type_check=True)) + _ = (p | beam.Create([1, 2]) | 'TypedStep' >> beam.ParDo(_AddOneDoFn())) + p.visit(typecheck.TypeCheckVisitor()) + + wrapped_dofns = [] + + class _Collector(PipelineVisitor): + def visit_transform(self, applied_transform): + transform = applied_transform.transform + if isinstance(transform, beam.ParDo) and isinstance( + getattr(transform, 'fn', None), typecheck.AbstractDoFnWrapper): + wrapped_dofns.append(transform.fn) + + p.visit(_Collector()) + self.assertTrue(wrapped_dofns) + for wrapper in wrapped_dofns: + self.assertIsInstance(wrapper, typecheck.RuntimeTypeCheckWrapperDoFn) + # The wrapped DoFn must be the user's DoFn, not another wrapper. + self.assertNotIsInstance(wrapper.dofn, typecheck.AbstractDoFnWrapper) + + def test_wrapper_labels_type_check_errors(self): + dofn = typecheck.RuntimeTypeCheckWrapperDoFn( + _AddOneDoFn(), _AddOneDoFn().get_type_hints(), 'MyStep') + with self.assertRaisesRegex( + typecheck.TypeCheckError, + r'Runtime type violation detected within ParDo\(MyStep\)'): + dofn.process('not-an-int') + + def test_wrapper_preserves_results(self): + dofn = typecheck.RuntimeTypeCheckWrapperDoFn( + _AddOneDoFn(), _AddOneDoFn().get_type_hints(), 'MyStep') + self.assertEqual(list(dofn.process(1)), [2]) + + +def _make_stateful_dofn(): + # Defined inside a function so that the class is serialized by value + # (like the DoFn created by GroupIntoBatches), which is the case where a + # cached bound method diverges from the reconstructed class's process. + count_state = userstate.CombiningValueStateSpec( + 'count', combiners.CountCombineFn()) + + class _CountingStatefulDoFn(beam.DoFn): + def process(self, element, count=beam.DoFn.StateParam(count_state)): + count.add(1) + yield element[1] + + return _CountingStatefulDoFn() + + +class RuntimeTypeCheckStatefulDoFnTest(unittest.TestCase): + """Regression tests for runtime_type_check with stateful DoFns. + + The type-check wrapper must not expose duplicate StateSpecs/TimerSpecs to + stateful DoFn validation, including after the runner serializes and + reconstructs the wrapped DoFn. + """ + def test_wrapper_does_not_cache_bound_process(self): + dofn = _make_stateful_dofn() + wrapper = typecheck.RuntimeTypeCheckWrapperDoFn( + dofn, dofn.get_type_hints(), 'Step') + # A bound method cached in the instance __dict__ is visible to + # userstate.get_dofn_specs and can diverge from self.dofn.process + # across (de)serialization. + for attr, value in wrapper.__dict__.items(): + self.assertFalse( + hasattr(value, '__self__'), + 'wrapper caches bound method %r, which breaks stateful DoFn ' + 'validation' % attr) + userstate.validate_stateful_dofn(wrapper) + + def test_stateful_dofn_with_runtime_type_check(self): + options = PipelineOptions() + options.view_as(TypeOptions).runtime_type_check = True + with TestPipeline(options=options) as p: + result = ( + p + | beam.Create([('k', 1), ('k', 2), ('k', 3)]) + | beam.ParDo(_make_stateful_dofn())) + assert_that(result, equal_to([1, 2, 3])) + + if __name__ == '__main__': unittest.main()