[BEAM-9296] Clean up and add type-hints to SDF API#10935
Conversation
|
cc: @robertwb |
chadrik
left a comment
There was a problem hiding this comment.
Looks good. Just two minor notes. Thanks for following up on this!
|
|
||
| if TYPE_CHECKING: | ||
| from apache_beam.io.iobase import RestrictionTracker | ||
| from apache_beam.io.iobase import RestrictionProgress |
There was a problem hiding this comment.
Our linters do not enforce alphanumeric order for modules inside the TYPE_CHECKING block, but we should still be diligent about it. Can you move this up one, please.
There was a problem hiding this comment.
Done. Thanks for mentioning that!
| return self.get_restriction_provider() is not None | ||
|
|
||
| def get_restriction_coder(self): | ||
| # type: () -> Optional[TupleCoder] |
There was a problem hiding this comment.
A slight refactor here will avoid introducing a new mypy error:
def get_restriction_coder(self):
# type: () -> Optional[TupleCoder]
"""Get coder for a restriction when processing an SDF. """
if self.is_splittable_dofn():
return TupleCoder([
(self.get_restriction_provider().restriction_coder()),
(self.get_watermark_estimator_provider().estimator_state_coder())
])
else:
return NoneThis avoids having to declare the restriction_coder variable as Optional[TupleCoder].
There was a problem hiding this comment.
Done. I think the return type should still be Optional[TupleCoder] given that it also returns None.
There was a problem hiding this comment.
correct. my comment about avoiding the declaration of Optional[TupleCoder] refers to the variable (which my edit does away with), not the the return type.
|
LGTM, assuming the tests pass. |
|
All tests passed. I'm going to merge it. Thanks, Chad! |
R: @chadrik
Thank you for your contribution! Follow this checklist to help us incorporate your contribution quickly and easily:
R: @username).[BEAM-XXX] Fixes bug in ApproximateQuantiles, where you replaceBEAM-XXXwith the appropriate JIRA issue, if applicable. This will automatically link the pull request to the issue.CHANGES.mdwith noteworthy changes.See the Contributor Guide for more tips on how to make review process smoother.
Post-Commit Tests Status (on master branch)
Pre-Commit Tests Status (on master branch)
See .test-infra/jenkins/README for trigger phrase, status and link of all Jenkins jobs.