Source code for SDF.data_model.image_dataset

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

from PIL.Image import Image

from SDF.data_model._helper_functions import pop_child_element, element_is_empty, pop_element_attribute
from SDF.data_model.comment import Comment
from SDF.data_model.dataset import Dataset
from SDF.data_model.date import Date
from SDF.data_model.image_data import ImageData
from SDF.data_model.instrument import Instrument
from SDF.data_model.name import NameElement
from SDF.data_model.owner import Owner
from SDF.data_model.parameter import ParameterType
from SDF.data_model.sample import Sample


class ImageDataset(Dataset):
    """Represents an SDF <dataset> element"""

    def __init__(self, name: Union[str, NameElement],
                 data: Union[ImageData, Image], *,
                 owner: Optional[Union[str, Owner]] = None,
                 date: Optional[Union[datetime, Date]] = None,
                 comment: Optional[Union[str, Comment]] = None,
                 samples: Optional[Iterable[Sample]] = None,
                 parameters: Optional[Iterable[ParameterType]] = None,
                 instruments: Optional[Iterable[Instrument]] = None):
        super().__init__(name=name, owner=owner, date=date, comment=comment, samples=samples, parameters=parameters,
                         instruments=instruments)
        if isinstance(data, ImageData):
            self.__data = data
        elif isinstance(data, Image):
            self.__data = ImageData(data=data)
        else:
            raise TypeError(f"Expected ArrayData1D or np.ndarray, got {type(data)}")

    @property
    def data(self) -> Image:
        return self.__data.data

    def to_xml_element(self) -> Element:
        element = Element("dataset", dict(type=self.__data.type_for_xml))
        super()._to_partial_xml_element(element)
        element.append(self.__data.to_xml_element())
        return element

    @classmethod
    def from_xml_element(cls, element: Element) -> "ImageDataset":
        if element.tag != "dataset":
            raise ValueError(f"Expected a <dataset> element, got {element.tag}")
        ds_type = pop_element_attribute(element, "type")
        if ds_type != "img":
            raise ValueError(f"Expected type 'img', got {ds_type}")

        name = NameElement.from_xml_element(pop_child_element(element, 0))
        data = ImageData.from_xml_element(pop_child_element(element, -1))

        ret = cls(name=name, data=data)
        ret._from_partial_xml_element(element)

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

        return ret

    def __repr__(self):
        return f"{self.__class__.__name__}({self.name!r}, data={self.data!r}, owner={self.owner!r}, " \
               f"date={self.date!r}, comment={self.comment!r}, samples={self.samples!r}, " \
               f"parameters={self.parameters!r}, instruments={self.instruments!r})"

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