github-actions[bot] commented on code in PR #40501:
URL: https://github.com/apache/beam/pull/40501#discussion_r4310010673


##########
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:
   **🟠 High**: `self._type_check_result(...)` is called inside the `try ... 
except TypeCheckError as e: raise self._add_label(e)` block in both `process()` 
(line 133) and `wrapper()` (line 114), while `type_check_output()` inside 
`_type_check_result()` (lines 174-177) also catches `TypeCheckError` and raises 
`self._add_label(e)`.
   
   When `process()` or `wrapper()` returns a non-generator iterable (e.g., a 
`list` or `tuple`, which includes `beam.Map` and list-returning 
`FlatMap`/`ParDo` transforms), `for o in transform_results: 
type_check_output(o)` runs synchronously inside `_type_check_result`. Any 
output `TypeCheckError` is wrapped by `self._add_label(e)` inside 
`type_check_output` and then caught and wrapped by `self._add_label(e)` a 
second time in `process()`/`wrapper()`, producing a duplicated prefix:
   `Runtime type violation detected within ParDo(<label>): Runtime type 
violation detected within ParDo(<label>): ...`
   
   Moving `self._type_check_result(result)` into the `else:` branch (in both 
`process` and `wrapper`) ensures `self._add_label(e)` is applied only once for 
both generator and non-generator outputs.
   
   ```suggestion
         result = self.dofn.process(*args, **kwargs)
   ```



##########
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:
   **🟠 High**: Pair with moving `self._type_check_result` out of the `try` 
block on line 133 (and similarly in `wrapper` on lines 114 and 118) so output 
`TypeCheckError`s are not labeled twice.
   
   ```suggestion
         return self._check_type(self._type_check_result(result))
   ```



-- 
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]

Reply via email to