udim commented on a change in pull request #12009:
URL: https://github.com/apache/beam/pull/12009#discussion_r461287903
##########
File path: sdks/python/apache_beam/typehints/decorators.py
##########
@@ -378,6 +379,67 @@ def has_simple_output_type(self):
self.output_types and len(self.output_types[0]) == 1 and
not self.output_types[1])
+ def strip_pcoll(self):
+ from apache_beam.pvalue import PBegin
+ from apache_beam.pvalue import PDone
+
+ return self.strip_pcoll_helper(self.input_types,
+ self._has_input_types,
+ 'input_types',
+ [PBegin],
+ 'An input typehint to a PTransform must be '
+ 'a single (or nested) type wrapped by '
+ 'a PCollection or PBegin. ',
+ 'strip_pcoll_input()').\
+ strip_pcoll_helper(self.output_types,
+ self.has_simple_output_type,
+ 'output_types',
+ [PDone, None],
+ 'An output typehint to a PTransform must be
'
+ 'a single (or nested) type wrapped by '
+ 'a PCollection, PDone, or None. ',
+ 'strip_pcoll_output()')
+
+ def strip_pcoll_helper(
+ self,
+ my_type, # type: any
+ has_my_type, # type: Callable[[], bool]
+ my_key, # type: str
+ special_containers, # type: List[Union[PBegin, PDone, PCollection]]
+ error_str, # type: str
+ source_str # type: str
+ ):
+ # type: (...) -> IOTypeHints
+ from apache_beam.pvalue import PCollection
+
+ if not has_my_type() or not my_type or len(my_type[0]) != 1:
+ return self
+
+ my_type = my_type[0][0]
+
+ if isinstance(my_type, typehints.AnyTypeConstraint):
+ return self
+
+ special_containers += [PCollection]
+
+ if (my_type not in special_containers and
+ getattr(my_type, '__origin__', None) != PCollection):
+ raise TypeCheckError(error_str)
Review comment:
Nit: This error message doesn't say what the incorrect type was. Ex:
```suggestion
raise TypeCheckError(error_str + ' Got: %s' % my_type)
```
----------------------------------------------------------------
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.
For queries about this service, please contact Infrastructure at:
[email protected]
