saavannanavati commented on a change in pull request #12009:
URL: https://github.com/apache/beam/pull/12009#discussion_r451733706
##########
File path: sdks/python/apache_beam/typehints/decorators.py
##########
@@ -378,6 +378,56 @@ 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):
+ return self.strip_pcoll_helper(self.input_types,
+ self._has_input_types,
+ {'input_types': None},
+ ['apache_beam.pvalue.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': None},
+ ['apache_beam.pvalue.PDone'],
+ 'An output typehint to a PTransform must be'
+ ' a single (or nested) type wrapped by '
+ 'a PCollection or PDone. ',
+ 'strip_pcoll_output()')
+
+ def strip_pcoll_helper(
+ self,
+ my_type, # type: any
+ has_my_type, # type: Callable[[], bool]
+ kwarg_dict, # type: Dict[str, any]
+ my_valid_classes, # type: List[str]
+ error_str, # type: str
+ source_str # type: str
+ ):
+ # type: (...) -> IOTypeHints
+
+ if not has_my_type() or len(my_type[0]) != 1:
+ return self
+
+ my_type = my_type[0][0]
+
+ if isinstance(my_type, typehints.AnyTypeConstraint):
+ return self
+
+ valid_classes = ['apache_beam.pvalue.PCollection'] + my_valid_classes
+
+ if not any(valid_class in str(my_type) for valid_class in valid_classes):
Review comment:
I tried using `isinstance` initially but it doesn't work well with
generic types
Another option is to use `__origin__` but I don't know if that's fully
backwards compatible
Strings are a wacky solution though.. do you have any other ideas?
----------------------------------------------------------------
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:
us...@infra.apache.org
