squash! make subclassing of ObservationInfo easier

make simple and json round-trip with extension properties

python/astro_metadata_translator/observationInfo.py

@@ -349,6 +349,10 @@ def __eq__(self, other):
         self_simple = self.to_simple()
         other_simple = other.to_simple()
 
+        # We don't care about the translator internal detail
+        self_simple.pop("_translator", None)
+        other_simple.pop("_translator", None)
+
         for k, self_value in self_simple.items():
             other_value = other_simple[k]
             if self_value != other_value:
@@ -380,7 +384,7 @@ def __getstate__(self):
         for p in self.all_properties:
             state[p] = getattr(self, p)
 
-        return self.extensions, state
+        return state, self.extensions
 
     def __setstate__(self, state):
         """Set object state from pickle
@@ -390,11 +394,18 @@ def __setstate__(self, state):
         state : `tuple`
             Pickled state.
         """
-        extensions, state = state
+        try:
+            state, extensions = state
+        except ValueError:
+            # Backwards compatibility for pickles generated before DM-34175
+            extensions = {}
         self._declare_extensions(extensions)
         for p in self.all_properties:
-            property = f"_{p}"
-            setattr(self, property, state[p])
+            if p.startswith("ext_"):
+                super().__setattr__(p, state[p])  # allows setting even write-protected extensions
+            else:
+                property = f"_{p}"
+                setattr(self, property, state[p])
 
     def to_simple(self):
         """Convert the contents of this object to simple dict form.
@@ -412,11 +423,20 @@ def to_simple(self):
         -------
         simple : `dict` of [`str`, `Any`]
             Simple dict of all properties.
+
+        Notes
+        -----
+        Round-tripping of extension properties requires that the
+        `ObservationInfo` was created with the help of a registered
+        `MetadataTranslator` (which contains the extension property
+        definitions).
         """
         simple = {}
+        if hasattr(self, "_translator") and self._translator and self._translator.name:
+            simple["_translator"] = self._translator.name
 
         for p in self.all_properties:
-            property = f"_{p}"
+            property = f"_{p}" if not p.startswith("ext_") else p
             value = getattr(self, property)
             if value is None:
                 continue
@@ -439,11 +459,18 @@ def to_json(self):
         -------
         j : `str`
             The properties of the ObservationInfo in JSON string form.
+
+        Notes
+        -----
+        Round-tripping of extension properties requires that the
+        `ObservationInfo` was created with the help of a registered
+        `MetadataTranslator` (which contains the extension property
+        definitions).
         """
         return json.dumps(self.to_simple())
 
     @classmethod
-    def from_simple(cls, simple, extensions=None):
+    def from_simple(cls, simple):
         """Convert the entity returned by `to_simple` back into an
         `ObservationInfo`.
 
@@ -456,7 +483,21 @@ def from_simple(cls, simple, extensions=None):
         -------
         obsinfo : `ObservationInfo`
             New object constructed from the dict.
+
+        Notes
+        -----
+        Round-tripping of extension properties requires that the
+        `ObservationInfo` was created with the help of a registered
+        `MetadataTranslator` (which contains the extension property
+        definitions).
         """
+        extensions = {}
+        translator = simple.pop("_translator", None)
+        if translator:
+            if translator not in MetadataTranslator.translators:
+                raise KeyError(f"Unrecognised translator: {translator}")
+            extensions = MetadataTranslator.translators[translator].extensions
+
         properties = cls._get_all_properties(extensions)
 
         processed = {}
@@ -488,6 +529,13 @@ def from_json(cls, json_str):
         -------
         obsinfo : `ObservationInfo`
             Reconstructed object.
+
+        Notes
+        -----
+        Round-tripping of extension properties requires that the
+        `ObservationInfo` was created with the help of a registered
+        `MetadataTranslator` (which contains the extension property
+        definitions).
         """
         simple = json.loads(json_str)
         return cls.from_simple(simple)

python/astro_metadata_translator/translator.py

@@ -1069,6 +1069,8 @@ def to_property(self):
 # Assigning to __abstractmethods__ directly does work but interacts
 # poorly with the metaclass automatically generating methods from
 # _trivialMap and _constMap.
+# Note that subclasses that provide extension properties are assumed to not
+# need abstract methods created for them.
 
 # Allow for concrete translator methods to exist in the base class
 # These translator methods can be defined in terms of other properties

tests/test_extensions.py

@@ -11,45 +11,92 @@
 
 import pickle
 import unittest
-from astro_metadata_translator import ObservationInfo, makeObservationInfo, PropertyDefinition
+from astro_metadata_translator import (
+    ObservationInfo,
+    StubTranslator,
+    makeObservationInfo,
+    PropertyDefinition
+)
 
 """Test that extensions to the core set of properties works"""
 
+FOO = "bar"
+NUMBER = 12345
+
+
+class DummyTranslator(StubTranslator):
+    name = "dummy"
+    supported_instrument = "dummy"
+    extensions = dict(
+        number=PropertyDefinition("A number", "int", int),
+        foo=PropertyDefinition("A string", "str", str),
+    )
+    _const_map = {
+        # "observation_id": 1,
+        # "telescope": "dummy",
+        "ext_foo": FOO,
+    }
+
+    @classmethod
+    def can_translate(cls, header, filename=None):
+        return "INSTRUME" in header and header["INSTRUME"] == "dummy"
+
+    def to_ext_number(self):
+        """Return the combination on my luggage"""
+        return NUMBER
+
 
 class ExtensionsTestCase(unittest.TestCase):
     def setUp(self):
-        self.extensions = dict(
-            myInfo=PropertyDefinition("My special information", "str", str),
-        )
+        self.header = dict(INSTRUME="dummy")
+        self.obsinfo = ObservationInfo(self.header)
+
+    def assert_observation_info(self, obsinfo):
+        """Check that the `ObservationInfo` is as expected"""
+        self.assertIsInstance(obsinfo, ObservationInfo)
+        self.assertEqual(obsinfo.ext_foo, FOO)
+        self.assertEqual(obsinfo.ext_number, NUMBER)
 
     def test_basic(self):
         """Test construction of extended ObservationInfo"""
-        string = "the usual"
-        obsinfo = makeObservationInfo(extensions=self.extensions, ext_myInfo=string)
-
         # Behaves like the original
-        self.assertIsInstance(obsinfo, ObservationInfo)
-        self.assertEqual(obsinfo.ext_myInfo, string)
+        self.assert_observation_info(self.obsinfo)
+
+        copy = makeObservationInfo(extensions=DummyTranslator.extensions, ext_foo=FOO, ext_number=NUMBER)
+        self.assertEqual(copy, self.obsinfo)
 
         with self.assertRaises(AttributeError):
             # Variable is read-only
-            obsinfo.ext_myInfo = "something completely different"
+            self.obsinfo.ext_foo = "something completely different"
 
         with self.assertRaises(KeyError):
             # Can't specify extension value without declaring extensions
-            obsinfo.makeObservationInfo(ext_myInfo="foobar")
+            makeObservationInfo(foo="foobar")
 
         with self.assertRaises(TypeError):
             # Type checking is applied, like in the original
-            obsinfo.makeObservationInfo(extensions=self.extensions, ext_myInfo=12345)
+            makeObservationInfo(extensions=DummyTranslator.extensions, ext_foo=98765)
 
     def test_pickle(self):
         """Test that pickling works on ObservationInfo with extensions"""
-        string = "foobar"
-        obsinfo = makeObservationInfo(extensions=self.extensions, ext_myInfo=string)
-        copy = pickle.loads(pickle.dumps(obsinfo))
-        self.assertIsInstance(copy, ObservationInfo)
-        self.assertEqual(obsinfo.ext_myInfo, string)
+        obsinfo = pickle.loads(pickle.dumps(self.obsinfo))
+        self.assert_observation_info(obsinfo)
+
+    def test_simple(self):
+        """Test that simple representation works"""
+        simple = self.obsinfo.to_simple()
+        self.assertIn("ext_foo", simple)
+        self.assertIn("ext_number", simple)
+        obsinfo = ObservationInfo.from_simple(simple)
+        self.assert_observation_info(obsinfo)
+
+    def test_json(self):
+        """Test that JSON representation works"""
+        json = self.obsinfo.to_json()
+        self.assertIn("ext_foo", json)
+        self.assertIn("ext_number", json)
+        obsinfo = ObservationInfo.from_json(json)
+        self.assert_observation_info(obsinfo)
 
 
 if __name__ == "__main__":