Fix warning message

sdks/python/apache_beam/transforms/core.py

@@ -1507,25 +1507,29 @@ def _check_fn_use_yield_and_return(fn):
     source_code = _get_function_body_without_inners(fn)
     has_yield = False
     has_return = False
-    return_none_warning = (
-        "No iterator is returned by the process method in %s.",
-        fn.__self__.__class__)
+    has_return_none = False
     for line in source_code.split("\n"):
       lstripped_line = line.lstrip()
       if lstripped_line.startswith("yield ") or lstripped_line.startswith(
           "yield("):
         has_yield = True
-      if lstripped_line.startswith("return ") or lstripped_line.startswith(
-          "return("):
+      elif lstripped_line.rstrip() == "return":
         has_return = True
-        if lstripped_line.startswith(
-            "return None") or lstripped_line.rstrip() == "return":
-          _LOGGER.warning(return_none_warning)
-      if has_yield and has_return:
+      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 or has_return_none):
         return True
 
-    if not has_yield and not has_return:
-      _LOGGER.warning(return_none_warning)
+    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.",
+          fn.__self__.__class__)
+    print(has_yield, has_return, has_return_none)
 
     return False
   except Exception as e:

sdks/python/apache_beam/transforms/core_test.py

@@ -40,7 +40,7 @@
 from apache_beam.typehints import row_type
 from apache_beam.typehints import typehints
 
-RETURN_NONE_PARTIAL_WARNING = "No iterator is returned"
+RETURN_NONE_PARTIAL_WARNING = "Process method returned None"
 
 
 class TestDoFn1(beam.DoFn):
@@ -115,15 +115,11 @@ def process(self, element):
 
 
 class TestDoFn11(beam.DoFn):
-  """test process returning None (no return and no yield)"""
+  """test process returning None in a filter pattern"""
   def process(self, element):
-    pass
-
-
-class TestDoFn12(beam.DoFn):
-  """test process returning None (return statement without a value)"""
-  def process(self, element):
-    return
+    if element == 0:
+      return
+    return element
 
 
 class TestDoFnStateful(beam.DoFn):
@@ -191,17 +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_missing_return_and_yield(self):
+  def test_dofn_with_implicit_return_none_and_value(self):
     with self._caplog.at_level(logging.WARNING):
       beam.ParDo(TestDoFn11())
-      assert RETURN_NONE_PARTIAL_WARNING in self._caplog.text
-      assert str(TestDoFn11) in self._caplog.text
-
-  def test_dofn_with_implicit_return_none_return_without_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
+      assert RETURN_NONE_PARTIAL_WARNING not in self._caplog.text
 
 
 class PartitionTest(unittest.TestCase):