Repository: yt-project/unyt
Pull request #78 (2019-03-26 16:05:15)

fix issues working with code units
This fixes some issues I found trying to get unyt working with yt. First, it adds an explicit test for a code unit system like yt's. It also makes sure that custom units will show up in a unit container object like i've created in yt via `ds.units`, to make sure that things like `ds.units.code_length` will work as expected. I also added a new `has_current_mks` property on the `UnitSystem` class that I found useful in the process of getting yt working with unyt.


Patch:
diff --git a/docs/usage.rst b/docs/usage.rst
index 70ffbf3..79d0953 100644
--- a/docs/usage.rst
+++ b/docs/usage.rst
@@ -938,6 +938,8 @@ contains units that are compatible with the ``Simulation`` instance we named
   >>> units = UnitContainer(s_cgs.registry)
   >>> units.kilometer
   km
+  >>> units.code_length
+  code_length
   >>> (10.0 * units.kilometer).in_base()
   unyt_quantity(1000000., 'cm')
   >>> (10.0 * units.kilometer).in_units('code_length')
diff --git a/unyt/unit_registry.py b/unyt/unit_registry.py
index 46ff81a..9b2edb5 100644
--- a/unyt/unit_registry.py
+++ b/unyt/unit_registry.py
@@ -34,7 +34,7 @@ def _sanitize_unit_system(unit_system, obj):
     elif hasattr(unit_system, "unit_registry"):
         unit_system = unit_system.unit_registry.unit_system_id
     elif unit_system == "code":
-        unit_system = obj.registry.unit_system_id
+        unit_system = obj.units.registry.unit_system_id
     return unit_system_registry[str(unit_system)]
 
 
diff --git a/unyt/unit_systems.py b/unyt/unit_systems.py
index ae770a0..e235600 100644
--- a/unyt/unit_systems.py
+++ b/unyt/unit_systems.py
@@ -55,6 +55,9 @@ def add_symbols(namespace, registry):
     for canonical_name, alt_names in name_alternatives.items():
         for alt_name in alt_names:
             namespace[alt_name] = Unit(canonical_name, registry=registry)
+    for name in registry.keys():
+        if name not in namespace:
+            namespace[name] = Unit(name, registry=registry)
 
 
 def add_constants(namespace, registry):
@@ -275,6 +278,11 @@ class UnitSystem(object):
                 repr += "  %s: %s\n" % (key, self.units_map[dim])
         return repr[:-1]
 
+    @property
+    def has_current_mks(self):
+        """Does this unit system have an MKS current dimension?"""
+        return self.units_map[cmks] is not None
+
 
 #: The CGS unit system
 cgs_unit_system = UnitSystem("cgs", "cm", "g", "s", current_mks_unit=None)


Test patch:
diff --git a/unyt/tests/test_unit_systems.py b/unyt/tests/test_unit_systems.py
index adf7541..58aa047 100644
--- a/unyt/tests/test_unit_systems.py
+++ b/unyt/tests/test_unit_systems.py
@@ -17,7 +17,12 @@ import unyt.unit_symbols as us
 
 from unyt.exceptions import IllDefinedUnitSystem, MissingMKSCurrent
 from unyt.unit_object import Unit
-from unyt.unit_systems import UnitSystem, cgs_unit_system, unit_system_registry
+from unyt.unit_systems import (
+    UnitSystem,
+    cgs_unit_system,
+    mks_unit_system,
+    unit_system_registry,
+)
 from unyt.unit_registry import UnitRegistry
 from unyt import dimensions
 
@@ -69,6 +74,30 @@ def test_bad_unit_system():
         UnitSystem("atomic", us.nm, us.fs, us.nK, us.rad, registry=UnitRegistry())
 
 
+def test_code_unit_system():
+    ureg = UnitRegistry()
+    ureg.add("code_length", 2.0, dimensions.length)
+    ureg.add("code_mass", 3.0, dimensions.mass)
+    ureg.add("code_time", 4.0, dimensions.time)
+    ureg.add("code_temperature", 5.0, dimensions.temperature)
+    code_unit_system = UnitSystem(
+        "my_unit_system",
+        "code_length",
+        "code_mass",
+        "code_time",
+        "code_temperature",
+        registry=ureg,
+    )
+    assert code_unit_system["length"] == Unit("code_length", registry=ureg)
+    assert code_unit_system["length"].base_value == 2
+    assert code_unit_system["mass"] == Unit("code_mass", registry=ureg)
+    assert code_unit_system["mass"].base_value == 3
+    assert code_unit_system["time"] == Unit("code_time", registry=ureg)
+    assert code_unit_system["time"].base_value == 4
+    assert code_unit_system["temperature"] == Unit("code_temperature", registry=ureg)
+    assert code_unit_system["temperature"].base_value == 5
+
+
 def test_mks_current():
     with pytest.raises(MissingMKSCurrent):
         cgs_unit_system[dimensions.current_mks]
@@ -78,6 +107,8 @@ def test_mks_current():
         cgs_unit_system[dimensions.current_mks] = "foo"
     with pytest.raises(MissingMKSCurrent):
         cgs_unit_system[dimensions.magnetic_field] = "bar"
+    assert cgs_unit_system.has_current_mks is False
+    assert mks_unit_system.has_current_mks is True
 
 
 def test_create_unit_system_from_unit_objects():