diff --git a/sdks/python/apache_beam/transforms/core.py b/sdks/python/apache_beam/transforms/core.py index 74773a4d7caf..ea11bca9474d 100644 --- a/sdks/python/apache_beam/transforms/core.py +++ b/sdks/python/apache_beam/transforms/core.py @@ -1515,7 +1515,8 @@ def _check_fn_use_yield_and_return(fn): "yield("): has_yield = True elif lstripped_line.rstrip() == "return": - has_return = True + # Return is likely used to exit the function - ok to use with 'yield'. + pass elif lstripped_line.startswith("return ") or lstripped_line.startswith( "return("): if lstripped_line.rstrip() == "return None" or lstripped_line.rstrip( diff --git a/sdks/python/apache_beam/transforms/core_test.py b/sdks/python/apache_beam/transforms/core_test.py index 80ab6a88afb4..73f004c130c2 100644 --- a/sdks/python/apache_beam/transforms/core_test.py +++ b/sdks/python/apache_beam/transforms/core_test.py @@ -43,6 +43,14 @@ RETURN_NONE_PARTIAL_WARNING = "Process method returned None" +class TestDoFn0(beam.DoFn): + """Returning without a value is allowed""" + def process(self, element): + if not element: + return + yield element + + class TestDoFn1(beam.DoFn): def process(self, element): yield element @@ -174,6 +182,7 @@ def test_dofn_with_yield_and_return(self): with self._caplog.at_level(logging.WARNING): assert beam.ParDo(sum) + assert beam.ParDo(TestDoFn0()) assert beam.ParDo(TestDoFn1()) assert beam.ParDo(TestDoFn2()) assert beam.ParDo(TestDoFn4())