Let save_to_db support saving BeliefsSeries, too (#523)

Plugins can save BeliefsSeries, too, instead of just BeliefsDataFrames.


* Let save_to_db support saving BeliefsSeries, too

Signed-off-by: F.N. Claessen <felix@seita.nl>

* Update type annotations

Signed-off-by: F.N. Claessen <felix@seita.nl>

* changelog entry

Signed-off-by: F.N. Claessen <felix@seita.nl>

* mypy: PEP 484 prohibits implicit Optional (#529)

Signed-off-by: F.N. Claessen <felix@seita.nl>

Signed-off-by: F.N. Claessen <felix@seita.nl>

Signed-off-by: F.N. Claessen <felix@seita.nl>

documentation/changelog.rst

@@ -24,6 +24,7 @@ Infrastructure / Support
 
 * Reduce size of Docker image (from 2GB to 1.4GB) [see `PR #512 <http://www.github.com/FlexMeasures/flexmeasures/pull/512>`_]
 * Remove bokeh dependency and obsolete UI views [see `PR #476 <http://www.github.com/FlexMeasures/flexmeasures/pull/476>`_]
+* Plugins can save BeliefsSeries, too, instead of just BeliefsDataFrames [see `PR #523 <http://www.github.com/FlexMeasures/flexmeasures/pull/523>`_]
 
 
 v0.11.3 | November 2, 2022

flexmeasures/api/common/utils/api_utils.py

@@ -1,3 +1,5 @@
+from __future__ import annotations
+
 from timely_beliefs.beliefs.classes import BeliefsDataFrame
 from typing import List, Sequence, Tuple, Union
 import copy
@@ -54,8 +56,8 @@ def contains_empty_items(groups: List[List[str]]):
 
 
 def parse_as_list(
-    connection: Union[Sequence[Union[str, float]], str, float], of_type: type = None
-) -> Sequence[Union[str, float, None]]:
+    connection: str | float | Sequence[str | float], of_type: type | None = None
+) -> Sequence[str | float | None]:
     """
     Return a list of connections (or values), even if it's just one connection (or value)
     """
@@ -141,7 +143,7 @@ def groups_to_dict(
     connection_groups: List[str],
     value_groups: List[List[str]],
     generic_asset_type_name: str,
-    plural_name: str = None,
+    plural_name: str | None = None,
     groups_name="groups",
 ) -> dict:
     """Put the connections and values in a dictionary and simplify if groups have identical values and/or if there is
@@ -343,7 +345,7 @@ def get_sensor_by_generic_asset_type_and_location(
 
 
 def enqueue_forecasting_jobs(
-    forecasting_jobs: List[Job] = None,
+    forecasting_jobs: list[Job] | None = None,
 ):
     """Enqueue forecasting jobs.
 
@@ -355,7 +357,7 @@ def enqueue_forecasting_jobs(
 
 def save_and_enqueue(
     data: Union[BeliefsDataFrame, List[BeliefsDataFrame]],
-    forecasting_jobs: List[Job] = None,
+    forecasting_jobs: list[Job] | None = None,
     save_changed_beliefs_only: bool = True,
 ) -> ResponseTuple:
 

flexmeasures/api/common/utils/validators.py

@@ -1,3 +1,5 @@
+from __future__ import annotations
+
 from datetime import datetime, timedelta
 from functools import wraps
 from typing import List, Tuple, Union, Optional
@@ -228,7 +230,7 @@ def decorated_service(*args, **kwargs):
 
 
 def optional_user_sources_accepted(
-    default_source: Union[int, str, List[Union[int, str]]] = None
+    default_source: int | str | list[int | str] | None = None,
 ):
     """Decorator which specifies that a GET or POST request accepts an optional source or list of data sources.
     It parses relevant form data and sets the "user_source_ids" keyword parameter.
@@ -539,7 +541,7 @@ def wrapper(*args, **kwargs):
 
 
 def assets_required(
-    generic_asset_type_name: str, plural_name: str = None, groups_name="groups"
+    generic_asset_type_name: str, plural_name: str | None = None, groups_name="groups"
 ):
     """Decorator which specifies that a GET or POST request must specify one or more assets.
     It parses relevant form data and sets the "generic_asset_name_groups" keyword param.

flexmeasures/api/v1/implementations.py

@@ -1,3 +1,5 @@
+from __future__ import annotations
+
 import isodate
 from typing import Dict, List, Optional, Tuple, Union
 from datetime import datetime as datetime_type, timedelta
@@ -154,8 +156,10 @@ def collect_connection_and_value_groups(
     start: datetime_type,
     duration: timedelta,
     connection_groups: List[List[str]],
-    user_source_ids: Union[int, List[int]] = None,  # None is interpreted as all sources
-    source_types: List[str] = None,
+    user_source_ids: int
+    | list[int]
+    | None = None,  # None is interpreted as all sources
+    source_types: list[str] | None = None,
 ) -> Tuple[dict, int]:
     """
     Code for GETting power values from the API.

flexmeasures/cli/data_add.py

@@ -1,4 +1,5 @@
 """CLI Tasks for populating the database - most useful in development"""
+from __future__ import annotations
 
 from datetime import datetime, timedelta
 from typing import Dict, List, Optional, Tuple
@@ -403,7 +404,7 @@ def add_beliefs(
     resample: bool = True,
     allow_overwrite: bool = False,
     skiprows: int = 1,
-    na_values: List[str] = None,
+    na_values: list[str] | None = None,
     nrows: Optional[int] = None,
     datecol: int = 0,
     valuecol: int = 1,

flexmeasures/data/queries/annotations.py

@@ -1,3 +1,5 @@
+from __future__ import annotations
+
 from datetime import datetime
 from typing import List, Optional
 
@@ -15,7 +17,7 @@ def query_asset_annotations(
     annotations_after: Optional[datetime] = None,
     annotations_before: Optional[datetime] = None,
     sources: Optional[List[DataSource]] = None,
-    annotation_type: str = None,
+    annotation_type: str | None = None,
 ) -> Query:
     """Match annotations assigned to the given asset."""
     query = Annotation.query.join(GenericAssetAnnotationRelationship).filter(

flexmeasures/data/utils.py

@@ -1,15 +1,15 @@
-from typing import List, Optional, Union
+from __future__ import annotations
 
 from flask import current_app
-from timely_beliefs import BeliefsDataFrame
+from timely_beliefs import BeliefsDataFrame, BeliefsSeries
 
 from flexmeasures.data import db
 from flexmeasures.data.models.data_sources import DataSource
 from flexmeasures.data.models.time_series import TimedBelief
 from flexmeasures.data.services.time_series import drop_unchanged_beliefs
 
 
-def save_to_session(objects: List[db.Model], overwrite: bool = False):
+def save_to_session(objects: list[db.Model], overwrite: bool = False):
     """Utility function to save to database, either efficiently with a bulk save, or inefficiently with a merge save."""
     if not overwrite:
         db.session.bulk_save_objects(objects)
@@ -20,8 +20,8 @@ def save_to_session(objects: List[db.Model], overwrite: bool = False):
 
 def get_data_source(
     data_source_name: str,
-    data_source_model: Optional[str] = None,
-    data_source_version: Optional[str] = None,
+    data_source_model: str | None = None,
+    data_source_version: str | None = None,
     data_source_type: str = "script",
 ) -> DataSource:
     """Make sure we have a data source. Create one if it doesn't exist, and add to session.
@@ -50,7 +50,7 @@ def get_data_source(
 
 
 def save_to_db(
-    data: Union[BeliefsDataFrame, List[BeliefsDataFrame]],
+    data: BeliefsDataFrame | BeliefsSeries | list[BeliefsDataFrame | BeliefsSeries],
     bulk_save_objects: bool = False,
     save_changed_beliefs_only: bool = True,
 ) -> str:
@@ -101,6 +101,10 @@ def save_to_db(
             # Nothing to save
             continue
 
+        # Convert series to frame if needed
+        if isinstance(timed_values, BeliefsSeries):
+            timed_values = timed_values.rename("event_value").to_frame()
+
         len_before = len(timed_values)
         if save_changed_beliefs_only:
 

flexmeasures/ui/crud/users.py

@@ -1,3 +1,5 @@
+from __future__ import annotations
+
 from typing import Optional, Union
 from datetime import datetime
 
@@ -36,7 +38,7 @@ class UserForm(FlaskForm):
     active = BooleanField("Activation Status", validators=[DataRequired()])
 
 
-def render_user(user: Optional[User], asset_count: int = 0, msg: str = None):
+def render_user(user: User | None, asset_count: int = 0, msg: str | None = None):
     user_form = UserForm()
     user_form.process(obj=user)
     return render_flexmeasures_template(