Skip to content

Commit db56ff6

Browse files
kangqiwangnstarman
kangqiwang
authored and
nstarman
committed
update mrt based on #14871 
1 parent 21810cf commit db56ff6

4 files changed

Lines changed: 333 additions & 2 deletions

File tree

‎astropy/cosmology/_src/io/builtin/__init__.py‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -5,4 +5,4 @@
55
"""
66

77
# Import to register with the I/O machinery
8-
from . import cosmology, ecsv, html, latex, mapping, model, row, table, yaml
8+
from . import cosmology, ecsv, html, latex, mapping, model, mrt, row, table, yaml

‎astropy/cosmology/_src/io/builtin/mrt.py‎

Lines changed: 148 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,148 @@
1+
r"""|Cosmology| <-> MRT I/O, using |Cosmology.read| and |Cosmology.write|.
2+
3+
This module provides functions to write/read a |Cosmology| object to/from an MRT file.
4+
The functions are registered with ``readwrite_registry`` under the format name "ascii.mrt".
5+
6+
"""
7+
8+
from __future__ import annotations
9+
10+
import json
11+
from typing import TYPE_CHECKING, Any, TypeVar
12+
13+
import astropy.cosmology.units as cu
14+
import astropy.units as u
15+
from astropy.cosmology._src.core import Cosmology
16+
from astropy.cosmology._src.io.connect import readwrite_registry
17+
from astropy.io.typing import PathLike, ReadableFileLike, WriteableFileLike
18+
from astropy.table import QTable
19+
20+
from .table import from_table, to_table
21+
22+
if TYPE_CHECKING:
23+
from astropy.cosmology._src.typing import _CosmoT
24+
from astropy.io.typing import PathLike, ReadableFileLike, WriteableFileLike
25+
from astropy.table import Table
26+
27+
_TableT = TypeVar("_TableT", "Table")
28+
29+
30+
def read_mrt(
31+
file: PathLike | ReadableFileLike[Table],
32+
index: int | str | None = None,
33+
*,
34+
move_to_meta: bool = False,
35+
cosmology: str | type[_CosmoT] | None = None,
36+
**kwargs: Any,
37+
) -> _CosmoT:
38+
r"""Read a `~astropy.cosmology.Cosmology` from an MRT file.
39+
40+
Parameters
41+
----------
42+
43+
Examples
44+
--------
45+
We assume the following setup:
46+
47+
>>> from pathlib import Path
48+
>>> from tempfile import TemporaryDirectory
49+
>>> temp_dir = TemporaryDirectory()
50+
51+
Writing a cosmology to a LaTeX file will produce a table with the cosmology's type,
52+
name, and parameters as columns.
53+
54+
>>> from astropy.cosmology import Planck18
55+
>>> file = Path(temp_dir.name) / "file.mrt"
56+
57+
>>> Planck18.write(file, format="ascii.mrt")
58+
>>> with open(file) as f: print(f.read())
59+
60+
61+
62+
63+
"""
64+
format = kwargs.pop("format", "ascii.mrt")
65+
if format != "ascii.mrt":
66+
raise ValueError(f"format must be 'ascii.mrt',not {format}")
67+
68+
with u.add_enabled_units(cu):
69+
table = QTable.read(file, format="ascii.mrt", **kwargs)
70+
71+
# Build the cosmology from table, using the private backend.
72+
return from_table(
73+
table, index=index, move_to_meta=move_to_meta, cosmology=cosmology
74+
)
75+
76+
77+
def write_mrt(
78+
cosmology: Cosmology,
79+
file: PathLike | WriteableFileLike[_TableT],
80+
*,
81+
overwrite: bool = False,
82+
cls: type[_TableT] = QTable,
83+
cosmology_in_meta: bool = True,
84+
**kwargs: Any,
85+
):
86+
format = kwargs.pop("format", "ascii.mrt")
87+
if format != "ascii.mrt":
88+
raise ValueError(f"format must be 'ascii.mrt', not {format}")
89+
90+
table = to_table(cosmology, cls=cls, cosmology_in_meta=cosmology_in_meta)
91+
for k, col in table.columns.items():
92+
# CDS can't serialize redshift units, so remove them # TODO: fix this
93+
if col.unit is cu.redshift:
94+
table[k] <<= u.dimensionless_unscaled
95+
## check if col is mutil den
96+
if len(col.shape) > 1:
97+
98+
def format_col_item(idx):
99+
obj = col[idx]
100+
try:
101+
obj = obj.value.tolist() if hasattr(obj, "value") else obj.tolist()
102+
except AttributeError as exc:
103+
pass
104+
return json.dumps(obj, separators=(",", ":"))
105+
else:
106+
107+
def format_col_item(idx):
108+
return str(col[idx])
109+
110+
try:
111+
table[k] = [format_col_item(idx) for idx in range(len(col))]
112+
113+
except TypeError as exc:
114+
raise TypeError(
115+
f"could not convert column {col.info.name!r} to string: {exc}"
116+
) from exc
117+
118+
# Write MRT
119+
table.write(file, overwrite=overwrite, format="ascii.mrt", **kwargs)
120+
121+
122+
def mrt_identify(
123+
origin: object, filepath: object, *args: object, **kwargs: object
124+
) -> bool:
125+
"""Identify if an object uses the HTML Table format.
126+
127+
Parameters
128+
----------
129+
origin : object
130+
Not used.
131+
filepath : object
132+
From where to read the Cosmology.
133+
*args : object
134+
Not used.
135+
**kwargs : object
136+
Not used.
137+
138+
Returns
139+
-------
140+
bool
141+
If the filepath is a string ending with '.mrt'.
142+
"""
143+
return isinstance(filepath, str) and filepath.endswith(".mrt")
144+
145+
146+
readwrite_registry.register_reader("ascii.mrt", Cosmology, read_mrt)
147+
readwrite_registry.register_writer("ascii.mrt", Cosmology, write_mrt)
148+
readwrite_registry.register_identifier("ascii.mrt", Cosmology, mrt_identify)

‎astropy/cosmology/_src/tests/io/test_.py‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -12,7 +12,7 @@ def test_expected_readwrite_io():
1212
"""Test that ONLY the expected I/O is registered."""
1313

1414
got = {k for k, _ in readwrite_registry._readers.keys()}
15-
expected = {"ascii.ecsv", "ascii.html"}
15+
expected = {"ascii.ecsv", "ascii.html", "ascii.mrt"}
1616

1717
assert got == expected
1818

‎astropy/cosmology/_src/tests/io/test_mrt.py‎

Lines changed: 183 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,183 @@
1+
# Licensed under a 3-clause BSD style license - see LICENSE.rst
2+
3+
# THIRD PARTY
4+
5+
import pytest
6+
7+
from astropy.cosmology._src.io.builtin.mrt import (
8+
read_mrt,
9+
write_mrt,
10+
)
11+
from astropy.table import QTable, Table
12+
13+
from .base import ReadWriteDirectTestBase, ReadWriteTestMixinBase
14+
15+
16+
class ReadWriteMRTTestMixin(ReadWriteTestMixinBase):
17+
"""
18+
Tests for a Cosmology[Read/Write] with ``format="mrt"``.
19+
This class will not be directly called by :mod:`pytest` since its name does
20+
not begin with ``Test``. To activate the contained tests this class must
21+
be inherited in a subclass. Subclasses must define a :func:`pytest.fixture`
22+
``cosmo`` that returns/yields an instance of a |Cosmology|.
23+
See ``TestCosmology`` for an example.
24+
"""
25+
26+
def test_to_mrt_bad_index(self, read, write, tmp_path):
27+
"""Test if argument ``index`` is incorrect"""
28+
fp = tmp_path / "test_to_mrt_bad_index.mrt"
29+
30+
write(fp, format="ascii.mrt")
31+
32+
# single-row table and has a non-0/None index
33+
with pytest.raises(IndexError, match="index 2 out of range"):
34+
read(fp, index=2, format="ascii.mrt")
35+
36+
# string index where doesn't match
37+
with pytest.raises(KeyError, match="No matches found for key"):
38+
read(fp, index="row 0", format="ascii.mrt")
39+
40+
# -----------------------
41+
42+
def test_to_mrt_failed_cls(self, write, tmp_path):
43+
"""Test failed table type."""
44+
fp = tmp_path / "test_to_mrt_failed_cls.mrt"
45+
46+
with pytest.raises(TypeError, match="'cls' must be"):
47+
write(fp, format="ascii.mrt", cls=list)
48+
49+
# -----------------------
50+
51+
@pytest.mark.parametrize("tbl_cls", [QTable, Table])
52+
def test_to_mrt_cls(self, write, tbl_cls, tmp_path):
53+
fp = tmp_path / "test_to_mrt_cls.mrt"
54+
write(fp, format="ascii.mrt", cls=tbl_cls)
55+
56+
# -----------------------
57+
58+
def test_readwrite_mrt_instance(self, cosmo_cls, cosmo, read, write, tmp_path):
59+
"""Test cosmology -> ascii.mrt -> cosmology."""
60+
fp = tmp_path / "test_readwrite_mrt_instance.mrt"
61+
62+
# ------------
63+
# To Table
64+
65+
write(fp, format="ascii.mrt")
66+
67+
# some checks on the saved file
68+
tbl = Table.read(fp, format="ascii.mrt")
69+
assert tbl["name"] == cosmo.name
70+
71+
# ------------
72+
# From Table
73+
74+
tbl["mismatching"] = "will error"
75+
tbl.write(fp, format="ascii.mrt", overwrite=True)
76+
77+
# tests are different if the last argument is a **kwarg
78+
if tuple(cosmo._init_signature.parameters.values())[-1].kind == 4:
79+
got = read(fp, format="ascii.mrt")
80+
81+
assert got.__class__ is cosmo_cls
82+
assert got.name == cosmo.name
83+
assert "mismatching" not in got.meta
84+
85+
return # don't continue testing
86+
87+
# read with mismatching parameters errors
88+
with pytest.raises(TypeError, match="there are unused parameters"):
89+
read(fp, format="ascii.mrt")
90+
91+
# unless mismatched are moved to meta
92+
got = read(fp, format="ascii.mrt", move_to_meta=True)
93+
assert got == cosmo
94+
assert got.meta["mismatching"] == "will error"
95+
96+
# it won't error if everything matches up
97+
tbl.remove_column("mismatching")
98+
tbl.write(fp, format="ascii.mrt", overwrite=True)
99+
got = read(fp, format="ascii.mrt")
100+
assert got == cosmo
101+
102+
# -----------------------
103+
104+
def test_readwrite_mrt_subclass_partial_info(
105+
self, cosmo_cls, cosmo, read, write, tmp_path
106+
):
107+
"""
108+
Test writing from an instance and reading from that class.
109+
This works with missing information.
110+
"""
111+
fp = tmp_path / "test_read_mrt_subclass_partial_info.mrt"
112+
113+
# test write
114+
write(fp, format="ascii.mrt")
115+
116+
# partial information
117+
tbl = Table.read(fp, format="ascii.mrt")
118+
del tbl["Tcmb0"]
119+
tbl.write(fp, format="ascii.mrt", overwrite=True)
120+
121+
# read with the same class that wrote fills in the missing info with
122+
# the default value
123+
got = cosmo_cls.read(fp, format="ascii.mrt")
124+
got2 = read(fp, format="ascii.mrt", cosmology=cosmo_cls)
125+
126+
assert got == got2 # internal consistency
127+
128+
# not equal, because Tcmb0 is changed, which also changes m_nu
129+
assert got != cosmo
130+
assert got.Tcmb0 == cosmo_cls._init_signature.parameters["Tcmb0"].default
131+
assert got.clone(name=cosmo.name, Tcmb0=cosmo.Tcmb0, m_nu=cosmo.m_nu) == cosmo
132+
133+
134+
class WriteMRTTestMixin(ReadWriteTestMixinBase):
135+
def test_to_mrt_failed_cls(self, write, tmp_path):
136+
"""Test failed table type."""
137+
fp = tmp_path / "test_to_mrt_failed_cls.mrt"
138+
139+
with pytest.raises(TypeError, match="'cls' must be"):
140+
write(fp, format="ascii.mrt", cls=list)
141+
142+
@pytest.mark.parametrize("tbl_cls", [QTable, Table])
143+
def test_to_mrt_cls(self, write, tbl_cls, tmp_path):
144+
fp = tmp_path / "test_to_mrt_cls.mrt"
145+
write(fp, format="ascii.mrt", cls=tbl_cls)
146+
147+
def test_to_mrt_bad_index(self, read, write, tmp_path):
148+
"""Test if argument ``index`` is incorrect"""
149+
fp = tmp_path / "test_to_mrt_bad_index.mrt"
150+
151+
write(fp, format="ascii.mrt")
152+
153+
# single-row table and has a non-0/None index
154+
with pytest.raises(IndexError, match="index 2 out of range"):
155+
read(fp, index=2, format="ascii.mrt")
156+
157+
# string index where doesn't match
158+
with pytest.raises(KeyError, match="No matches found for key"):
159+
read(fp, index="row 0", format="ascii.mrt")
160+
161+
def test_write_mrt_invalid_path(self, write):
162+
"""Test passing an invalid path"""
163+
invalid_fp = ""
164+
with pytest.raises(FileNotFoundError, match="No such file or directory"):
165+
write(invalid_fp, format="ascii.mrt")
166+
167+
def test_readwrite_mrt_instance(self, cosmo_cls, cosmo, read, write, tmp_path):
168+
fp = tmp_path / "test_readwrite_mrt_instance.mrt"
169+
write(fp, format="ascii.mrt")
170+
tb1 = Table.read(fp, format="ascii.mrt")
171+
assert tb1["name"] == cosmo.name
172+
173+
174+
class TestReadWriteMRT(ReadWriteDirectTestBase, WriteMRTTestMixin):
175+
"""
176+
Directly test ``read/write_mrt``.
177+
These are not public API and are discouraged from use, in favor of
178+
``Cosmology.read/write(..., format="mrt")``, but should be
179+
tested regardless b/c they are used internally.
180+
"""
181+
182+
def setup_class(self):
183+
self.functions = {"read": read_mrt, "write": write_mrt}

0 commit comments

Comments
 (0)