Override extension type in ASDF

override_compound.py

#!/usr/bin/env python3
"""A simple example demonstrating how a compound model can be overridden"""

import sys

from astropy.modeling.models import Polynomial1D
from astropy.modeling.core import _CompoundModel

import asdf
from asdf import CustomType
from asdf.extension import BuiltinExtension
from asdf.yamlutil import tagged_tree_to_custom_tree

class NewCompound:
    def __init__(self, op1, op2):
        self.op1 = op1
        self.op2 = op2

class NewCompoundType(CustomType):
    """In practice this tag class would be extended to cover all compound
    types, but this is just a simple example."""

    name = 'transform/add'
    version = '1.1.0'
    types = [NewCompound]

    @classmethod
    def from_tree(cls, node, ctx):
        forward = node['forward']
        op1 = tagged_tree_to_custom_tree(forward[0], ctx)
        op2 = tagged_tree_to_custom_tree(forward[1], ctx)
        return NewCompound(op1, op2)

class NewExtension(BuiltinExtension):
    """We inherit from ASDF's BuiltinExtension here since NewCompoundType does
    not require any new schemas. We simply use the schemas that are provided by
    ASDF, and we override the type we care about."""

    @property
    def types(self):
        return [NewCompoundType]


def main():

    p1 = Polynomial1D(degree=3)
    p2 = Polynomial1D(degree=5)
    compound = p1 + p2

    tree = dict(p1=p1, p2=p2, compound=compound)

    # Create the file without using the new extension
    with asdf.AsdfFile(tree) as af:
        af.write_to('compound.asdf')

    # Open the file without using the new extension
    with asdf.open('compound.asdf') as old_af:
        assert isinstance(old_af.tree['p1'], Polynomial1D)
        assert isinstance(old_af.tree['p2'], Polynomial1D)
        print("without extension:", old_af.tree['compound'])
        assert isinstance(old_af.tree['compound'], _CompoundModel)

    # Open the file using the new extension
    with asdf.open('compound.asdf', extensions=NewExtension()) as new_af:
        # Make sure we didn't override the original Polynomial1D models
        assert isinstance(new_af.tree['p1'], Polynomial1D)
        assert isinstance(new_af.tree['p2'], Polynomial1D)

        # Make sure we did actually override the compound model
        print("with extension:", new_af.tree['compound'])
        assert isinstance(new_af.tree['compound'], NewCompound)
        assert isinstance(new_af.tree['compound'].op1, Polynomial1D)
        assert isinstance(new_af.tree['compound'].op2, Polynomial1D)

if __name__ == '__main__':
    sys.exit(main())