diff --git a/glue/core/state.py b/glue/core/state.py index 449005c6a..b5174a27d 100644 --- a/glue/core/state.py +++ b/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()) diff --git a/glue/core/tests/test_state.py b/glue/core/tests/test_state.py index 514e9174c..e70a5ec51 100644 --- a/glue/core/tests/test_state.py +++ b/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