Prototype serializing Table/QTable with the ASDF core table tag

Author: eslavich

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)

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-*",)

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-")

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
+]