Mazen050 commented on code in PR #40501:
URL: https://github.com/apache/beam/pull/40501#discussion_r4368337376
##########
sdks/python/apache_beam/typehints/typecheck.py:
##########
@@ -84,25 +84,67 @@ 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 wrapper around a DoFn which performs type-checking of input and
output."""
+ def __init__(self, dofn, type_hints, full_label):
super().__init__(dofn)
self.full_label = full_label
+ process_fn = self.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)
+
Review Comment:
I tested this change with the regression test I just pushed. It fails
because the `TypeCheckError` raised in `else` bypasses the preceding `except`,
so the `ParDo` label is lost. I'll keep `_type_check_result()` inside `try`.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
