From 8bc12593ad1356afd33b9aed57f672d53980d262 Mon Sep 17 00:00:00 2001
From: Nathaniel Starkman <nstarman@users.noreply.github.com>
Date: Wed, 17 May 2023 18:38:53 -0400
Subject: [PATCH] Add rename to ECSV

Signed-off-by: Nathaniel Starkman <nstarman@users.noreply.github.com>
---
 astropy/cosmology/io/ecsv.py            | 39 ++++++++++++++++++-------
 astropy/cosmology/io/tests/test_ecsv.py | 23 +++++++++++++++
 2 files changed, 51 insertions(+), 11 deletions(-)

diff --git a/astropy/cosmology/io/ecsv.py b/astropy/cosmology/io/ecsv.py
index 06603a95ed82..fb56f631e299 100644
--- a/astropy/cosmology/io/ecsv.py
+++ b/astropy/cosmology/io/ecsv.py
@@ -9,8 +9,10 @@
 from .table import from_table, to_table
 
 
-def read_ecsv(filename, index=None, *, move_to_meta=False, cosmology=None, **kwargs):
-    """Read a `~astropy.cosmology.Cosmology` from an ECSV file.
+def read_ecsv(
+    filename, index=None, *, move_to_meta=False, cosmology=None, rename=None, **kwargs
+):
+    """Read a |Cosmology| from an ECSV file.
 
     Parameters
     ----------
@@ -31,10 +33,8 @@ def read_ecsv(filename, index=None, *, move_to_meta=False, cosmology=None, **kwa
         (e.g. for ``Cosmology(meta={'key':10}, key=42)``, the ``Cosmology.meta``
         will be ``{'key': 10}``).
 
-    cosmology : str, `~astropy.cosmology.Cosmology` class, or None (optional, keyword-only)
-        The cosmology class (or string name thereof) to use when constructing
-        the cosmology instance. The class also provides default parameter values,
-        filling in any non-mandatory arguments missing in 'table'.
+    rename : dict or None (optional keyword-only)
+        A dictionary mapping column names to fields of the Cosmology.
 
     **kwargs
         Passed to :attr:`astropy.table.QTable.read`
@@ -49,18 +49,30 @@ def read_ecsv(filename, index=None, *, move_to_meta=False, cosmology=None, **kwa
 
     # build cosmology from table
     return from_table(
-        table, index=index, move_to_meta=move_to_meta, cosmology=cosmology
+        table,
+        index=index,
+        move_to_meta=move_to_meta,
+        cosmology=cosmology,
+        rename=rename,
     )
 
 
 def write_ecsv(
-    cosmology, file, *, overwrite=False, cls=QTable, cosmology_in_meta=True, **kwargs
+    cosmology,
+    file,
+    *,
+    overwrite=False,
+    cls=QTable,
+    cosmology_in_meta=True,
+    rename=None,
+    **kwargs
 ):
     """Serialize the cosmology into a ECSV.
 
     Parameters
     ----------
-    cosmology : `~astropy.cosmology.Cosmology` subclass instance
+    cosmology : |Cosmology|
+        The cosmology instance to convert to a mapping.
     file : path-like or file-like
         Location to save the serialized cosmology.
 
@@ -69,9 +81,12 @@ def write_ecsv(
     cls : type (optional, keyword-only)
         Astropy :class:`~astropy.table.Table` (sub)class to use when writing.
         Default is :class:`~astropy.table.QTable`.
-    cosmology_in_meta : bool
+    cosmology_in_meta : bool (optional, keyword-only)
         Whether to put the cosmology class in the Table metadata (if `True`,
         default) or as the first column (if `False`).
+    rename : dict or None (optional keyword-only)
+        A dictionary mapping fields of the Cosmology to columns of the table.
+
     **kwargs
         Passed to ``cls.write``
 
@@ -80,7 +95,9 @@ def write_ecsv(
     TypeError
         If kwarg (optional) 'cls' is not a subclass of `astropy.table.Table`
     """
-    table = to_table(cosmology, cls=cls, cosmology_in_meta=cosmology_in_meta)
+    table = to_table(
+        cosmology, cls=cls, cosmology_in_meta=cosmology_in_meta, rename=rename
+    )
 
     kwargs["format"] = "ascii.ecsv"
     table.write(file, overwrite=overwrite, **kwargs)
diff --git a/astropy/cosmology/io/tests/test_ecsv.py b/astropy/cosmology/io/tests/test_ecsv.py
index 1c5274c25c55..9c2c703996b2 100644
--- a/astropy/cosmology/io/tests/test_ecsv.py
+++ b/astropy/cosmology/io/tests/test_ecsv.py
@@ -127,6 +127,29 @@ def test_readwrite_ecsv_instance(
         got = read(fp)
         assert got == cosmo
 
+    def test_readwrite_ecsv_rename(
+        self, cosmo_cls, cosmo, read, write, tmp_path, add_cu
+    ):
+        """Test rename argument to read/write."""
+        fp = tmp_path / "test_readwrite_ecsv_rename.ecsv"
+        rename = {"name": "cosmo_name"}
+
+        write(fp, format="ascii.ecsv", rename=rename)
+
+        tbl = QTable.read(fp, format="ascii.ecsv")
+
+        assert "name" not in tbl.colnames
+        assert "cosmo_name" in tbl.colnames
+
+        # Errors if reading
+        with pytest.raises(TypeError, match="there are unused parameters"):
+            read(fp)
+
+        # Roundtrips
+        inv_rename = {v: k for k, v in rename.items()}
+        got = read(fp, rename=inv_rename)
+        assert got == cosmo
+
     def test_readwrite_ecsv_subclass_partial_info(
         self, cosmo_cls, cosmo, read, write, tmp_path, add_cu
     ):