Fix beam.Filter type hints for case where filter_fn specifies output type hint but not input type hint (#35283)

* remove unnecessary output_hint inferring logic

* extract input_hint and set it as the output_hint

* only pull out input hint if it's defined

* only conditionally set output hint if input hitn is not Any or None

* tidy up test

Author: hjtran

sdks/python/apache_beam/pipeline_test.py

@@ -639,6 +639,35 @@ def mux_input(x):
     self.assertNotIn(multi.letters, visitor.visited)
     self.assertNotIn(multi.numbers, visitor.visited)
 
+  def test_filter_typehint(self):
+    # Check input type hint and output type hint are both specified.
+    def always_true_with_all_typehints(x: int) -> bool:
+      return True
+
+    # Check only input type hint is specified.
+    def always_true_only_inptype(x: int):
+      return True
+
+    # Check only output type hint is specified.
+    def always_true_only_outptype(x) -> bool:
+      return True
+
+    # Check if inp type hint is Any that we can still infer
+    # from the input pcollection type
+    def always_true_any_inptype(x: typehints.Any) -> bool:
+      return True
+
+    for filter_fn in [always_true_with_all_typehints,
+                      always_true_only_inptype,
+                      always_true_only_outptype,
+                      always_true_any_inptype]:
+      with TestPipeline() as p:
+        pcoll = (
+            p
+            | beam.Create([1, 2, 3]).with_input_types(int)
+            | beam.Filter(filter_fn))
+      self.assertEqual(pcoll.element_type, int)
+
   def test_kv_ptransform_honor_type_hints(self):
 
     # The return type of this DoFn cannot be inferred by the default

sdks/python/apache_beam/transforms/core.py

@@ -2681,18 +2681,27 @@ def Filter(fn, *args, **kwargs):  # pylint: disable=invalid-name
 
   # Proxy the type-hint information from the function being wrapped, setting the
   # output type to be the same as the input type.
-  if type_hints.input_types is not None:
+  def has_simple_input_type(th):
+    return (
+        th.input_types is not None and len(th.input_types[0]) == 1 and
+        not th.input_types[1])
+
+  def simple_input_type(th):
+    return th.input_types[0][0] if has_simple_input_type(th) else None
+
+  if type_hints.input_types is not None and simple_input_type(
+      type_hints) is not typehints.Any:
     wrapper = with_input_types(
         *type_hints.input_types[0], **type_hints.input_types[1])(
             wrapper)
-  output_hint = type_hints.simple_output_type(label)
-  if (output_hint is None and get_type_hints(wrapper).input_types and
-      get_type_hints(wrapper).input_types[0]):
-    output_hint = get_type_hints(wrapper).input_types[0][0]
-  if output_hint:
-    wrapper = with_output_types(
-        typehints.Iterable[_strip_output_annotations(output_hint)])(
-            wrapper)
+    output_hint = type_hints.simple_output_type(label)
+    if (output_hint is None and get_type_hints(wrapper).input_types and
+        get_type_hints(wrapper).input_types[0]):
+      output_hint = get_type_hints(wrapper).input_types[0][0]
+    if output_hint:
+      wrapper = with_output_types(
+          typehints.Iterable[_strip_output_annotations(output_hint)])(
+              wrapper)
   # pylint: disable=protected-access
   wrapper._argspec_fn = fn
   # pylint: enable=protected-access