diff --git a/asdf_astropy/__init__.py b/asdf_astropy/__init__.py index 11d266e8..b655e1ee 100644 --- a/asdf_astropy/__init__.py +++ b/asdf_astropy/__init__.py @@ -4,3 +4,9 @@ from ._astropy_init import * # noqa: F403 # ---------------------------------------------------------------------------- + +def configure_core_table_support(asdf_config): + from .extensions import CORE_TABLE_EXTENSIONS + + for extension in CORE_TABLE_EXTENSIONS: + asdf_config.add_extension(extension) diff --git a/asdf_astropy/converters/table/table.py b/asdf_astropy/converters/table/table.py index 2e771f9e..42ac64c1 100644 --- a/asdf_astropy/converters/table/table.py +++ b/asdf_astropy/converters/table/table.py @@ -1,5 +1,6 @@ from asdf.extension import Converter from asdf.tags.core.ndarray import NDArrayType +from asdf.tagged import TaggedDict class ColumnConverter(Converter): @@ -45,18 +46,52 @@ def from_yaml_tree(self, node, tag, ctx): class AsdfTableConverter(Converter): - tags = ("tag:stsci.edu:asdf/core/table-*",) - types = () - - def to_yaml_tree(self, obj, tag, ctx): - msg = "astropy does not support writing astropy.table.Table with the ASDF table-1.0.0 tag" - raise NotImplementedError(msg) + tags = ( + "tag:stsci.edu:asdf/core/table-1.0.0", + "tag:stsci.edu:asdf/core/table-1.1.0", + ) + types = ( + "astropy.table.table.Table", + "astropy.table.table.QTable", + ) def from_yaml_tree(self, node, tag, ctx): from astropy.table import Table return Table(node["columns"], meta=node.get("meta")) + def to_yaml_tree(self, obj, tag, ctx): + from astropy.table import Column, QTable, Table + + if isinstance(obj, QTable): + obj = Table(obj) + + for column in obj.columns.values(): + if not isinstance(column, Column): + msg = f"The ASDF table converter does not support columns of type '{type(column)}'"; + raise NotImplementedError(msg) + + node = { + "columns": [c for c in obj.columns.values()] + } + + if obj.meta: + node["meta"] = obj.meta + + return node + + +class AsdfTableConverterReadOnly(AsdfTableConverter): + types = () + + def to_yaml_tree(self, obj, tag, ctx): + msg = ( + "The default configuration of asdf-astropy does not support writing " + f"astropy.table.Table with the {tag} tag. " + "Call asdf_astropy.configure_core_table_support(asdf.get_config()) to use the ASDF core table tags." + ) + raise NotImplementedError(msg) + class AstropyTableConverter(Converter): tags = ("tag:astropy.org:astropy/table/table-*",) diff --git a/asdf_astropy/converters/table/tests/test_table.py b/asdf_astropy/converters/table/tests/test_table.py index ad1c70e6..28e4de3b 100644 --- a/asdf_astropy/converters/table/tests/test_table.py +++ b/asdf_astropy/converters/table/tests/test_table.py @@ -1,6 +1,7 @@ import warnings import asdf +import asdf_astropy import astropy.units as u import numpy as np import pytest @@ -270,3 +271,22 @@ def test_asdf_table(): assert table["c"].name == "c" assert table["c"].description == "The target name" assert_array_equal(table["c"].data, np.array([b"d", b"e", b"f"])) + + +def test_write_asdf_table(tmp_path): + with asdf.config_context() as config: + asdf_astropy.configure_core_table_support(config) + + table = Table() + table["a"] = [1, 2] + table["b"] = ["x", "y"] + + helpers.assert_table_roundtrip(table, tmp_path) + + file_path = tmp_path / "test_tag.asdf" + with asdf.AsdfFile() as af: + af["table"] = table + af.write_to(file_path) + + with asdf.open(file_path, _force_raw_types=True) as af: + assert af["table"]._tag.startswith("tag:stsci.edu:asdf/core/table-") diff --git a/asdf_astropy/extensions.py b/asdf_astropy/extensions.py index 5f5dba55..90970896 100644 --- a/asdf_astropy/extensions.py +++ b/asdf_astropy/extensions.py @@ -14,7 +14,7 @@ from .converters.coordinates.sky_coord import SkyCoordConverter from .converters.coordinates.spectral_coord import SpectralCoordConverter from .converters.fits.fits import AsdfFitsConverter, AstropyFitsConverter -from .converters.table.table import AsdfTableConverter, AstropyTableConverter, ColumnConverter, NdarrayMixinConverter +from .converters.table.table import AsdfTableConverter, AsdfTableConverterReadOnly, AstropyTableConverter, ColumnConverter, NdarrayMixinConverter from .converters.time.time import TimeConverter from .converters.time.time_delta import TimeDeltaConverter from .converters.transform.compound import CompoundConverter @@ -516,11 +516,15 @@ QuantityConverter(), TimeConverter(), ColumnConverter(), - AsdfTableConverter(), + AsdfTableConverterReadOnly(), AsdfFitsConverter(), ] -UNIT_CONVETERS = [ +CORE_TABLE_CONVERTERS = [ + AsdfTableConverter(), +] + +UNIT_CONVERTERS = [ UnitConverter(), EquivalencyConverter(), MagUnitConverter(), @@ -542,8 +546,14 @@ CompoundManifestExtension( [ ManifestExtension.from_uri(u, converters=CORE_CONVERTERS), - ManifestExtension.from_uri("asdf://astropy.org/astropy/manifests/units-1.0.0", converters=UNIT_CONVETERS), + ManifestExtension.from_uri("asdf://astropy.org/astropy/manifests/units-1.0.0", converters=UNIT_CONVERTERS), ], ) for u in CORE_MANIFEST_URIS ] + + +CORE_TABLE_EXTENSIONS = [ + ManifestExtension.from_uri(u, converters=CORE_TABLE_CONVERTERS) + for u in CORE_MANIFEST_URIS +]