Add a trait for astropy quantities

docs/changes/2524.feature.rst

@@ -0,0 +1 @@
+Add an ``AstroQuantity`` trait which can hold any ``astropy.units.Quantity``.

src/ctapipe/core/tests/test_traits.py

@@ -9,6 +9,7 @@
 
 from ctapipe.core import Component
 from ctapipe.core.traits import (
+    AstroQuantity,
     AstroTime,
     List,
     Path,
@@ -278,6 +279,46 @@ class NoNone(Component):
         c.time = None
 
 
+def test_quantity():
+    import astropy.units as u
+
+    class SomeComponentWithQuantityTrait(Component):
+        quantity = AstroQuantity()
+
+    c = SomeComponentWithQuantityTrait()
+    c.quantity = -1.754 * u.m / (u.s * u.deg)
+    assert isinstance(c.quantity, u.Quantity)
+    assert c.quantity.value == -1.754
+    assert c.quantity.unit == u.Unit("m / (deg s)")
+
+    c.quantity = "1337 erg / centimeter**2 second"
+    assert isinstance(c.quantity, u.Quantity)
+    assert c.quantity.value == 1337
+    assert c.quantity.unit == u.Unit("erg / (s cm2)")
+
+    with pytest.raises(TraitError):
+        c.quantity = "No quantity"
+
+    # Try misspelled/ non-existent unit
+    with pytest.raises(TraitError):
+        c.quantity = "5 meters"
+
+
+def test_quantity_none():
+    class AllowNone(Component):
+        quantity = AstroQuantity(default_value=None, allow_none=True)
+
+    c = AllowNone()
+    assert c.quantity is None
+
+    class NoNone(Component):
+        quantity = AstroQuantity(default_value="5 meter", allow_none=False)
+
+    c = NoNone()
+    with pytest.raises(TraitError):
+        c.quantity = None
+
+
 def test_component_name():
     from ctapipe.core.traits import ComponentName, ComponentNameList
 

src/ctapipe/core/traits.py

@@ -5,6 +5,7 @@
 import pathlib
 from urllib.parse import urlparse
 
+import astropy.units as u
 import traitlets
 import traitlets.config
 from astropy.time import Time
@@ -17,6 +18,7 @@
 
 __all__ = [
     # Implemented here
+    "AstroQuantity",
     "AstroTime",
     "BoolTelescopeParameter",
     "IntTelescopeParameter",
@@ -72,8 +74,27 @@
 flag = traitlets.config.boolean_flag
 
 
+class AstroQuantity(TraitType):
+    """A trait containing an ``astropy.units`` quantity."""
+
+    def info(self):
+        info = "An ``astropy.units.Quantity`` instance"
+        if self.allow_none:
+            info += "or None"
+        return info
+
+    def validate(self, obj, value):
+        try:
+            quantity = u.Quantity(value)
+            return quantity
+        except TypeError:
+            return self.error(obj, value)
+        except ValueError:
+            return self.error(obj, value)
+
+
 class AstroTime(TraitType):
-    """A trait representing a point in Time, as understood by `astropy.time`"""
+    """A trait representing a point in Time, as understood by ``astropy.time``."""
 
     def validate(self, obj, value):
         """try to parse and return an ISO time string"""