[BEAM-7798] Typehints fix for bundle based runner.

apache · Jul 23, 2019 · 16eebb5 · 16eebb5
1 parent f357c48
commit 16eebb5
Showing 1 changed file with 6 additions and 5 deletions.
diff --git a/sdks/python/apache_beam/runners/direct/helper_transforms.py b/sdks/python/apache_beam/runners/direct/helper_transforms.py
@@ -22,6 +22,7 @@
 import typing
 
 import apache_beam as beam
+from apache_beam import typehints
 from apache_beam.internal.util import ArgumentPlaceholder
 from apache_beam.transforms.combiners import _CurriedFn
 from apache_beam.utils.windowed_value import WindowedValue
@@ -78,11 +79,11 @@ def default_type_hints(self):
     K = typing.TypeVar('K')
     if hints.input_types:
       args, kwargs = hints.input_types
-      args = (typing.Tuple[K, args[0]],) + args[1:]
+      args = (typehints.Tuple[K, args[0]],) + args[1:]
       hints.set_input_types(*args, **kwargs)
     else:
-      hints.set_input_types(typing.Tuple[K, typing.Any])
-    hints.set_output_types(typing.Tuple[K, typing.Any])
+      hints.set_input_types(typehints.Tuple[K, typing.Any])
+    hints.set_output_types(typehints.Tuple[K, typing.Any])
     return hints
 
 
@@ -102,8 +103,8 @@ def process(self, element):
   def default_type_hints(self):
     hints = self._combine_fn.get_type_hints().copy()
     K = typing.TypeVar('K')
-    hints.set_input_types(typing.Tuple[K, typing.Any])
+    hints.set_input_types(typehints.Tuple[K, typing.Any])
     if hints.output_types:
       main_output_type = hints.simple_output_type('')
-      hints.set_output_types(typing.Tuple[K, main_output_type])
+      hints.set_output_types(typehints.Tuple[K, main_output_type])
     return hints