[#20489] Add support for re.Pattern and re.Match type checking (#34604)

* [#20489] Add support for re.Pattern and re.Match type checking

* remove debugging statement

Author: jrmccluskey

sdks/python/apache_beam/typehints/native_type_compatibility_test.py

@@ -21,6 +21,7 @@
 
 import collections.abc
 import enum
+import re
 import typing
 import unittest
 
@@ -300,16 +301,17 @@ def test_newtype(self):
         typehints.Any, convert_to_beam_type(typing.NewType('Number', int)))
 
   def test_pattern(self):
-    # TODO(https://github.com/apache/beam/issues/20489): Unsupported.
-    self.assertEqual(typehints.Any, convert_to_beam_type(typing.Pattern))
-    self.assertEqual(typehints.Any, convert_to_beam_type(typing.Pattern[str]))
-    self.assertEqual(typehints.Any, convert_to_beam_type(typing.Pattern[bytes]))
+    self.assertEqual(re.Pattern, convert_to_beam_type(re.Pattern))
+    self.assertEqual(re.Pattern[str], convert_to_beam_type(re.Pattern[str]))
+    self.assertEqual(re.Pattern[bytes], convert_to_beam_type(re.Pattern[bytes]))
+    self.assertNotEqual(
+        re.Pattern[bytes], convert_to_beam_type(re.Pattern[str]))
 
   def test_match(self):
-    # TODO(https://github.com/apache/beam/issues/20489): Unsupported.
-    self.assertEqual(typehints.Any, convert_to_beam_type(typing.Match))
-    self.assertEqual(typehints.Any, convert_to_beam_type(typing.Match[str]))
-    self.assertEqual(typehints.Any, convert_to_beam_type(typing.Match[bytes]))
+    self.assertEqual(re.Match, convert_to_beam_type(re.Match))
+    self.assertEqual(re.Match[str], convert_to_beam_type(re.Match[str]))
+    self.assertEqual(re.Match[bytes], convert_to_beam_type(re.Match[bytes]))
+    self.assertNotEqual(re.Match[bytes], convert_to_beam_type(re.Match[str]))
 
   def test_forward_reference(self):
     self.assertEqual(typehints.Any, convert_to_beam_type('int'))

sdks/python/apache_beam/typehints/typehints.py

@@ -1516,15 +1516,40 @@ def is_consistent_with(sub, base):
   elif isinstance(sub, TypeConstraint):
     # Nothing but object lives above any type constraints.
     return base == object
+  elif getattr(base, '__module__', None) == 're':
+    return regex_consistency(sub, base)
   elif is_typing_generic(base):
     # Cannot check unsupported parameterized generic which will cause issubclass
     # to fail with an exception.
     return False
   return issubclass(sub, base)
 
 
+def regex_consistency(sub, base) -> bool:
+  """Checks whether two regular expression (re) type hints are consistent
+  with each other.
+
+  Either the sub or base hint can be parameterized generics since the set of
+  possible parameters is restricted to str | bytes. A base hint without a
+  parameter is treated as re.Class[str|bytes] so any sub param with a matching
+  base class is consistent. On the flip side, a sub hint without a parameter is
+  treated as inconsistent with a parameterized base hint.
+  """
+  base_generic = getattr(base, '__origin__', None)
+  sub_class = getattr(sub, '__origin__', sub)
+  if base_generic:
+    if sub_class == sub:
+      # if the sub hint is not parameterized but the base hint is we
+      # auto-fail
+      return False
+    return issubclass(sub_class,
+                      base_generic) and (sub.__args__ == base.__args__)
+  else:
+    return issubclass(sub_class, base)
+
+
 def get_yielded_type(type_hint):
-  """Obtains the type of elements yielded by an iterable.
+  """Obtains the type of elements yielded by an iterable.s
 
   Note that "iterable" here means: can be iterated over in a for loop, excluding
   strings and dicts.

sdks/python/apache_beam/typehints/typehints_test.py

@@ -21,6 +21,7 @@
 
 import collections.abc
 import functools
+import re
 import sys
 import typing
 import unittest
@@ -1220,6 +1221,35 @@ def increment(a):
         e.exception.args[0])
 
 
+class RegexTestCase(TypeHintTestCase):
+  def test_pattern(self):
+    self.assertCompatible(re.Pattern, re.Pattern)
+    self.assertCompatible(re.Pattern, re.Pattern[str])
+    self.assertCompatible(re.Pattern, re.Pattern[bytes])
+    self.assertCompatible(re.Pattern[str], re.Pattern[str])
+    self.assertCompatible(re.Pattern[bytes], re.Pattern[bytes])
+    self.assertNotCompatible(re.Pattern[str], re.Pattern)
+    self.assertNotCompatible(re.Pattern[bytes], re.Pattern)
+    self.assertNotCompatible(re.Pattern[str], re.Pattern[bytes])
+    self.assertNotCompatible(re.Pattern[bytes], re.Pattern[str])
+
+  def test_match(self):
+    self.assertCompatible(re.Match, re.Match)
+    self.assertCompatible(re.Match, re.Match[str])
+    self.assertCompatible(re.Match, re.Match[bytes])
+    self.assertCompatible(re.Match[str], re.Match[str])
+    self.assertCompatible(re.Match[bytes], re.Match[bytes])
+    self.assertNotCompatible(re.Match[str], re.Match)
+    self.assertNotCompatible(re.Match[bytes], re.Match)
+    self.assertNotCompatible(re.Match[str], re.Match[bytes])
+    self.assertNotCompatible(re.Match[bytes], re.Match[str])
+
+  def test_mix_fails(self):
+    self.assertNotCompatible(re.Pattern, re.Match)
+    self.assertNotCompatible(re.Match, re.Pattern)
+    self.assertNotCompatible(re.Pattern[str], re.Match[str])
+
+
 class TakesDecoratorTestCase(TypeHintTestCase):
   def test_must_be_primitive_type_or_constraint(self):
     with self.assertRaises(TypeError) as e: