compound_converter.py

class CompoundConverter(TransformConverter):
    _tags_add = {
        "http://asdf-format.org/schemas/transform/add-2.0.0",
        "tag:stsci.edu:asdf/transform/add-1.2.0",
        "tag:stsci.edu:asdf/transform/add-1.1.0",
        "tag:stsci.edu:asdf/transform/add-1.0.0",
    }
    _tags_subtract = {
        "http://asdf-format.org/schemas/transform/subtract-2.0.0",
        "tag:stsci.edu:asdf/transform/subtract-1.2.0",
        "tag:stsci.edu:asdf/transform/subtract-1.1.0",
        "tag:stsci.edu:asdf/transform/subtract-1.0.0",
    }
    _tags_multiply = {
        "http://asdf-format.org/schemas/transform/multiply-2.0.0",
        "tag:stsci.edu:asdf/transform/multiply-1.2.0",
        "tag:stsci.edu:asdf/transform/multiply-1.1.0",
        "tag:stsci.edu:asdf/transform/multiply-1.0.0",
    }
    _tags_divide = {
        "http://asdf-format.org/schemas/transform/divide-2.0.0",
        "tag:stsci.edu:asdf/transform/divide-1.2.0",
        "tag:stsci.edu:asdf/transform/divide-1.1.0",
        "tag:stsci.edu:asdf/transform/divide-1.0.0",
    }
    _tags_power = {
        "http://asdf-format.org/schemas/transform/power-2.0.0",
        "tag:stsci.edu:asdf/transform/power-1.2.0",
        "tag:stsci.edu:asdf/transform/power-1.1.0",
        "tag:stsci.edu:asdf/transform/power-1.0.0",
    }
    _tags_compose = {
        "http://asdf-format.org/schemas/transform/compose-2.0.0",
        "tag:stsci.edu:asdf/transform/compose-1.2.0",
        "tag:stsci.edu:asdf/transform/compose-1.1.0",
        "tag:stsci.edu:asdf/transform/compose-1.0.0",
    }
    _tags_concatenate = {
        "http://asdf-format.org/schemas/transform/concatenate-2.0.0",
        "tag:stsci.edu:asdf/transform/concatenate-1.2.0",
        "tag:stsci.edu:asdf/transform/concatenate-1.1.0",
        "tag:stsci.edu:asdf/transform/concatenate-1.0.0",
    }
    _tags_fix_inputs = {
        "http://asdf-format.org/schemas/transform/fix_inputs-2.0.0",
        "tag:stsci.edu:asdf/transform/fix_inputs-1.2.0",
        "tag:stsci.edu:asdf/transform/fix_inputs-1.1.0",
    }

    tags = set().union(
        _tags_add,
        _tags_subtract,
        _tags_multiply,
        _tags_divide,
        _tags_power,
        _tags_compose,
        _tags_concatenate,
        _tags_fix_inputs,
    )

    types = {CompoundModel}

    _operator_to_tags = {
        '+':  _tags_add,
        '-':  _tags_subtract,
        '*':  _tags_multiply,
        '/':  _tags_divide,
        '**': _tags_power,
        '|':  _tags_compose,
        '&':  _tags_concatenate,
        'fix_inputs': _tags_fix_inputs,
    }
        
    _tag_to_method = {}
    _tag_to_method.update({t: "__add__" for t in _tags_add})
    _tag_to_method.update({t: "__sub__" for t in _tags_subtract})
    _tag_to_method.update({t: "__mul__" for t in _tags_multiply})
    _tag_to_method.update({t: "__truediv__" for t in _tags_divide})
    _tag_to_method.update({t: "__pow__" for t in _tags_power})
    _tag_to_method.update({t: "__or__" for t in _tags_compose})
    _tag_to_method.update({t: "__and__" for t in _tags_concatenate})
    _tag_to_method.update({t: "fix_inputs" for t in _tags_fix_inputs})

    def from_tree_transform(self, node, tag):
        oper = self._tag_to_method[tag]
        left = node['forward'][0]
        if not isinstance(left, Model):
            raise TypeError("Unknown model type '{0}'".format(
                node['forward'][0].__class__.__name__))
        right = node['forward'][1]
        if (not isinstance(right, Model) and
                not (oper == 'fix_inputs' and isinstance(right, dict))):
            raise TypeError("Unknown model type '{0}'".format(
                node['forward'][1].__class__.__name__))
        if oper == 'fix_inputs':
            right = dict(zip(right['keys'], right['values']))
            model = CompoundModel('fix_inputs', left, right)
        else:
            model = getattr(left, oper)(right)
        return model

    def to_tree_transform(self, model, tags):
        left = model.left

        if isinstance(model.right, dict):
            right = {
                'keys': list(model.right.keys()),
                'values': list(model.right.values())
            }
        else:
            right = model.right

        node = {
            'forward': [left, right]
        }

        tag = self._operator_to_tags[model.op].intersection(tags).pop()
        return tagged.tag_object(tag, node)