Add some more cloudpickle tests. (#34683)

claudevdm · web-flow · commit 9abd394d3fb0 · 2025-04-22T10:01:11.000-07:00
* Add tests.

* Comments.
diff --git a/sdks/python/apache_beam/internal/cloudpickle_pickler_test.py b/sdks/python/apache_beam/internal/cloudpickle_pickler_test.py
@@ -27,12 +27,89 @@
 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.utils import shared
+
+GLOBAL_DICT_REF = module_test.GLOBAL_DICT
+
+
+# Allow weakref to dict
+class DictWrapper(dict):
+  pass
+
+
+MAIN_MODULE_DICT = DictWrapper()
+
+
+def acquire_dict():
+  return DictWrapper()
 
 
 class PicklerTest(unittest.TestCase):
 
   NO_MAPPINGPROXYTYPE = not hasattr(types, "MappingProxyType")
 
+  def test_globals_main_are_pickled_by_value(self):
+    self.assertIsNot(MAIN_MODULE_DICT, loads(dumps(lambda: MAIN_MODULE_DICT))())
+
+  def test_globals_shared_are_pickled_by_reference(self):
+    shared_handler = shared.Shared()
+    original_dict = shared_handler.acquire(acquire_dict)
+
+    unpickled_dict = loads(
+        dumps(lambda: shared_handler.acquire(acquire_dict)))()
+
+    self.assertIs(original_dict, unpickled_dict)
+
+  def test_module_globals_are_pickled_by_value_when_directly_referenced(self):
+    global_dict = loads(dumps(module_test.GLOBAL_DICT))
+
+    self.assertIsNot(module_test.GLOBAL_DICT, global_dict)
+
+  def test_function_main_with_explicit_module_reference_pickles_by_reference(
+      self):
+    def returns_global_dict():
+      return module_test.GLOBAL_DICT
+
+    self.assertIs(module_test.GLOBAL_DICT, loads(dumps(returns_global_dict))())
+
+  def test_function_main_with_indirect_module_reference_pickles_by_value(self):
+    def returns_global_dict():
+      return GLOBAL_DICT_REF
+
+    self.assertIsNot(
+        module_test.GLOBAL_DICT, loads(dumps(returns_global_dict))())
+
+  def test_function_referencing_unpicklable_object_works_when_imported(self):
+    self.assertEqual(
+        module_test.UNPICKLABLE_INSTANCE,
+        loads(dumps(module_test.fn_returns_unpicklable))())
+
+  def test_closure_with_unpicklable_object_fails_when_imported(self):
+    # The entire closure is pickled by value, and therefore module_test is
+    # not imported. Requires the module global to be pickled by value.
+    with self.assertRaises(Exception):
+      loads(dumps(module_test.closure_contains_unpicklable()))
+
+  def test_closure_with_explicit_self_import_can_reference_unpicklable_objects(
+      self):
+    # The closure imports module_test within the function definition
+    # and returns self.UNPICKLABLE_INSTANCE. This allows cloudpickle
+    # to use submimort to reference module_test.UNPICKLABLE_INSTANCE
+    self.assertIs(
+        module_test.UNPICKLABLE_INSTANCE,
+        loads(dumps(module_test.closure_contains_unpicklable_imports_self()))())
+
+  def test_closure_main_can_reference_unpicklable_module_objects(self):
+    def outer():
+      def inner():
+        return module_test.UNPICKLABLE_INSTANCE
+
+      return inner
+
+    # Uses subimport to reference module_test.UNPICKLABLE_INSTANCE rather than
+    # recreate.
+    self.assertIs(module_test.UNPICKLABLE_INSTANCE, loads(dumps(outer()))())
+
   def test_pickle_nested_enum_descriptor(self):
     NestedEnum = proto2_coder_test_messages_pb2.MessageD.NestedEnum
 
diff --git a/sdks/python/apache_beam/internal/module_test.py b/sdks/python/apache_beam/internal/module_test.py
@@ -23,6 +23,38 @@
 import sys
 from typing import Any
 
+GLOBAL_DICT = {}
+
+
+class UnPicklable:
+  def __init__(self, x):
+    self.x = x
+
+  def __reduce__(self):
+    return NotImplemented
+
+
+UNPICKLABLE_INSTANCE = UnPicklable(5)
+
+
+def closure_contains_unpicklable():
+  def inner():
+    return UNPICKLABLE_INSTANCE
+
+  return inner
+
+
+def closure_contains_unpicklable_imports_self():
+  def inner():
+    import apache_beam.internal.module_test as self_module
+    return self_module.UNPICKLABLE_INSTANCE
+
+  return inner
+
+
+def fn_returns_unpicklable():
+  return UNPICKLABLE_INSTANCE
+
 
 class TopClass(object):
   class NestedClass(object):
