[DO NOT MERGE] Fix NewType and NamedTuple (#3)

sdks/python/apache_beam/runners/interactive/sql/utils.py

@@ -115,8 +115,7 @@ def pformat_namedtuple(schema: NamedTuple) -> str:
   return '{}({})'.format(
       schema.__name__,
       ', '.join([
-          '{}: {}'.format(k, v.__name__ if hasattr(v, '__name__') else repr(v))
-          for k,
+          '{}: {}'.format(k, repr(v)) for k,
           v in schema.__annotations__.items()
       ]))
 

sdks/python/apache_beam/runners/interactive/sql/utils_test.py

@@ -77,7 +77,7 @@ def test_replace_single_pcoll_token(self):
 
   def test_pformat_namedtuple(self):
     actual = pformat_namedtuple(ANamedTuple)
-    self.assertEqual('ANamedTuple(a: int, b: str)', actual)
+    self.assertEqual("ANamedTuple(a: <class 'int'>, b: <class 'str'>)", actual)
 
   def test_pformat_namedtuple_with_unnamed_fields(self):
     actual = pformat_namedtuple(OptionalUnionType)

sdks/python/apache_beam/runners/portability/fn_api_runner/fn_runner_test.py

@@ -191,7 +191,8 @@ def test_batch_rebatch_pardos(self):
     # - The output batch type of the producer
     # - The input batch type of the consumer
     with self.assertWarnsRegex(InefficientExecutionWarning,
-                               r'ListPlusOneDoFn.*NumpyArray.*List\[int64\]'):
+                               (r'ListPlusOneDoFn.*NumpyArray.*List\[<class '
+                                r'\'numpy.int64\'>\]')):
       with self.create_pipeline() as p:
         res = (
             p

sdks/python/apache_beam/transforms/batch_dofn_test.py

@@ -240,7 +240,8 @@ def test_mismatched_batch_producer_raises(self):
     with self.assertRaisesRegex(
         TypeError,
         (r'(?ms)MismatchedBatchProducingDoFn.*'
-         r'process: List\[int\].*process_batch: List\[float\]')):
+         r'process: List\[<class \'int\'>\].*process_batch: '
+         r'List\[<class \'float\'>\]')):
       _ = pc | beam.ParDo(MismatchedBatchProducingDoFn())
 
   def test_mismatched_element_producer_raises(self):

sdks/python/apache_beam/transforms/ptransform_test.py

@@ -1584,7 +1584,7 @@ def test_group_by_does_not_type_check(self):
         e.exception.args[0],
         "Input type hint violation at T: "
         "expected Tuple[TypeVariable[K], TypeVariable[V]], "
-        "got Iterable[int]")
+        "got Iterable[<class 'int'>]")
 
   def test_pipeline_checking_pardo_insufficient_type_information(self):
     self.p._options.view_as(TypeOptions).type_check_strictness = 'ALL_REQUIRED'
@@ -1705,9 +1705,9 @@ def is_even_as_key(a):
     self.assertStartswith(
         e.exception.args[0],
         "Runtime type violation detected within ParDo(IsEven): "
-        "Tuple[bool, int] hint type-constraint violated. "
+        "Tuple[<class 'bool'>, <class 'int'>] hint type-constraint violated. "
         "The type of element #0 in the passed tuple is incorrect. "
-        "Expected an instance of type bool, "
+        "Expected an instance of type <class 'bool'>, "
         "instead received an instance of type int.")
 
   def test_pipeline_checking_satisfied_run_time_checking_satisfied(self):
@@ -1771,10 +1771,10 @@ def test_pipeline_runtime_checking_violation_composite_type_input(self):
         e.exception.args[0],
         "Runtime type violation detected within ParDo(Add): "
         "Type-hint for argument: 'x_y' violated: "
-        "Tuple[int, int] hint type-constraint violated. "
+        "Tuple[<class 'int'>, <class 'int'>] hint type-constraint violated. "
         "The type of element #1 in the passed tuple is incorrect. "
-        "Expected an instance of type int, instead received an instance "
-        "of type float.")
+        "Expected an instance of type <class 'int'>, instead received an "
+        "instance of type float.")
 
   def test_pipeline_runtime_checking_violation_simple_type_output(self):
     self.p._options.view_as(TypeOptions).runtime_type_check = True
@@ -2097,8 +2097,9 @@ def test_mean_globally_pipeline_checking_violated(self):
 
     expected_msg = \
       "Type hint violation for 'CombinePerKey': " \
-      "requires Tuple[TypeVariable[K], Union[float, float64, int, int64]] " \
-      "but got Tuple[None, str] for element"
+      "requires Tuple[TypeVariable[K], Union[<class 'float'>, <class 'int'>, " \
+      "<class 'numpy.float64'>, <class 'numpy.int64'>]] " \
+      "but got Tuple[None, <class 'str'>] for element"
 
     self.assertStartswith(e.exception.args[0], expected_msg)
 
@@ -2163,8 +2164,9 @@ def test_mean_per_key_pipeline_checking_violated(self):
 
     expected_msg = \
       "Type hint violation for 'CombinePerKey(MeanCombineFn)': " \
-      "requires Tuple[TypeVariable[K], Union[float, float64, int, int64]] " \
-      "but got Tuple[str, str] for element"
+      "requires Tuple[TypeVariable[K], Union[<class 'float'>, <class 'int'>, " \
+      "<class 'numpy.float64'>, <class 'numpy.int64'>]] " \
+      "but got Tuple[<class 'str'>, <class 'str'>] for element"
 
     self.assertStartswith(e.exception.args[0], expected_msg)
 
@@ -2203,8 +2205,10 @@ def test_mean_per_key_runtime_checking_violated(self):
       "Runtime type violation detected within " \
       "OddMean/CombinePerKey(MeanCombineFn): " \
       "Type-hint for argument: 'element' violated: " \
-      "Union[float, float64, int, int64] type-constraint violated. " \
-      "Expected an instance of one of: ('float', 'float64', 'int', 'int64'), " \
+      "Union[<class 'float'>, <class 'int'>, <class 'numpy.float64'>, <class " \
+      "'numpy.int64'>] type-constraint violated. " \
+      "Expected an instance of one of: (\"<class 'float'>\", \"<class " \
+      "'int'>\", \"<class 'numpy.float64'>\", \"<class 'numpy.int64'>\"), " \
       "received str instead"
 
     self.assertStartswith(e.exception.args[0], expected_msg)
@@ -2570,9 +2574,9 @@ def test_inferred_bad_kv_type(self):
 
     self.assertStartswith(
         e.exception.args[0],
-        'Input type hint violation at GroupByKey: '
-        'expected Tuple[TypeVariable[K], TypeVariable[V]], '
-        'got Tuple[str, int, float]')
+        "Input type hint violation at GroupByKey: "
+        "expected Tuple[TypeVariable[K], TypeVariable[V]], "
+        "got Tuple[<class 'str'>, <class 'int'>, <class 'float'>]")
 
   def test_type_inference_command_line_flag_toggle(self):
     self.p._options.view_as(TypeOptions).pipeline_type_check = False

sdks/python/apache_beam/typehints/native_type_compatibility.py

@@ -192,6 +192,12 @@ def convert_to_beam_type(typ):
     # TODO(https://github.com/apache/beam/issues/19954): Currently unhandled.
     _LOGGER.info('Converting string literal type hint to Any: "%s"', typ)
     return typehints.Any
+  elif sys.version_info >= (3, 10) and isinstance(typ, typing.NewType):  # pylint: disable=isinstance-second-argument-not-valid-type
+    # Special case for NewType, where, since Python 3.10, NewType is now a class
+    # rather than a function.
+    # TODO(https://github.com/apache/beam/issues/20076): Currently unhandled.
+    _LOGGER.info('Converting NewType type hint to Any: "%s"', typ)
+    return typehints.Any
   elif getattr(typ, '__module__', None) != 'typing':
     # Only translate types from the typing module.
     return typ

sdks/python/apache_beam/typehints/row_type.py

@@ -171,8 +171,7 @@ def __hash__(self):
 
   def __repr__(self):
     return 'Row(%s)' % ', '.join(
-        '%s=%s' % (name, typehints._unified_repr(t)) for name,
-        t in self._fields)
+        '%s=%s' % (name, repr(t)) for name, t in self._fields)
 
   def get_type_for(self, name):
     return dict(self._fields)[name]

sdks/python/apache_beam/typehints/sharded_key_type.py

@@ -61,10 +61,8 @@ def type_check(self, instance):
       raise typehints.CompositeTypeHintError(
           "%s type-constraint violated. The type of key in 'ShardedKey' "
           "is incorrect. Expected an instance of type '%s', "
-          "instead received an instance of type '%s'." % (
-              repr(self),
-              typehints._unified_repr(self.key_type),
-              instance.key.__class__.__name__))
+          "instead received an instance of type '%s'." %
+          (repr(self), repr(self.key_type), instance.key.__class__.__name__))
 
   def match_type_variables(self, concrete_type):
     if isinstance(concrete_type, ShardedKeyTypeConstraint):
@@ -80,7 +78,7 @@ def __hash__(self):
     return hash(self.key_type)
 
   def __repr__(self):
-    return 'ShardedKey[%s]' % typehints._unified_repr(self.key_type)
+    return 'ShardedKey[%s]' % repr(self.key_type)
 
 
 ShardedKeyType = ShardedKeyTypeConstraint

sdks/python/apache_beam/typehints/sharded_key_type_test.py

@@ -37,7 +37,7 @@ def test_compatibility(self):
 
   def test_repr(self):
     constraint = ShardedKeyType[int]
-    self.assertEqual('ShardedKey[int]', repr(constraint))
+    self.assertEqual('ShardedKey[<class \'int\'>]', repr(constraint))
 
   def test_type_check_not_sharded_key(self):
     constraint = ShardedKeyType[int]
@@ -55,9 +55,9 @@ def test_type_check_invalid_key_type(self):
     with self.assertRaises((TypeError, TypeError)) as e:
       constraint.type_check(obj)
     self.assertEqual(
-        "ShardedKey[int] type-constraint violated. The type of key in "
-        "'ShardedKey' is incorrect. Expected an instance of type 'int', "
-        "instead received an instance of type 'str'.",
+        "ShardedKey[<class \'int\'>] type-constraint violated. The type of key "
+        "in 'ShardedKey' is incorrect. Expected an instance of type \'<class "
+        "\'int\'>\', instead received an instance of type 'str'.",
         e.exception.args[0])
 
   def test_type_check_valid_simple_type(self):

sdks/python/apache_beam/typehints/typecheck_test.py

@@ -263,9 +263,9 @@ def is_even_as_key(a):
         e.exception.args[0],
         "Runtime type violation detected within ParDo(IsEven): "
         "Type-hint for return type violated: "
-        "Tuple[bool, int] hint type-constraint violated. "
-        "The type of element #0 in the passed tuple is incorrect. "
-        "Expected an instance of type bool, "
+        "Tuple[<class \'bool\'>, <class \'int\'>] hint type-constraint "
+        "violated. The type of element #0 in the passed tuple is incorrect. "
+        "Expected an instance of type <class \'bool\'>, "
         "instead received an instance of type int. ")
 
   def test_pipeline_runtime_checking_violation_composite_type_output(self):

sdks/python/apache_beam/typehints/typed_pipeline_test.py

@@ -126,12 +126,14 @@ def process(self, element: int) -> typehints.Tuple[str]:
     result = [(1, 2)] | beam.ParDo(MyDoFn())
     self.assertEqual([1], sorted(result))
 
-    with self.assertRaisesRegex(typehints.TypeCheckError,
-                                r'requires.*Tuple\[int, int\].*got.*str'):
+    with self.assertRaisesRegex(
+        typehints.TypeCheckError,
+        r'requires.*Tuple\[<class \'int\'>, <class \'int\'>\].*got.*str'):
       _ = ['a', 'b', 'c'] | beam.ParDo(MyDoFn())
 
-    with self.assertRaisesRegex(typehints.TypeCheckError,
-                                r'requires.*Tuple\[int, int\].*got.*int'):
+    with self.assertRaisesRegex(
+        typehints.TypeCheckError,
+        r'requires.*Tuple\[<class \'int\'>, <class \'int\'>\].*got.*int'):
       _ = [1, 2, 3] | (beam.ParDo(MyDoFn()) | 'again' >> beam.ParDo(MyDoFn()))
 
   def test_typed_callable_iterable_output(self):
@@ -745,7 +747,8 @@ def repeat(s, *times):
 
     with self.assertRaisesRegex(
         typehints.TypeCheckError,
-        r'requires Tuple\[int, ...\] but got Tuple\[str, ...\]'):
+        (r'requires Tuple\[<class \'int\'>, ...\] but got '
+         r'Tuple\[<class \'str\'>, ...\]')):
       ['a', 'bb', 'c'] | beam.Map(repeat, 'z')
 
   def test_var_positional_only_side_input_hint(self):
@@ -762,8 +765,8 @@ def test_var_positional_only_side_input_hint(self):
 
     with self.assertRaisesRegex(
         typehints.TypeCheckError,
-        r'requires Tuple\[Union\[int, str\], ...\] but got '
-        r'Tuple\[Union\[float, int\], ...\]'):
+        r'requires Tuple\[Union\[<class \'int\'>, <class \'str\'>\], ...\] but '
+        r'got Tuple\[Union\[<class \'float\'>, <class \'int\'>\], ...\]'):
       _ = [1.2] | beam.Map(lambda *_: 'a', 5).with_input_types(int, str)
 
   def test_var_keyword_side_input_hint(self):
@@ -783,7 +786,8 @@ def test_var_keyword_side_input_hint(self):
 
     with self.assertRaisesRegex(
         typehints.TypeCheckError,
-        r'requires Dict\[str, str\] but got Dict\[str, int\]'):
+        r'requires Dict\[<class \'str\'>, <class \'str\'>\] but got '
+        r'Dict\[<class \'str\'>, <class \'int\'>\]'):
       _ = (['a', 'b', 'c']
            | beam.Map(lambda e, **_: 'a', kw=5).with_input_types(
                str, ignored=str))