Query time series by source misses conversion to BeliefSource

flexmeasures/data/models/time_series.py

@@ -2,6 +2,7 @@
 from datetime import datetime as datetime_type, timedelta
 import json
 
+from flask import current_app
 from sqlalchemy.ext.declarative import declared_attr
 from sqlalchemy.orm import Query, Session
 import timely_beliefs as tb
@@ -18,6 +19,7 @@
 from flexmeasures.data.services.time_series import collect_time_series_data
 from flexmeasures.utils.entity_address_utils import build_entity_address
 from flexmeasures.data.models.charts import chart_type_to_chart_specs
+from flexmeasures.data.models.data_sources import DataSource
 from flexmeasures.data.models.generic_assets import GenericAsset
 from flexmeasures.utils.time_utils import server_now
 from flexmeasures.utils.flexmeasures_inflection import capitalize
@@ -67,7 +69,9 @@ def search_beliefs(
         event_ends_before: Optional[datetime_type] = None,
         beliefs_after: Optional[datetime_type] = None,
         beliefs_before: Optional[datetime_type] = None,
-        source: Optional[Union[int, List[int], str, List[str]]] = None,
+        source: Optional[
+            Union[DataSource, List[DataSource], int, List[int], str, List[str]]
+        ] = None,
         as_json: bool = False,
     ):
         """Search all beliefs about events for this sensor.
@@ -76,7 +80,7 @@ def search_beliefs(
         :param event_ends_before: only return beliefs about events that end before this datetime (inclusive)
         :param beliefs_after: only return beliefs formed after this datetime (inclusive)
         :param beliefs_before: only return beliefs formed before this datetime (inclusive)
-        :param source: search only beliefs by this source (pass its name or id) or list of sources
+        :param source: search only beliefs by this source (pass the DataSource, or its name or id) or list of sources
         :param as_json: return beliefs in JSON format (e.g. for use in charts) rather than as BeliefsDataFrame
         """
         bdf = TimedBelief.search(
@@ -100,7 +104,9 @@ def chart(
         event_ends_before: Optional[datetime_type] = None,
         beliefs_after: Optional[datetime_type] = None,
         beliefs_before: Optional[datetime_type] = None,
-        source: Optional[Union[int, List[int], str, List[str]]] = None,
+        source: Optional[
+            Union[DataSource, List[DataSource], int, List[int], str, List[str]]
+        ] = None,
         include_data: bool = False,
         dataset_name: Optional[str] = None,
         **kwargs,
@@ -112,7 +118,7 @@ def chart(
         :param event_ends_before: only return beliefs about events that end before this datetime (inclusive)
         :param beliefs_after: only return beliefs formed after this datetime (inclusive)
         :param beliefs_before: only return beliefs formed before this datetime (inclusive)
-        :param source: search only beliefs by this source (pass its name or id) or list of sources
+        :param source: search only beliefs by this source (pass the DataSource, or its name or id) or list of sources
         :param include_data: if True, include data in the chart, or if False, exclude data
         :param dataset_name: optionally name the dataset used in the chart (the default name is sensor_<id>)
         """
@@ -209,7 +215,9 @@ def search(
         event_ends_before: Optional[datetime_type] = None,
         beliefs_after: Optional[datetime_type] = None,
         beliefs_before: Optional[datetime_type] = None,
-        source: Optional[Union[int, List[int], str, List[str]]] = None,
+        source: Optional[
+            Union[DataSource, List[DataSource], int, List[int], str, List[str]]
+        ] = None,
     ) -> tb.BeliefsDataFrame:
         """Search all beliefs about events for a given sensor.
 
@@ -218,16 +226,17 @@ def search(
         :param event_ends_before: only return beliefs about events that end before this datetime (inclusive)
         :param beliefs_after: only return beliefs formed after this datetime (inclusive)
         :param beliefs_before: only return beliefs formed before this datetime (inclusive)
-        :param source: search only beliefs by this source (pass its name or id) or list of sources
+        :param source: search only beliefs by this source (pass the DataSource, or its name or id) or list of sources
         """
+        parsed_sources = parse_source_arg(source)
         return cls.search_session(
             session=db.session,
             sensor=sensor,
             event_starts_after=event_starts_after,
             event_ends_before=event_ends_before,
             beliefs_after=beliefs_after,
             beliefs_before=beliefs_before,
-            source=source,
+            source=parsed_sources,
         )
 
     @classmethod
@@ -407,3 +416,38 @@ def collect(
             resolution=resolution,
             sum_multiple=sum_multiple,
         )
+
+
+def parse_source_arg(
+    source: Optional[
+        Union[DataSource, List[DataSource], int, List[int], str, List[str]]
+    ]
+) -> Optional[List[DataSource]]:
+    """Parse the "source" argument by looking up DataSources corresponding to any given ids or names."""
+    if source is None:
+        return source
+    if not isinstance(source, list):
+        sources = [source]
+    else:
+        sources = source
+    parsed_sources: List[DataSource] = []
+    for source in sources:
+        if isinstance(source, int):
+            parsed_source = DataSource.query.filter_by(id=source).one_or_none()
+            if parsed_source is None:
+                current_app.logger.warning(
+                    f"Beliefs searched for unknown source {source}"
+                )
+            else:
+                parsed_sources.append(parsed_source)
+        elif isinstance(source, str):
+            _parsed_sources = DataSource.query.filter_by(name=source).all()
+            if _parsed_sources is []:
+                current_app.logger.warning(
+                    f"Beliefs searched for unknown source {source}"
+                )
+            else:
+                parsed_sources.extend(_parsed_sources)
+        else:
+            parsed_sources.append(source)
+    return parsed_sources