Merge pull request #27 from twosigma/pickle-compatibility

Author: daniel-shields

pyproject.toml

@@ -14,7 +14,7 @@ packages = [
 ]
 readme = "README.md"
 repository = "https://github.com/twosigma/uberjob"
-version = "1.0.1"
+version = "1.0.2"
 
 [tool.poetry.dependencies]
 networkx = "^2.5"

src/uberjob/__init__.py

@@ -29,7 +29,7 @@
 __author__ = "Daniel Shields, Timothy Shields"
 __maintainer__ = "Daniel Shields, Timothy Shields"
 __email__ = "[email protected], [email protected]"
-__version__ = "1.0.1"
+__version__ = "1.0.2"
 
 from uberjob import graph, progress, stores
 from uberjob._errors import CallError, NotTransformedError

src/uberjob/_plan.py

@@ -49,6 +49,14 @@ def __init__(self):
         self._scope = ()
         self._scope_lock = RLock()
 
+    def __getstate__(self) -> dict:
+        return {"graph": self.graph}
+
+    def __setstate__(self, state: dict) -> None:
+        self.graph = state["graph"]
+        self._scope = ()
+        self._scope_lock = RLock()
+
     def _call(self, stack_frame, fn: Callable, *args, **kwargs) -> Call:
         call = Call(fn, scope=self._scope, stack_frame=stack_frame)
         self.graph.add_node(call)

src/uberjob/_transformations/caching.py

@@ -29,13 +29,17 @@
 
 
 class BarrierType:
-    __slots__ = ()
+    def __new__(cls):
+        return Barrier
+
+    def __getnewargs__(self):
+        return ()
 
     def __repr__(self):
         return "Barrier"
 
 
-Barrier = BarrierType()
+Barrier = object.__new__(BarrierType)
 
 
 def _to_naive_utc_time(value: dt.datetime | None) -> dt.datetime | None:

src/uberjob/_util/__init__.py

@@ -37,23 +37,31 @@ def is_ipython():
 
 
 class OmittedType:
-    __slots__ = ()
+    def __new__(cls):
+        return Omitted
+
+    def __getnewargs__(self):
+        return ()
 
     def __repr__(self):
         return "<...>"
 
 
-Omitted = OmittedType()
+Omitted = object.__new__(OmittedType)
 
 
 class MissingType:
-    __slots__ = ()
+    def __new__(cls):
+        return Missing
+
+    def __getnewargs__(self):
+        return ()
 
     def __repr__(self):
         return "Missing"
 
 
-Missing = MissingType()
+Missing = object.__new__(MissingType)
 
 
 class Slot:

src/uberjob/_util/traceback.py

@@ -48,20 +48,18 @@ def __repr__(self):
         )
 
 
-TruncatedStackFrame = None
+class TruncatedStackFrameType:
+    def __new__(cls):
+        return TruncatedStackFrame
 
+    def __getnewargs__(self):
+        return ()
 
-class TruncatedStackFrameType:
     def __repr__(self):
         return "TruncatedStackFrame"
 
-    def __new__(cls, *args, **kwargs):
-        if TruncatedStackFrame is not None:
-            return TruncatedStackFrame
-        return super().__new__(cls, *args, **kwargs)
-
 
-TruncatedStackFrame = TruncatedStackFrameType()
+TruncatedStackFrame = object.__new__(TruncatedStackFrameType)
 
 
 MAX_TRACEBACK_DEPTH = 3

tests/test_atoms.py

@@ -0,0 +1,52 @@
+#
+# Copyright 2025 Two Sigma Open Source, LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+
+import pickle
+import unittest
+
+from uberjob._transformations.caching import BarrierType, Barrier
+from uberjob._util import MissingType, Missing, OmittedType, Omitted
+from uberjob._util.traceback import TruncatedStackFrameType, TruncatedStackFrame
+
+ATOMS = [
+    (BarrierType, Barrier, "Barrier"),
+    (MissingType, Missing, "Missing"),
+    (OmittedType, Omitted, "<...>"),
+    (TruncatedStackFrameType, TruncatedStackFrame, "TruncatedStackFrame"),
+]
+
+
+class AtomTestCase(unittest.TestCase):
+    def test_atom_singleton(self):
+        for atom_type, atom, _ in ATOMS:
+            with self.subTest(atom_type_name=atom_type.__name__):
+                self.assertIs(atom_type(), atom_type())
+                self.assertIs(atom_type(), atom)
+
+    def test_atom_pickle_round_trip_is_atom(self):
+        for atom_type, atom, _ in ATOMS:
+            with self.subTest(atom_type_name=atom_type.__name__):
+                self.assertIs(pickle.loads(pickle.dumps(atom)), atom)
+
+    def test_type_of_atom_is_atom_type(self):
+        for atom_type, atom, _ in ATOMS:
+            with self.subTest(atom_type_name=atom_type.__name__):
+                self.assertIs(type(atom), atom_type)
+
+    def test_atom_repr(self):
+        for atom_type, atom, atom_repr in ATOMS:
+            with self.subTest(atom_type_name=atom_type.__name__):
+                self.assertEqual(repr(atom), atom_repr)

tests/test_plan.py

@@ -405,3 +405,10 @@ def test_serialize_node_error(self):
         self.assertIsInstance(unpickled_exception, NodeError)
         self.assertIsInstance(unpickled_exception.node, Call)
         self.assertIs(unpickled_exception.node.fn, pow)
+
+    def test_plan_pickle_round_trip(self):
+        plan = uberjob.Plan()
+        result = plan.call(pow, 3, 2)
+        plan2, result2 = pickle.loads(pickle.dumps([plan, result]))
+        self.assertEqual(plan2._scope, ())
+        self.assertEqual(uberjob.run(plan2, output=result2), 9)

tests/test_version.py

@@ -18,7 +18,7 @@
 
 import uberjob
 
-EXPECTED_VERSION = "1.0.1"
+EXPECTED_VERSION = "1.0.2"
 REPOSITORY_ROOT = os.path.abspath(os.path.join(os.path.dirname(__file__), ".."))