Merge pull request #1366 from keewis/application_registry

improve the application registry

5 files changed, 100 insertions, 41 deletions

diff --git a/CHANGES b/CHANGES
index 477aa09..e53cc5b 100644
--- a/CHANGES
+++ b/CHANGES
@@ -16,6 +16,7 @@ Pint Changelog
 - Fix string formatting of numpy array scalars
 - Fix default format for Measurement class (Issue #1300)
 - Fix parsing of pretty units with same exponents but different sign. (Issue #1360)
+- Convert the application registry to a wrapper object (Issue #1365)
 
 ### Breaking Changes
 
diff --git a/pint/__init__.py b/pint/__init__.py
index 3db4bb5..9f126d5 100644
--- a/pint/__init__.py
+++ b/pint/__init__.py
@@ -25,9 +25,9 @@ from .errors import (  # noqa: F401
 from .formatting import formatter
 from .measurement import Measurement
 from .quantity import Quantity
-from .registry import LazyRegistry, UnitRegistry
+from .registry import ApplicationRegistry, LazyRegistry, UnitRegistry
 from .unit import Unit
-from .util import logger, pi_theorem
+from .util import logger, pi_theorem  # noqa: F401
 
 try:
     from importlib.metadata import version
@@ -47,7 +47,7 @@ except Exception:  # pragma: no cover
 _DEFAULT_REGISTRY = LazyRegistry()
 
 #: Registry used for unpickling operations.
-_APP_REGISTRY = _DEFAULT_REGISTRY
+application_registry = ApplicationRegistry(_DEFAULT_REGISTRY)
 
 
 def _unpickle(cls, *args):
@@ -72,24 +72,24 @@ def _unpickle(cls, *args):
         # We need to make sure that this happens before using.
         if isinstance(arg, UnitsContainer):
             for name in arg:
-                _APP_REGISTRY.parse_units(name)
+                application_registry.parse_units(name)
 
     return cls(*args)
 
 
 def _unpickle_quantity(cls, *args):
     """Rebuild quantity upon unpickling using the application registry."""
-    return _unpickle(_APP_REGISTRY.Quantity, *args)
+    return _unpickle(application_registry.Quantity, *args)
 
 
 def _unpickle_unit(cls, *args):
     """Rebuild unit upon unpickling using the application registry."""
-    return _unpickle(_APP_REGISTRY.Unit, *args)
+    return _unpickle(application_registry.Unit, *args)
 
 
 def _unpickle_measurement(cls, *args):
     """Rebuild measurement upon unpickling using the application registry."""
-    return _unpickle(_APP_REGISTRY.Measurement, *args)
+    return _unpickle(application_registry.Measurement, *args)
 
 
 def set_application_registry(registry):
@@ -100,11 +100,7 @@ def set_application_registry(registry):
     ----------
     registry : pint.UnitRegistry
     """
-    if not isinstance(registry, (LazyRegistry, UnitRegistry)):
-        raise TypeError("Expected UnitRegistry; got %s" % type(registry))
-    global _APP_REGISTRY
-    logger.debug("Changing app registry from %r to %r.", _APP_REGISTRY, registry)
-    _APP_REGISTRY = registry
+    application_registry.set(registry)
 
 
 def get_application_registry():
@@ -116,7 +112,7 @@ def get_application_registry():
     -------
     pint.UnitRegistry
     """
-    return _APP_REGISTRY
+    return application_registry
 
 
 def test():
diff --git a/pint/registry.py b/pint/registry.py
index 4428ad4..f53becc 100644
--- a/pint/registry.py
+++ b/pint/registry.py
@@ -2323,3 +2323,54 @@ class LazyRegistry:
     def __call__(self, *args, **kwargs):
         self.__init()
         return self(*args, **kwargs)
+
+
+class ApplicationRegistry:
+    """A wrapper class used to distribute changes to the application registry."""
+
+    def __init__(self, registry):
+        self._registry = registry
+
+    def get(self):
+        """Get the wrapped registry"""
+        return self._registry
+
+    def set(self, new_registry):
+        """Set the new registry
+
+        Parameters
+        ----------
+        new_registry : ApplicationRegistry or LazyRegistry or UnitRegistry
+            The new registry.
+
+        See Also
+        --------
+        set_application_registry
+        """
+        if isinstance(new_registry, type(self)):
+            new_registry = new_registry.get()
+
+        if not isinstance(new_registry, (LazyRegistry, UnitRegistry)):
+            raise TypeError("Expected UnitRegistry; got %s" % type(new_registry))
+        logger.debug(
+            "Changing app registry from %r to %r.", self._registry, new_registry
+        )
+        self._registry = new_registry
+
+    def __getattr__(self, name):
+        return getattr(self._registry, name)
+
+    def __dir__(self):
+        return dir(self._registry)
+
+    def __getitem__(self, item):
+        return self._registry[item]
+
+    def __call__(self, *args, **kwargs):
+        return self._registry(*args, **kwargs)
+
+    def __contains__(self, item):
+        return self._registry.__contains__(item)
+
+    def __iter__(self):
+        return iter(self._registry)
diff --git a/pint/testsuite/test_application_registry.py b/pint/testsuite/test_application_registry.py
index 59c3521..2ec349c 100644
--- a/pint/testsuite/test_application_registry.py
+++ b/pint/testsuite/test_application_registry.py
@@ -5,6 +5,7 @@ import pickle
 import pytest
 
 from pint import (
+    ApplicationRegistry,
     Measurement,
     Quantity,
     UndefinedUnitError,
@@ -16,6 +17,14 @@ from pint import (
 from pint.testsuite import helpers
 
 
+@pytest.fixture
+def isolate_application_registry(monkeypatch):
+    import pint
+
+    new = ApplicationRegistry(UnitRegistry())
+    monkeypatch.setattr(pint, "application_registry", new)
+
+
 class TestDefaultApplicationRegistry:
     @helpers.allprotos
     def test_unit(self, protocol):
@@ -94,20 +103,16 @@ class TestDefaultApplicationRegistry:
             pickle.loads(b)
 
 
-class TestCustomApplicationRegistry:
-    @classmethod
-    def setup_class(cls):
-        cls.ureg_bak = get_application_registry()
-        cls.ureg = UnitRegistry(None)
-        cls.ureg.define("foo = []")
-        cls.ureg.define("bar = foo / 2")
-        set_application_registry(cls.ureg)
-        assert get_application_registry() is cls.ureg
-
-    @classmethod
-    def teardown_class(cls):
-        set_application_registry(cls.ureg_bak)
+@pytest.fixture
+def custom_registry(isolate_application_registry):
+    ureg = UnitRegistry(None)
+    ureg.define("foo = []")
+    ureg.define("bar = foo / 2")
+    set_application_registry(ureg)
+
 
+@pytest.mark.usefixtures("custom_registry")
+class TestCustomApplicationRegistry:
     @helpers.allprotos
     def test_unit(self, protocol):
         u = Unit("foo")
@@ -158,6 +163,7 @@ class TestCustomApplicationRegistry:
         assert str(m.units) == "foo"
 
 
+@pytest.mark.usefixtures("isolate_application_registry")
 class TestSwapApplicationRegistry:
     """Test that the constructors of Quantity, Unit, and Measurement capture
     the registry that is set as the application registry at creation time
@@ -170,19 +176,12 @@ class TestSwapApplicationRegistry:
 
     """
 
-    @classmethod
-    def setup_class(cls):
-        cls.ureg_bak = get_application_registry()
-        cls.ureg1 = UnitRegistry(None)
-        cls.ureg1.define("foo = [dim1]")
-        cls.ureg1.define("bar = foo / 2")
-        cls.ureg2 = UnitRegistry(None)
-        cls.ureg2.define("foo = [dim2]")
-        cls.ureg2.define("bar = foo / 3")
-
-    @classmethod
-    def teardown_class(cls):
-        set_application_registry(cls.ureg_bak)
+    ureg1 = UnitRegistry(None)
+    ureg1.define("foo = [dim1]")
+    ureg1.define("bar = foo / 2")
+    ureg2 = UnitRegistry(None)
+    ureg2.define("foo = [dim2]")
+    ureg2.define("bar = foo / 3")
 
     @helpers.allprotos
     def test_quantity_1arg(self, protocol):
@@ -259,3 +258,15 @@ class TestSwapApplicationRegistry:
         assert m1.to("bar").error.magnitude == 2
         assert m2.to("bar").error.magnitude == 3
         assert m3.to("bar").error.magnitude == 3
+
+
+@pytest.mark.usefixtures("isolate_application_registry")
+def test_update_saved_registries():
+    ureg1 = get_application_registry()
+    ureg2 = get_application_registry()
+
+    new = UnitRegistry()
+    new.define("foo = []")
+    set_application_registry(new)
+
+    assert ureg1.Unit("foo") == ureg2.Unit("foo")
diff --git a/pint/util.py b/pint/util.py
index b0fed09..94d1193 100644
--- a/pint/util.py
+++ b/pint/util.py
@@ -816,9 +816,9 @@ class SharedRegistryObject:
         if not hasattr(cls, "_REGISTRY"):
             # Base class, not subclasses dynamically by
             # UnitRegistry._init_dynamic_classes
-            from . import _APP_REGISTRY
+            from . import application_registry
 
-            inst._REGISTRY = _APP_REGISTRY
+            inst._REGISTRY = application_registry.get()
         return inst
 
     def _check(self, other) -> bool: