refactor code

Signed-off-by: Victor Garcia Reolid <victor@seita.nl>

flexmeasures/data/models/planning/tests/test_solver.py

@@ -20,7 +20,7 @@
 from flexmeasures.data.models.planning.linear_optimization import device_scheduler
 from flexmeasures.data.models.planning.tests.utils import check_constraints
 from flexmeasures.data.models.planning.utils import initialize_series, initialize_df
-from flexmeasures.data.schemas.scheduling.storage import TimedEventSchema
+from flexmeasures.data.schemas.sensors import TimedEventSchema
 from flexmeasures.utils.calculations import (
     apply_stock_changes_and_losses,
     integrate_time_series,

flexmeasures/data/schemas/scheduling/storage.py

@@ -2,7 +2,6 @@
 
 from datetime import datetime
 
-import isodate
 from flask import current_app
 from marshmallow import (
     Schema,
@@ -12,22 +11,14 @@
     fields,
     validates,
 )
-from marshmallow.validate import OneOf, ValidationError, Validator
-import pandas as pd
+from marshmallow.validate import OneOf, ValidationError
 
 from flexmeasures.data.models.time_series import Sensor
-from flexmeasures.data.schemas.times import AwareDateTimeField, DurationField
 from flexmeasures.data.schemas.units import QuantityField
-from flexmeasures.data.schemas.sensors import QuantityOrSensor, SensorIdField
+from flexmeasures.data.schemas.sensors import QuantityOrSensor, TimeSeriesOrSensor
 
 from flexmeasures.utils.unit_utils import ur
 
-from flexmeasures.data.schemas.utils import (
-    FMValidationError,
-    MarshmallowClickMixin,
-    with_appcontext_if_needed,
-)
-
 
 class EfficiencyField(QuantityField):
     """Field that deserializes to a Quantity with % units. Must be greater than 0% and less than or equal to 100%.
@@ -56,135 +47,6 @@ def __init__(self, *args, **kwargs):
         )
 
 
-class TimedEventSchema(Schema):
-    value = fields.Float(required=True)
-    datetime = AwareDateTimeField(required=False)
-    start = AwareDateTimeField(required=False)
-    end = AwareDateTimeField(required=False)
-    duration = DurationField(required=False)
-
-    def __init__(
-        self,
-        timezone: str | None = None,
-        value_validator: Validator | None = None,
-        *args,
-        **kwargs,
-    ):
-        """A time period (or single point) with a target value.
-
-        :param timezone:  Optionally, set a timezone to be able to interpret nominal durations.
-        """
-        self.timezone = timezone
-        self.value_validator = value_validator
-        super().__init__(*args, **kwargs)
-
-    @validates("value")
-    def validate_value(self, _value):
-        if self.value_validator is not None:
-            self.value_validator(_value)
-
-    @validates_schema
-    def check_time_window(self, data: dict, **kwargs):
-        """Checks whether a complete time interval can be derived from the timing fields.
-
-        The data is updated in-place, guaranteeing that the 'start' and 'end' fields are filled out.
-        """
-        dt = data.get("datetime")
-        start = data.get("start")
-        end = data.get("end")
-        duration = data.get("duration")
-
-        if dt is not None:
-            if any([p is not None for p in (start, end, duration)]):
-                raise ValidationError(
-                    "If using the 'datetime' field, no 'start', 'end' or 'duration' is expected."
-                )
-            data["start"] = dt
-            data["end"] = dt
-        elif duration is not None:
-            if self.timezone is None and isinstance(duration, isodate.Duration):
-                raise ValidationError(
-                    "Cannot interpret nominal duration used in the 'duration' field without a known timezone."
-                )
-            elif all([p is None for p in (start, end)]) or all(
-                [p is not None for p in (start, end)]
-            ):
-                raise ValidationError(
-                    "If using the 'duration' field, either 'start' or 'end' is expected."
-                )
-            if start is not None:
-                grounded = DurationField.ground_from(
-                    duration, pd.Timestamp(start).tz_convert(self.timezone)
-                )
-                data["start"] = start
-                data["end"] = start + grounded
-            else:
-                grounded = DurationField.ground_from(
-                    -duration, pd.Timestamp(end).tz_convert(self.timezone)
-                )
-                data["start"] = end + grounded
-                data["end"] = end
-        else:
-            if any([p is None for p in (start, end)]):
-                raise ValidationError(
-                    "Missing field(s) to describe timing: use the 'datetime' field, "
-                    "or a combination of 2 fields of 'start', 'end' and 'duration'."
-                )
-            data["start"] = start
-            data["end"] = end
-
-
-class TimeSeriesOrSensor(MarshmallowClickMixin, fields.Field):
-    def __init__(
-        self,
-        unit,
-        *args,
-        timezone: str | None = None,
-        value_validator: Validator | None = None,
-        **kwargs,
-    ):
-        """
-        The timezone is only used in case a time series is specified and one
-        of the *timed events* in the time series uses a nominal duration, such as "P1D".
-        """
-        super().__init__(*args, **kwargs)
-        self.timezone = timezone
-        self.value_validator = value_validator
-        self.unit = ur.Quantity(unit)
-
-    @with_appcontext_if_needed()
-    def _deserialize(
-        self, value: str | dict[str, int], attr, obj, **kwargs
-    ) -> list[dict] | Sensor:
-
-        if isinstance(value, dict):
-            if "sensor" not in value:
-                raise FMValidationError(
-                    "Dictionary provided but `sensor` key not found."
-                )
-
-            sensor = SensorIdField(unit=self.unit)._deserialize(
-                value["sensor"], None, None
-            )
-
-            return sensor
-
-        elif isinstance(value, list):
-            field = fields.List(
-                fields.Nested(
-                    TimedEventSchema(
-                        timezone=self.timezone, value_validator=self.value_validator
-                    )
-                )
-            )
-
-            return field._deserialize(value, None, None)
-        else:
-            raise FMValidationError(
-                f"Unsupported value type. `{type(value)}` was provided but only dict and list are supported."
-            )
-
-
 class StorageFlexModelSchema(Schema):
     """
     This schema lists fields we require when scheduling storage assets.

flexmeasures/data/schemas/sensors.py

@@ -1,9 +1,12 @@
 from __future__ import annotations
-from marshmallow import Schema, fields, validates, ValidationError
+from marshmallow import Schema, fields, validates, ValidationError, validates_schema
+from marshmallow.validate import Validator
 from pint import DimensionalityError
 
 import json
 import re
+import isodate
+import pandas as pd
 
 from flexmeasures.data import ma, db
 from flexmeasures.data.models.generic_assets import GenericAsset
@@ -14,7 +17,7 @@
     with_appcontext_if_needed,
 )
 from flexmeasures.utils.unit_utils import is_valid_unit, ur, units_are_convertible
-from flexmeasures.data.schemas.times import DurationField
+from flexmeasures.data.schemas.times import DurationField, AwareDateTimeField
 
 
 class JSON(fields.Field):
@@ -28,6 +31,84 @@ def _serialize(self, value, attr, data, **kwargs) -> str:
         return json.dumps(value)
 
 
+class TimedEventSchema(Schema):
+    value = fields.Float(required=True)
+    datetime = AwareDateTimeField(required=False)
+    start = AwareDateTimeField(required=False)
+    end = AwareDateTimeField(required=False)
+    duration = DurationField(required=False)
+
+    def __init__(
+        self,
+        timezone: str | None = None,
+        value_validator: Validator | None = None,
+        *args,
+        **kwargs,
+    ):
+        """A time period (or single point) with a value.
+
+        :param timezone:  Optionally, set a timezone to be able to interpret nominal durations.
+        """
+        self.timezone = timezone
+        self.value_validator = value_validator
+        super().__init__(*args, **kwargs)
+
+    @validates("value")
+    def validate_value(self, _value):
+        if self.value_validator is not None:
+            self.value_validator(_value)
+
+    @validates_schema
+    def check_time_window(self, data: dict, **kwargs):
+        """Checks whether a complete time interval can be derived from the timing fields.
+
+        The data is updated in-place, guaranteeing that the 'start' and 'end' fields are filled out.
+        """
+        dt = data.get("datetime")
+        start = data.get("start")
+        end = data.get("end")
+        duration = data.get("duration")
+
+        if dt is not None:
+            if any([p is not None for p in (start, end, duration)]):
+                raise ValidationError(
+                    "If using the 'datetime' field, no 'start', 'end' or 'duration' is expected."
+                )
+            data["start"] = dt
+            data["end"] = dt
+        elif duration is not None:
+            if self.timezone is None and isinstance(duration, isodate.Duration):
+                raise ValidationError(
+                    "Cannot interpret nominal duration used in the 'duration' field without a known timezone."
+                )
+            elif all([p is None for p in (start, end)]) or all(
+                [p is not None for p in (start, end)]
+            ):
+                raise ValidationError(
+                    "If using the 'duration' field, either 'start' or 'end' is expected."
+                )
+            if start is not None:
+                grounded = DurationField.ground_from(
+                    duration, pd.Timestamp(start).tz_convert(self.timezone)
+                )
+                data["start"] = start
+                data["end"] = start + grounded
+            else:
+                grounded = DurationField.ground_from(
+                    -duration, pd.Timestamp(end).tz_convert(self.timezone)
+                )
+                data["start"] = end + grounded
+                data["end"] = end
+        else:
+            if any([p is None for p in (start, end)]):
+                raise ValidationError(
+                    "Missing field(s) to describe timing: use the 'datetime' field, "
+                    "or a combination of 2 fields of 'start', 'end' and 'duration'."
+                )
+            data["start"] = start
+            data["end"] = end
+
+
 class SensorSchemaMixin(Schema):
     """
     Base sensor schema.
@@ -185,3 +266,54 @@ def convert(self, value, param, ctx, **kwargs):
             _value = value
 
         return super().convert(_value, param, ctx, **kwargs)
+
+
+class TimeSeriesOrSensor(MarshmallowClickMixin, fields.Field):
+    def __init__(
+        self,
+        unit,
+        *args,
+        timezone: str | None = None,
+        value_validator: Validator | None = None,
+        **kwargs,
+    ):
+        """
+        The timezone is only used in case a time series is specified and one
+        of the *timed events* in the time series uses a nominal duration, such as "P1D".
+        """
+        super().__init__(*args, **kwargs)
+        self.timezone = timezone
+        self.value_validator = value_validator
+        self.unit = ur.Quantity(unit)
+
+    @with_appcontext_if_needed()
+    def _deserialize(
+        self, value: str | dict[str, int], attr, obj, **kwargs
+    ) -> list[dict] | Sensor:
+
+        if isinstance(value, dict):
+            if "sensor" not in value:
+                raise FMValidationError(
+                    "Dictionary provided but `sensor` key not found."
+                )
+
+            sensor = SensorIdField(unit=self.unit)._deserialize(
+                value["sensor"], None, None
+            )
+
+            return sensor
+
+        elif isinstance(value, list):
+            field = fields.List(
+                fields.Nested(
+                    TimedEventSchema(
+                        timezone=self.timezone, value_validator=self.value_validator
+                    )
+                )
+            )
+
+            return field._deserialize(value, None, None)
+        else:
+            raise FMValidationError(
+                f"Unsupported value type. `{type(value)}` was provided but only dict and list are supported."
+            )

flexmeasures/data/schemas/tests/test_scheduling.py

@@ -10,9 +10,9 @@
     ProcessType,
 )
 from flexmeasures.data.schemas.scheduling.storage import (
-    TimedEventSchema,
     StorageFlexModelSchema,
 )
+from flexmeasures.data.schemas.sensors import TimedEventSchema
 
 
 @pytest.mark.parametrize(