Merge pull request #2475 from Carifio24/serialize-units

Allow serializing and deserializing named astropy units

glue/core/state.py

@@ -65,6 +65,8 @@ def load(rec, context)
 import numpy as np
 from matplotlib.colors import Colormap
 from matplotlib import cm
+import astropy.units as u
+from astropy.units import UnitBase, Unit
 from astropy.wcs import WCS
 import shapely
 
@@ -621,6 +623,23 @@ def _load_slice(rec, context):
     return slice(rec['start'], rec['stop'], rec['step'])
 
 
+@saver(UnitBase)
+def _save_unit_base(unit, context):
+    unit_str = unit.to_string()
+    # Check that unit can be parsed back with default enabled systems
+    try:
+        with u.set_enabled_units([u.si, u.cgs, u.astrophys]):
+            _ = u.Unit(unit_str)
+    except ValueError:
+        raise GlueSerializeError(f"Serializing units of '{unit}' is not yet supported")
+    return dict(unit_base=unit_str)
+
+
+@loader(UnitBase)
+def _load_unit_base(rec, context):
+    return Unit(rec["unit_base"])
+
+
 @saver(WCS)
 def _save_wcs(wcs, context):
     return dict(header=wcs.to_header_string())

glue/core/tests/test_state.py

@@ -305,6 +305,34 @@ def test_datetime_component():
     assert isinstance(c2.data[0], np.datetime64)
 
 
+@requires_astropy
+def test_astropy_units():
+    import astropy.units as u
+    unit = u.m
+    unit2 = clone(unit)
+    assert unit2 is unit
+
+    unit = u.km
+    unit2 = clone(unit)
+    assert unit2 is unit
+
+
+@requires_astropy
+def test_astropy_compound_units():
+    import astropy.units as u
+    unit = u.m / u.s
+    unit2 = clone(unit)
+    assert unit2 == unit
+
+    unit = u.W / u.m**2 / u.nm
+    unit2 = clone(unit)
+    assert unit2 == unit
+
+    unit = u.km
+    unit2 = clone(unit)
+    assert unit2 is unit
+
+
 class DummyClass(object):
     pass