Source code for SDF.data_model.date

from datetime import datetime
from typing import Optional
from xml.etree.ElementTree import Element

from SDF.data_model._helper_functions import pop_element_attribute, pop_element_text, element_is_empty
from SDF.data_model.abstract import XMLWritable
from SDF.sdf_rc import DEFAULT_DATE_FORMAT


class Date(XMLWritable):
    """Wraps a datetime.datetime object. Represents an SDF <date> element."""

    def __init__(self, date: datetime, dateformat: Optional[str] = None):
        if not isinstance(date, datetime):
            raise TypeError(f"Expected a datetime object, got {type(date)}")
        self.date = date

        if dateformat is None:
            self.format = DEFAULT_DATE_FORMAT
        elif isinstance(dateformat, str):
            self.format = dateformat.strip()
        else:
            raise TypeError(f"Expected a string, got {type(dateformat)}")

    def to_xml_element(self) -> Element:
        element = Element("date")
        if self.format != DEFAULT_DATE_FORMAT:
            element.attrib["format"] = self.format
        element.text = self.date.strftime(self.format)
        return element

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

        dateformat = pop_element_attribute(element, "format") if "format" in element.attrib else DEFAULT_DATE_FORMAT
        text = pop_element_text(element)
        if not element_is_empty(element):
            raise ValueError("Element is not empty")
        return cls(datetime.strptime(text.strip(), dateformat), dateformat)

    def __repr__(self):
        return f"{self.__class__.__name__}({self.date!r}, {self.format!r})"

    def __eq__(self, other):
        if isinstance(other, Date):
            # milliseconds are not stored in XML, so we cannot always expect exact equality
            return (self.date.year == other.date.year
                    and self.date.month == other.date.month
                    and self.date.day == other.date.day
                    and self.date.hour == other.date.hour
                    and self.date.minute == other.date.minute)  # don't check seconds because of rounding
        return False