Source code for SDF.data_model.array_dataset_2d

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

import numpy as np

from SDF.data_model._helper_functions import pop_child_element, element_is_empty, pop_element_attribute
from SDF.data_model.array_data_2d import ArrayData2D
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.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 ArrayDataset2D(Dataset):
    """Represents an SDF <dataset> element"""

    def __init__(self, name: Union[str, NameElement],
                 data: Union[ArrayData2D, np.ndarray], *,
                 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, ArrayData2D):
            self.__data = data
        elif isinstance(data, np.ndarray):
            self.__data = ArrayData2D(data=data)
        else:
            raise TypeError(f"Expected ArrayData1D or np.ndarray, got {type(data)}")

    @property
    def data(self) -> np.ndarray:
        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) -> "ArrayDataset2D":
        if element.tag != "dataset":
            raise ValueError(f"Expected a <dataset> element, got {element.tag}")
        ds_type = pop_element_attribute(element, "type")
        if ds_type != "mc":
            raise ValueError(f"Expected type 'mc', got {ds_type}")

        name = NameElement.from_xml_element(pop_child_element(element, 0))
        data = ArrayData2D.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, ArrayDataset2D):
            return super().__eq__(other) and self.__data == other.__data
        return False