Merge pull request #12710 from nstarman/cosmo-add-arguments-to-to_map…

…ping

Add kwargs to ``to_mapping``

astropy/cosmology/io/mapping.py

@@ -116,7 +116,7 @@ class itself.
     return cosmology(*ba.args, **ba.kwargs)
 
 
-def to_mapping(cosmology, *args, cls=dict):
+def to_mapping(cosmology, *args, cls=dict, cosmology_as_str=False, move_from_meta=False):
     """Return the cosmology class, parameters, and metadata as a `dict`.
 
     Parameters
@@ -128,13 +128,22 @@ def to_mapping(cosmology, *args, cls=dict):
     cls : type (optional, keyword-only)
         `dict` or `collections.Mapping` subclass.
         The mapping type to return. Default is `dict`.
+    cosmology_as_str : bool (optional, keyword-only)
+        Whether the cosmology value is the class (if `False`, default) or
+        the semi-qualified name (if `True`).
+    move_from_meta : bool (optional, keyword-only)
+        Whether to add the Cosmology's metadata as an item to the mapping (if
+        `False`, default) or to merge with the rest of the mapping, preferring
+        the original values (if `True`)
 
     Returns
     -------
     dict
         with key-values for the cosmology parameters and also:
         - 'cosmology' : the class
-        - 'meta' : the contents of the cosmology's metadata attribute
+        - 'meta' : the contents of the cosmology's metadata attribute.
+                   If ``move_from_meta`` is `True`, this key is missing and the
+                   contained metadata are added to the main `dict`.
 
     Examples
     --------
@@ -148,19 +157,58 @@ def to_mapping(cosmology, *args, cls=dict):
          'Tcmb0': <Quantity 2.7255 K>, 'Neff': 3.046,
          'm_nu': <Quantity [0.  , 0.  , 0.06] eV>, 'Ob0': 0.04897,
          'meta': ...
+
+    The dictionary type may be changed with the ``cls`` keyword argument:
+
+        >>> from collections import OrderedDict
+        >>> Planck18.to_format('mapping', cls=OrderedDict)
+        OrderedDict([('cosmology', <class 'astropy.cosmology.flrw.FlatLambdaCDM'>),
+          ('name', 'Planck18'), ('H0', <Quantity 67.66 km / (Mpc s)>),
+          ('Om0', 0.30966), ('Tcmb0', <Quantity 2.7255 K>), ('Neff', 3.046),
+          ('m_nu', <Quantity [0.  , 0.  , 0.06] eV>), ('Ob0', 0.04897),
+          ('meta', ...
+
+    Sometimes it is more useful to have the name of the cosmology class, not
+    the object itself. The keyword argument ``cosmology_as_str`` may be used:
+
+        >>> Planck18.to_format('mapping', cosmology_as_str=True)
+        {'cosmology': 'FlatLambdaCDM', ...
+
+    The metadata is normally included as a nested mapping. To move the metadata
+    into the main mapping, use the keyword argument ``move_from_meta``. This
+    kwarg inverts ``move_to_meta`` in
+    ``Cosmology.to_format("mapping", move_to_meta=...)`` where extra items
+    are moved to the metadata (if the cosmology constructor does not have a
+    variable keyword-only argument -- ``**kwargs``).
+
+        >>> from astropy.cosmology import Planck18
+        >>> Planck18.to_format('mapping', move_from_meta=True)
+        {'cosmology': <class 'astropy.cosmology.flrw.FlatLambdaCDM'>,
+         'name': 'Planck18', 'Oc0': 0.2607, 'n': 0.9665, 'sigma8': 0.8102, ...
     """
     if not issubclass(cls, (dict, Mapping)):
-        raise TypeError(f"'cls' must be a (sub)class of dict or Mapping, not {type(cls)}")
+        raise TypeError(f"'cls' must be a (sub)class of dict or Mapping, not {cls}")
 
     m = cls()
     # start with the cosmology class & name
-    m["cosmology"] = cosmology.__class__
+    m["cosmology"] = cosmology.__class__.__qualname__ if cosmology_as_str else cosmology.__class__
     m["name"] = cosmology.name  # here only for dict ordering
-    # get all the immutable inputs
+
+    meta = copy.deepcopy(cosmology.meta)  # metadata (mutable)
+    if move_from_meta:
+        # Merge the mutable metadata. Since params are added later they will
+        # be preferred in cases of overlapping keys. Likewise, need to pop
+        # cosmology and name from meta.
+        meta.pop("cosmology", None)
+        meta.pop("name", None)
+        m.update(meta)
+
+    # Add all the immutable inputs
     m.update({k: v for k, v in cosmology._init_arguments.items()
               if k not in ("meta", "name")})
-    # add the mutable metadata
-    m["meta"] = copy.deepcopy(cosmology.meta)
+    # Lastly, add the metadata, if haven't already (above)
+    if not move_from_meta:
+        m["meta"] = meta  # TODO? should meta be type(cls)
 
     return m
 

astropy/cosmology/io/tests/test_mapping.py

@@ -2,9 +2,11 @@
 
 # STDLIB
 import copy
+import inspect
 from collections import OrderedDict
 
 # THIRD PARTY
+import numpy as np
 import pytest
 
 # LOCAL
@@ -34,69 +36,126 @@ class ToFromMappingTestMixin(IOTestMixinBase):
     See ``TestCosmology`` for an example.
     """
 
-    def test_failed_cls_to_mapping(self, cosmo, to_format):
+    def test_to_mapping_default(self, cosmo, to_format):
+        """Test default usage of Cosmology -> mapping."""
+        m = to_format('mapping')
+        keys = tuple(m.keys())
+
+        assert isinstance(m, dict)
+        # Check equality of all expected items
+        assert keys[0] == "cosmology"
+        assert m.pop("cosmology") is cosmo.__class__
+        assert keys[1] == "name"
+        assert m.pop("name") == cosmo.name
+        for i, k in enumerate(cosmo.__parameters__, start=2):
+            assert keys[i] == k
+            assert np.array_equal(m.pop(k), getattr(cosmo, k))
+        assert keys[-1] == "meta"
+        assert m.pop("meta") == cosmo.meta
+
+        # No unexpected items
+        assert not m
+
+    def test_to_mapping_wrong_cls(self, to_format):
         """Test incorrect argument ``cls`` in ``to_mapping()``."""
         with pytest.raises(TypeError, match="'cls' must be"):
             to_format('mapping', cls=list)
 
     @pytest.mark.parametrize("map_cls", [dict, OrderedDict])
-    def test_to_mapping_cls(self, cosmo, to_format, map_cls):
+    def test_to_mapping_cls(self, to_format, map_cls):
         """Test argument ``cls`` in ``to_mapping()``."""
-        params = to_format('mapping', cls=map_cls)
-        assert isinstance(params, map_cls)  # test type
+        m = to_format('mapping', cls=map_cls)
+        assert isinstance(m, map_cls)  # test type
 
-    def test_from_not_mapping(self, cosmo, from_format):
-        """Test incorrect map type in ``from_mapping()``."""
-        with pytest.raises((TypeError, ValueError)):
-            from_format("NOT A MAP", format="mapping")
+    def test_to_mapping_cosmology_as_str(self, cosmo_cls, to_format):
+        """Test argument ``cosmology_as_str`` in ``to_mapping()``."""
+        default = to_format('mapping')
+
+        # Cosmology is the class
+        m = to_format('mapping', cosmology_as_str=False)
+        assert inspect.isclass(m["cosmology"])
+        assert cosmo_cls is m["cosmology"]
 
-    def test_tofrom_mapping_instance(self, cosmo, to_format, from_format):
-        """Test cosmology -> Mapping -> cosmology."""
-        # ------------
-        # To Mapping
+        assert m == default  # False is the default option
+
+        # Cosmology is a string
+        m = to_format('mapping', cosmology_as_str=True)
+        assert isinstance(m["cosmology"], str)
+        assert m["cosmology"] == cosmo_cls.__qualname__  # Correct class
+        assert tuple(m.keys())[0] == "cosmology"  # Stayed at same index
+
+    def test_tofrom_mapping_cosmology_as_str(self, cosmo, to_format, from_format):
+        """Test roundtrip with ``cosmology_as_str=True``.
+
+        The test for the default option (`False`) is in ``test_tofrom_mapping_instance``.
+        """
+        m = to_format('mapping', cosmology_as_str=True)
+
+        got = from_format(m, format="mapping")
+        assert got == cosmo
+        assert got.meta == cosmo.meta
 
-        params = to_format('mapping')
-        assert isinstance(params, dict)  # test type
-        assert params["cosmology"] is cosmo.__class__
-        assert params["name"] == cosmo.name
+    def test_to_mapping_move_from_meta(self, to_format):
+        """Test argument ``move_from_meta`` in ``to_mapping()``."""
+        default = to_format('mapping')
 
-        # ------------
-        # From Mapping
+        # Metadata is 'separate' from main mapping
+        m = to_format('mapping', move_from_meta=False)
+        assert "meta" in m.keys()
+        assert not any([k in m for k in m["meta"]])  # Not added to main
 
-        params["mismatching"] = "will error"
+        assert m == default  # False is the default option
 
-        # tests are different if the last argument is a **kwarg
+        # Metadata is mixed into main mapping.
+        m = to_format('mapping', move_from_meta=True)
+        assert "meta" not in m.keys()
+        assert all([k in m for k in default["meta"]])  # All added to main
+        #  The parameters take precedence over the metadata
+        assert all([np.array_equal(v, m[k]) for k, v in default.items() if k != "meta"])
+
+    def test_tofrom_mapping_move_tofrom_meta(self, cosmo, to_format, from_format):
+        """Test roundtrip of ``move_from/to_meta`` in ``to/from_mapping()``."""
+        # Metadata is mixed into main mapping.
+        m = to_format('mapping', move_from_meta=True)
+        # (Just adding something to ensure there's 'metadata')
+        m["mismatching"] = "will error"
+
+        # (Tests are different if the last argument is a **kwarg)
         if tuple(cosmo._init_signature.parameters.values())[-1].kind == 4:
-            got = from_format(params, format="mapping")
+            got = from_format(m, format="mapping")
 
             assert got.name == cosmo.name
             assert "mismatching" not in got.meta
 
             return  # don't continue testing
 
-        # read with mismatching parameters errors
+        # Reading with mismatching parameters errors...
         with pytest.raises(TypeError, match="there are unused parameters"):
-            from_format(params, format="mapping")
+            from_format(m, format="mapping")
 
-        # unless mismatched are moved to meta
-        got = from_format(params, format="mapping", move_to_meta=True)
-        assert got == cosmo
+        # unless mismatched are moved to meta.
+        got = from_format(m, format="mapping", move_to_meta=True)
+        assert got == cosmo  # (Doesn't check metadata)
         assert got.meta["mismatching"] == "will error"
 
-        # it won't error if everything matches up
-        params.pop("mismatching")
-        got = from_format(params, format="mapping")
-        assert got == cosmo
-        assert got.meta == cosmo.meta
+    # -----------------------------------------------------
 
-        # and it will also work if the cosmology is a string
-        params["cosmology"] = params["cosmology"].__qualname__
-        got = from_format(params, format="mapping")
+    def test_from_not_mapping(self, cosmo, from_format):
+        """Test incorrect map type in ``from_mapping()``."""
+        with pytest.raises((TypeError, ValueError)):
+            from_format("NOT A MAP", format="mapping")
+
+    def test_from_mapping_default(self, cosmo, to_format, from_format):
+        """Test (cosmology -> Mapping) -> cosmology."""
+        m = to_format('mapping')
+
+        # Read from exactly as given.
+        got = from_format(m, format="mapping")
         assert got == cosmo
         assert got.meta == cosmo.meta
 
-        # also it auto-identifies 'format'
-        got = from_format(params)
+        # Reading auto-identifies 'format'
+        got = from_format(m)
         assert got == cosmo
         assert got.meta == cosmo.meta
 
@@ -105,9 +164,7 @@ def test_fromformat_subclass_partial_info_mapping(self, cosmo):
         Test writing from an instance and reading from that class.
         This works with missing information.
         """
-        # test to_format
         m = cosmo.to_format("mapping")
-        assert isinstance(m, dict)
 
         # partial information
         m.pop("cosmology", None)
@@ -134,3 +191,8 @@ class TestToFromMapping(IOFormatTestBase, ToFromMappingTestMixin):
 
     def setup_class(self):
         self.functions = {"to": to_mapping, "from": from_mapping}
+
+    @pytest.mark.skip("N/A")
+    def test_fromformat_subclass_partial_info_mapping(self):
+        """This test does not apply to the direct functions."""
+        pass

docs/changes/cosmology/12710.feature.rst

@@ -0,0 +1,5 @@
+For converting a cosmology to a mapping, two new boolean keyword arguments are
+added: ``cosmology_as_str`` for turning the the class reference to a string,
+instead of the class object itself, and ``move_from_meta`` to merge the
+metadata with the rest of the returned mapping instead of adding it as a
+nested dictionary.

docs/cosmology/io.rst

@@ -259,8 +259,7 @@ a ``*args`` to absorb unneeded information passed by
     >>> from astropy.table import QTable
 
     >>> def to_table_row(cosmology, *args):
-    ...     p = cosmology.to_format("mapping")
-    ...     p["cosmology"] = p["cosmology"].__qualname__  # as string
+    ...     p = cosmology.to_format("mapping", cosmology_as_str=True)
     ...     meta = p.pop("meta")
     ...     # package parameters into lists for Table parsing
     ...     params = {k: [v] for k, v in p.items()}
@@ -363,8 +362,7 @@ boolean flag "overwrite" to set behavior for existing files. Note that
 .. code-block:: python
 
     >>> def write_json(cosmology, file, *, overwrite=False, **kwargs):
-    ...    data = cosmology.to_format("mapping")  # start by turning into dict
-    ...    data["cosmology"] = data["cosmology"].__name__  # change class field to str
+    ...    data = cosmology.to_format("mapping", cosmology_as_str=True)  # start by turning into dict
     ...    # serialize Quantity
     ...    for k, v in data.items():
     ...        if isinstance(v, u.Quantity):