Source code for SDF.data_model.sample

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
from SDF.data_model.abstract import XMLWritable
from SDF.data_model.comment import Comment
from SDF.data_model.element_set import ElementSet
from SDF.data_model.name import NameElement


class Sample(XMLWritable):
    """Represents an SDF <sample> element"""

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

        if comment is None:
            self.__comment = None
        elif isinstance(comment, Comment):
            self.__comment = comment
        else:
            self.__comment = Comment(comment)

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

    @property
    def comment(self) -> Optional[str]:
        if self.__comment is None:
            return None
        return self.__comment.comment

    @comment.setter
    def comment(self, _comment: Optional[Union[str, Comment]]) -> None:
        if _comment is None:
            self.__comment = None
        elif isinstance(_comment, Comment):
            self.__comment = _comment
        else:
            self.__comment = Comment(_comment)

    def to_xml_element(self) -> Element:
        element = Element("sample")
        element.append(self.__name.to_xml_element())
        if self.__comment is not None:
            element.append(self.__comment.to_xml_element())
        return element

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

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

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

        return cls(name=name, comment=comment)

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

    def __repr__(self):
        return f"Sample({self.name!r}, {self.comment!r})"


class SampleSet(ElementSet[Sample]):
    """Behaves like `Dict[str, str]`, where the keys are sample names and the values are the associated comments"""

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

    def __getitem__(self, key) -> str:
        sample = super().__getitem__(key)
        return sample.comment

    def __setitem__(self, key, value):
        try:
            sample = self[key]
        except KeyError:
            sample = Sample(name=key)
            self.add(sample)
        sample.comment = value