Validate units are in the unit registry on set (#1013)

* Validate units are in the unit registry on set

Closes #1010

* add test for units, or rather add unit, unit tests ;P

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* actually do the change on datasets not group

* remove extraneous array in unit test

Co-authored-by: Morrow <dmorrow@anl.gov>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>

CHANGELOG.md

@@ -11,6 +11,7 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/).
 ### Changed
 - complete units overhaul, now using pint library
 - explicitly store axes as fixed string dtype
+- validate units are in the unit registry on set
 
 ### Fixed
 - Fixed bug in `from_COLORS` that misordered the relationship between variables and channels for 1D datasets.

WrightTools/_dataset.py

@@ -185,6 +185,8 @@ def units(self, value):
         if value is None:
             if "units" in self.attrs.keys():
                 self.attrs.pop("units")
+        elif value not in wt_units.ureg:
+            raise ValueError(f"'{value}' is not in the unit registry")
         else:
             try:
                 self.attrs["units"] = value

WrightTools/data/_axis.py

@@ -141,6 +141,18 @@ def size(self) -> int:
         """Size."""
         return functools.reduce(operator.mul, self.shape)
 
+    @property
+    def units(self):
+        return self._units
+
+    @units.setter
+    def units(self, value):
+        if value == "None":
+            value = None
+        if value is not None and value not in wt_units.ureg:
+            raise ValueError(f"'{value}' is not in the unit registry")
+        self._units = value
+
     @property
     def units_kind(self) -> str:
         """Units kind."""

tests/data/from_Aramis.py

@@ -11,6 +11,8 @@
 
 here = pathlib.Path(__file__).parent
 
+wt.units.ureg.define("@alias count = Cnt")
+
 # --- test ----------------------------------------------------------------------------------------
 
 

tests/data/from_PyCMDS.py

@@ -13,6 +13,7 @@
 
 here = os.path.abspath(os.path.dirname(__file__))
 
+wt.units.ureg.define("steps = []")
 
 # --- test ----------------------------------------------------------------------------------------
 

tests/units.py

@@ -26,3 +26,27 @@ def test_axis_convert_exception():
 def test_in_mm_conversion():
     assert np.isclose(wt.units.convert(25.4, "mm", "in"), 1.0)
     assert np.isclose(wt.units.convert(1.0, "in", "mm"), 25.4)
+
+
+def unit_registry_test1():
+    values = np.linspace(-1, 1, 51)
+    d = wt.Data(name="test")
+    try:
+        d.create_variable("Bgood", values=values, units="tesla")
+        d.transform("Bgood")
+    except ValueError:
+        assert False
+    else:
+        assert True
+
+
+def unit_registry_test2():
+    values = np.linspace(-1, 1, 51)
+    d = wt.Data(name="test")
+    try:
+        d.create_variable("Bbad", values=values, units="Tesla")
+        d.transform("Bbad")
+    except ValueError:
+        assert True
+    else:
+        assert False