Mazen050 commented on code in PR #40501:
URL: https://github.com/apache/beam/pull/40501#discussion_r4368339859
##########
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)
+
+ result = self._type_check_result(self.dofn.process(*args, **kwargs))
+ except TypeCheckError as e:
+ raise self._add_label(e)
+ else:
Review Comment:
The same issue applies here: exceptions raised in `else` bypass the
preceding `except`. Since `type_check_output()` doesn't call `_add_label()`
itself, I couldn't reproduce the double-labeling rationale in my current
implementation.
--
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]
