flexmeasures-main/flexmeasures/data/schemas/times.py

from __future__ import annotations

import json
from datetime import datetime, timedelta

from flask import current_app
from marshmallow import fields, Schema, validates_schema
from marshmallow.exceptions import ValidationError
import isodate
from isodate.isoerror import ISO8601Error
import pandas as pd

from flexmeasures.data.schemas.utils import FMValidationError, MarshmallowClickMixin


class DurationValidationError(FMValidationError):
    status = "INVALID_PERIOD"  # USEF error status


class DurationField(MarshmallowClickMixin, fields.Str):
    """Field that deserializes to a ISO8601 Duration
    and serializes back to a string."""

    def _deserialize(self, value, attr, obj, **kwargs) -> timedelta | isodate.Duration:
        """
        Use the isodate library to turn an ISO8601 string into a timedelta.
        For some non-obvious cases, it will become an isodate.Duration, see
        ground_from for more.
        This method throws a ValidationError if the string is not ISO norm.
        """
        try:
            duration_value = isodate.parse_duration(value)
        except ISO8601Error as iso_err:
            raise DurationValidationError(
                f"Cannot parse {value} as ISO8601 duration: {iso_err}"
            )
        if duration_value.seconds % 60 != 0 or duration_value.microseconds != 0:
            raise DurationValidationError(
                "FlexMeasures only support multiples of 1 minute."
            )
        return duration_value

    def _serialize(self, value, attr, data, **kwargs):
        """
        An implementation of _serialize.
        It is not guaranteed to return the same string as was input,
        if ground_from has been used!
        """
        return isodate.strftime(value, "P%P")

    @staticmethod
    def ground_from(
        duration: timedelta | isodate.Duration, start: datetime | None
    ) -> timedelta:
        """
        For some valid duration strings (such as "P1M", a month),
        converting to a datetime.timedelta is not possible (no obvious
        number of days). In this case, `_deserialize` returned an
        `isodate.Duration`. We can derive the timedelta by grounding to an
        actual time span, for which we require a timezone-aware start datetime.
        """
        if isinstance(duration, isodate.Duration) and start:
            years = duration.years
            months = duration.months
            days = duration.days
            seconds = duration.tdelta.seconds
            offset = pd.DateOffset(
                years=years, months=months, days=days, seconds=seconds
            )
            return (pd.Timestamp(start) + offset).to_pydatetime() - start
        return duration


class PlanningDurationField(DurationField):
    @classmethod
    def load_default(cls):
        """
        Use this with the load_default arg to __init__ if you want the default FlexMeasures planning horizon.
        """
        return current_app.config.get("FLEXMEASURES_PLANNING_HORIZON")


class AwareDateTimeField(MarshmallowClickMixin, fields.AwareDateTime):
    """Field that de-serializes to a timezone aware datetime
    and serializes back to a string."""

    def _deserialize(self, value: str, attr, obj, **kwargs) -> datetime:
        """
        Work-around until this PR lands:
        https://github.com/marshmallow-code/marshmallow/pull/1787
        """
        value = value.replace(" ", "+")
        return fields.AwareDateTime._deserialize(self, value, attr, obj, **kwargs)


class TimeIntervalSchema(Schema):
    start = AwareDateTimeField(required=True)
    duration = DurationField(required=True)


class TimeIntervalField(MarshmallowClickMixin, fields.Dict):
    """Field that de-serializes to a TimeInverval defined with start and duration."""

    def _deserialize(self, value: str, attr, obj, **kwargs) -> dict:
        try:
            v = json.loads(value)
        except json.JSONDecodeError:
            raise ValidationError()

        return TimeIntervalSchema().load(v)


class StartEndTimeSchema(Schema):
    start_time = AwareDateTimeField(required=False)
    end_time = AwareDateTimeField(required=False)

    @validates_schema
    def validate(self, data, **kwargs):
        if not (data.get("start_time") or data.get("end_time")):
            return
        if not (data.get("start_time") and data.get("end_time")):
            raise ValidationError(
                "Both start_time and end_time must be provided together."
            )
        if data["start_time"] >= data["end_time"]:
            raise ValidationError("start_time must be before end_time.")