Allow yield and return in the same 'process' if return doesn't have a value. (#37190)

Author: tvalentyn

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(

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())