address comments - make raw return a valid return

Author: Abacn

sdks/python/apache_beam/transforms/core.py

@@ -1514,19 +1514,20 @@ def _check_fn_use_yield_and_return(fn):
           "yield("):
         has_yield = True
       elif lstripped_line.rstrip() == "return":
-        has_return_none = True
+        has_return = True
       elif lstripped_line.startswith("return ") or lstripped_line.startswith(
           "return("):
         if lstripped_line.startswith("return None"):
           has_return_none = True
         else:
           has_return = True
-      if has_yield and has_return:
+      if has_yield and (has_return or has_return_none):
         return True
 
     if not has_yield and not has_return and has_return_none:
       _LOGGER.warning(
-          "Process method returned None (element won't be emitted): %s. Check if intended.",
+          "Process method returned None (element won't be emitted): %s."
+          " Check if intended.",
           fn.__self__.__class__)
     print(has_yield, has_return, has_return_none)
 

sdks/python/apache_beam/transforms/core_test.py

@@ -114,10 +114,12 @@ def process(self, element):
     return None
 
 
-class TestDoFn12(beam.DoFn):
-  """test process returning None (return statement without a value)"""
+class TestDoFn11(beam.DoFn):
+  """test process returning None in a filter pattern"""
   def process(self, element):
-    return
+    if element == 0:
+      return
+    return element
 
 
 class TestDoFnStateful(beam.DoFn):
@@ -185,11 +187,10 @@ def test_dofn_with_explicit_return_none(self):
       assert RETURN_NONE_PARTIAL_WARNING in self._caplog.text
       assert str(TestDoFn10) in self._caplog.text
 
-  def test_dofn_with_implicit_return_none_return_without_value(self):
+  def test_dofn_with_implicit_return_none_and_value(self):
     with self._caplog.at_level(logging.WARNING):
-      beam.ParDo(TestDoFn12())
-      assert RETURN_NONE_PARTIAL_WARNING in self._caplog.text
-      assert str(TestDoFn12) in self._caplog.text
+      beam.ParDo(TestDoFn11())
+      assert RETURN_NONE_PARTIAL_WARNING not in self._caplog.text
 
 
 class PartitionTest(unittest.TestCase):