Source code for SDF.data_model.instrument

from typing import Optional, Union, Iterable
from xml.etree.ElementTree import Element

from SDF.data_model._helper_functions import pop_child_element, element_is_empty, pop_all_child_elements
from SDF.data_model.abstract import XMLWritable
from SDF.data_model.element_set import ElementSet
from SDF.data_model.name import NameElement
from SDF.data_model.parameter import AnonymousParameterSet, ParameterType


class Instrument(AnonymousParameterSet, XMLWritable):
    """Represents an SDF <instrument> element"""

    def __init__(self, name: Optional[Union[str, NameElement]] = None,
                 parameters: Optional[Iterable[ParameterType]] = None):
        if isinstance(name, NameElement):
            self.__name = name
        else:
            self.__name = NameElement(name)

        super().__init__(items=parameters)

    @property
    def name(self) -> str:
        return self.__name.name

    def to_xml_element(self) -> Element:
        element = Element("instrument")
        element.append(self.__name.to_xml_element())
        for parameter in self:
            element.append(parameter.to_xml_element())
        return element

    @classmethod
    def from_xml_element(cls, element: Element) -> "Instrument":
        if element.tag != "instrument":
            raise ValueError(f"Expected an <instrument> element, got {element.tag}")

        name = NameElement.from_xml_element(pop_child_element(element, 0))

        parameters = []
        for child in pop_all_child_elements(element):
            if child.tag == "par":
                parameters.append(ParameterType.from_xml_element(child))
            else:
                raise ValueError(f"Unexpected child of type '{child.tag}'")

        if not element_is_empty(element):
            raise ValueError("Element is not empty")

        return cls(name, parameters)

    def __repr__(self):
        return f"{self.__class__.__name__}({self.name!r}, {[item for item in self]!r})"

    def __eq__(self, other):
        if isinstance(other, Instrument):
            return self.__name == other.__name and super().__eq__(other)
        return False


class InstrumentSet(ElementSet[Instrument]):
    """Extends `ElementSet[Instrument]` to allow adding parameters more intuitively"""

    def __init__(self, items: Optional[Iterable[Instrument]] = None):
        super().__init__(key_func=lambda inst: inst.name, check_func=lambda inst: isinstance(inst, Instrument),
                         items=items)

    def __setitem__(self, key, value):
        if key not in self.keys():
            self.add(Instrument(name=key))
        self[key].add_from_structure(value)