Add reducers for threading.Lock and EnumDescriptor. (#34537)

Author: claudevdm

sdks/python/apache_beam/internal/cloudpickle_pickler.py

@@ -30,6 +30,7 @@
 import base64
 import bz2
 import io
+import sys
 import threading
 import zlib
 
@@ -40,10 +41,65 @@
 except (ImportError, ModuleNotFoundError):
   pass
 
+
+def _get_proto_enum_descriptor_class():
+  try:
+    from google.protobuf.internal import api_implementation
+  except ImportError:
+    return None
+
+  implementation_type = api_implementation.Type()
+
+  if implementation_type == 'upb':
+    try:
+      from google._upb._message import EnumDescriptor
+      return EnumDescriptor
+    except ImportError:
+      pass
+  elif implementation_type == 'cpp':
+    try:
+      from google.protobuf.pyext._message import EnumDescriptor
+      return EnumDescriptor
+    except ImportError:
+      pass
+  elif implementation_type == 'python':
+    try:
+      from google.protobuf.internal.python_message import EnumDescriptor
+      return EnumDescriptor
+    except ImportError:
+      pass
+
+  return None
+
+
+EnumDescriptor = _get_proto_enum_descriptor_class()
+
 # Pickling, especially unpickling, causes broken module imports on Python 3
 # if executed concurrently, see: BEAM-8651, http://bugs.python.org/issue38884.
 _pickle_lock = threading.RLock()
 RLOCK_TYPE = type(_pickle_lock)
+LOCK_TYPE = type(threading.Lock())
+
+
+def _reconstruct_enum_descriptor(full_name):
+  for _, module in sys.modules.items():
+    if not hasattr(module, 'DESCRIPTOR'):
+      continue
+
+    for _, attr_value in vars(module).items():
+      if not hasattr(attr_value, 'DESCRIPTOR'):
+        continue
+
+      if hasattr(attr_value.DESCRIPTOR, 'enum_types_by_name'):
+        for (_, enum_desc) in attr_value.DESCRIPTOR.enum_types_by_name.items():
+          if enum_desc.full_name == full_name:
+            return enum_desc
+  raise ImportError(f'Could not find enum descriptor: {full_name}')
+
+
+def _pickle_enum_descriptor(obj):
+  full_name = obj.full_name
+  return _reconstruct_enum_descriptor, (full_name, )
 
 
 def dumps(o, enable_trace=True, use_zlib=False) -> bytes:
@@ -59,6 +115,12 @@ def dumps(o, enable_trace=True, use_zlib=False) -> bytes:
         pickler.dispatch_table[RLOCK_TYPE] = _pickle_rlock
       except NameError:
         pass
+      try:
+        pickler.dispatch_table[LOCK_TYPE] = _lock_reducer
+      except NameError:
+        pass
+      if EnumDescriptor is not None:
+        pickler.dispatch_table[EnumDescriptor] = _pickle_enum_descriptor
       pickler.dump(o)
       s = file.getvalue()
 
@@ -106,6 +168,10 @@ def _pickle_rlock(obj):
   return RLOCK_TYPE, tuple([])
 
 
+def _lock_reducer(obj):
+  return threading.Lock, tuple([])
+
+
 def dump_session(file_path):
   # It is possible to dump session with cloudpickle. However, since references
   # are saved it should not be necessary. See https://s.apache.org/beam-picklers

sdks/python/apache_beam/internal/cloudpickle_pickler_test.py

@@ -26,12 +26,21 @@
 from apache_beam.internal import module_test
 from apache_beam.internal.cloudpickle_pickler import dumps
 from apache_beam.internal.cloudpickle_pickler import loads
+from apache_beam.portability.api import beam_runner_api_pb2
 
 
 class PicklerTest(unittest.TestCase):
 
   NO_MAPPINGPROXYTYPE = not hasattr(types, "MappingProxyType")
 
+  def test_pickle_enum_descriptor(self):
+    TimeDomain = beam_runner_api_pb2.TimeDomain.Enum
+
+    def fn():
+      return TimeDomain.EVENT_TIME
+
+    self.assertEqual(fn(), loads(dumps(fn))())
+
   def test_basics(self):
     self.assertEqual([1, 'a', ('z', )], loads(dumps([1, 'a', ('z', )])))
     fun = lambda x: 'xyz-%s' % x
@@ -97,6 +106,12 @@ def test_pickle_rlock(self):
 
     self.assertIsInstance(loads(dumps(rlock_instance)), rlock_type)
 
+  def test_pickle_lock(self):
+    lock_instance = threading.Lock()
+    lock_type = type(lock_instance)
+
+    self.assertIsInstance(loads(dumps(lock_instance)), lock_type)
+
   @unittest.skipIf(NO_MAPPINGPROXYTYPE, 'test if MappingProxyType introduced')
   def test_dump_and_load_mapping_proxy(self):
     self.assertEqual(