Merge pull request #476 from neutrinoceros/hotfix_registry_cache_error

unyt/tests/test_unit_registry.py

@@ -57,6 +57,33 @@ def test_modify_symbol_from_default_unit_registry():
         default_unit_registry.remove("cm")
 
 
+def test_modify_cache_clear():
+    # see https://github.com/yt-project/unyt/issues/473
+
+    ureg = UnitRegistry()
+    ureg.add("celery", 1.0, length)
+    u0 = Unit("m", registry=ureg)
+    u1 = Unit("celery", registry=ureg)
+    assert u1 == u0
+
+    ureg.modify("celery", 0.5)
+
+    # check that this equality still holds post registry modification
+    assert u1 == u0
+
+    u2 = Unit("celery", registry=ureg)
+    assert u2 != u1
+    assert 1.0 * u2 == 0.5 * u0
+
+
+def test_remove_unit():
+    ureg = UnitRegistry()
+    ureg.add("celery", 1.0, length)
+    ureg.remove("celery")
+    with pytest.raises(UnitParseError):
+        Unit("celery", registry=ureg)
+
+
 def test_keys():
     ureg = UnitRegistry()
     assert sorted(ureg.keys()) == sorted(ureg.lut.keys())

unyt/unit_registry.py

@@ -185,6 +185,8 @@ def remove(self, symbol):
             )
 
         del self.lut[symbol]
+        if symbol in self._unit_object_cache:
+            del self._unit_object_cache[symbol]
 
     def modify(self, symbol, base_value):
         """
@@ -216,6 +218,8 @@ def modify(self, symbol, base_value):
             new_dimensions = self.lut[symbol][1]
 
         self.lut[symbol] = (float(base_value), new_dimensions) + self.lut[symbol][2:]
+        if symbol in self._unit_object_cache:
+            del self._unit_object_cache[symbol]
 
     def keys(self):
         """