Merge #1176

1176: Makes unpickling work properly using the Application Registry r=hgrecco a=hgrecco - [x] Closes #1175 (insert issue number) - [x] Executed ``black -t py36 . && isort -rc . && flake8`` with no errors - [x] The change is fully covered by automated unit tests - [ ] ~Documented in docs/ as appropriate~ - [x] Added an entry to the CHANGES file Co-authored-by: Hernan <hernan.grecco@gmail.com>
hgrecco · Sep 22, 2020 · b9fed4f · b9fed4f
2 parents ffc05dc + 2fbd3f8
commit b9fed4f
Show file tree

Hide file tree

Showing 6 changed files with 35 additions and 9 deletions.
diff --git a/CHANGES b/CHANGES
@@ -4,8 +4,8 @@ Pint Changelog
 0.17 (unreleased)
 --------------
 
-- Nothing changed yet.
-
+- Fix unpickling, now it is using the APP_REGISTRY as expected.
+  (Issue #1175)
 
 0.16 (2020-09-13)
 -----------------

diff --git a/pint/__init__.py b/pint/__init__.py
@@ -89,6 +89,24 @@ def _unpickle(cls, *args):
     return cls(*args)
 
 
+def _unpickle_quantity(cls, *args):
+    """Rebuild quantity upon unpickling using the application registry.
+    """
+    return _unpickle(_APP_REGISTRY.Quantity, *args)
+
+
+def _unpickle_unit(cls, *args):
+    """Rebuild unit upon unpickling using the application registry.
+    """
+    return _unpickle(_APP_REGISTRY.Unit, *args)
+
+
+def _unpickle_measurement(cls, *args):
+    """Rebuild measurement upon unpickling using the application registry.
+    """
+    return _unpickle(_APP_REGISTRY.Measurement, *args)
+
+
 def set_application_registry(registry):
     """Set the application registry, which is used for unpickling operations
     and when invoking pint.Quantity or pint.Unit directly.

diff --git a/pint/measurement.py b/pint/measurement.py
@@ -70,9 +70,9 @@ def rel(self):
 
     def __reduce__(self):
         # See notes in Quantity.__reduce__
-        from . import _unpickle
+        from . import _unpickle_measurement
 
-        return _unpickle, (Measurement, self.magnitude, self._units)
+        return _unpickle_measurement, (Measurement, self.magnitude, self._units)
 
     def __repr__(self):
         return "<Measurement({}, {}, {})>".format(

diff --git a/pint/quantity.py b/pint/quantity.py
@@ -192,11 +192,11 @@ def __reduce__(self):
         """Allow pickling quantities. Since UnitRegistries are not pickled, upon
         unpickling the new object is always attached to the application registry.
         """
-        from . import _unpickle
+        from . import _unpickle_quantity
 
         # Note: type(self) would be a mistake as subclasses built by
         # build_quantity_class can't be pickled
-        return _unpickle, (Quantity, self.magnitude, self._units)
+        return _unpickle_quantity, (Quantity, self.magnitude, self._units)
 
     def __new__(cls, value, units=None):
         if is_upcast_type(type(value)):

diff --git a/pint/testsuite/test_issues.py b/pint/testsuite/test_issues.py
@@ -5,7 +5,7 @@
 
 import pytest
 
-from pint import Context, DimensionalityError, UnitRegistry
+from pint import Context, DimensionalityError, UnitRegistry, get_application_registry
 from pint.compat import np
 from pint.testsuite import QuantityTestCase, helpers
 from pint.unit import UnitsContainer
@@ -775,6 +775,14 @@ def test_issue_1136(self):
         with pytest.raises(DimensionalityError):
             2 ** ureg.Quantity([2, 3], "m")
 
+    def test_issue1175(self):
+        import pickle
+
+        foo1 = get_application_registry().Quantity(1, "s")
+        foo2 = pickle.loads(pickle.dumps(foo1))
+        self.assertIsInstance(foo1, foo2.__class__)
+        self.assertIsInstance(foo2, foo1.__class__)
+
 
 if np is not None:
 

diff --git a/pint/unit.py b/pint/unit.py
@@ -28,9 +28,9 @@ class Unit(PrettyIPython, SharedRegistryObject):
 
     def __reduce__(self):
         # See notes in Quantity.__reduce__
-        from . import _unpickle
+        from . import _unpickle_unit
 
-        return _unpickle, (Unit, self._units)
+        return _unpickle_unit, (Unit, self._units)
 
     def __init__(self, units):
         super().__init__()