Annotations

flexmeasures/cli/data_add.py

@@ -12,12 +12,18 @@
 import getpass
 from sqlalchemy.exc import IntegrityError
 import timely_beliefs as tb
+from workalendar.registry import registry
 
 from flexmeasures.data import db
 from flexmeasures.data.services.forecasting import create_forecasting_jobs
 from flexmeasures.data.services.users import create_user
 from flexmeasures.data.models.user import Account, AccountRole, RolesAccounts
-from flexmeasures.data.models.time_series import Sensor, TimedBelief
+from flexmeasures.data.models.time_series import (
+    Sensor,
+    TimedBelief,
+    Annotation,
+    GenericAssetAnnotationRelationship,
+)
 from flexmeasures.data.schemas.sensors import SensorSchema
 from flexmeasures.data.schemas.generic_assets import (
     GenericAssetSchema,
@@ -30,6 +36,7 @@
     get_or_create_source,
     get_source_or_none,
 )
+from flexmeasures.utils import flexmeasures_inflection
 from flexmeasures.utils.time_utils import server_now
 
 
@@ -480,6 +487,79 @@ def add_beliefs(
             print("As a possible workaround, use the --allow-overwrite flag.")
 
 
+@fm_add_data.command("holidays")
+@with_appcontext
+@click.option(
+    "--year",
+    type=click.INT,
+    help="The year for which to look up holidays",
+)
+@click.option(
+    "--country",
+    "countries",
+    type=click.STRING,
+    multiple=True,
+    help="The ISO 3166-1 country/region or ISO 3166-2 sub-region for which to look up holidays (such as US, BR and DE). This argument can be given multiple times.",
+)
+@click.option(
+    "--asset-id",
+    "generic_asset_ids",
+    type=click.INT,
+    multiple=True,
+    help="Add annotations to this asset. Follow up with the asset's ID. This argument can be given multiple times.",
+)
+@click.option(
+    "--account-id",
+    "account_ids",
+    type=click.INT,
+    multiple=True,
+    help="Add annotations to all assets of this account. Follow up with the account's ID. This argument can be given multiple times.",
+)
+def add_holidays(
+    year: int,
+    countries: List[str],
+    generic_asset_ids: List[int],
+    account_ids: List[int],
+):
+    """Add holiday annotations to assets."""
+    calendars = registry.get_calendars(countries)
+    num_holidays = {}
+    asset_query = db.session.query(GenericAsset)
+    if generic_asset_ids:
+        asset_query = asset_query.filter(GenericAsset.id.in_(generic_asset_ids))
+    if account_ids:
+        asset_query = asset_query.filter(GenericAsset.account_id.in_(account_ids))
+    assets = asset_query.all()
+    for country, calendar in calendars.items():
+        _source = get_or_create_source(
+            "workalendar", model=country, source_type="CLI script"
+        )
+        holidays = calendar().holidays(year)
+        for holiday in holidays:
+            start = pd.Timestamp(holiday[0])
+            end = start + pd.offsets.DateOffset(days=1)
+            annotation = Annotation(
+                name=holiday[1],
+                start=start,
+                end=end,
+                source=_source,
+                type="holiday",
+            )
+            db.session.add(annotation)
+            db.session.flush()
+            for asset in assets:
+                relation = GenericAssetAnnotationRelationship(
+                    annotation=annotation,
+                    asset=asset,
+                )
+                db.session.add(relation)
+        num_holidays[country] = len(holidays)
+    db.session.commit()
+    print(
+        f"Successfully added holidays to {len(assets)} {flexmeasures_inflection.pluralize('asset', len(assets))}:\n{num_holidays}"
+    )
+
+
 @fm_add_data.command("forecasts")
 @with_appcontext
 @click.option(

flexmeasures/data/migrations/versions/7f8b8920355f_create_annotation_table.py

@@ -0,0 +1,79 @@
+"""create annotation table
+
+Revision ID: 7f8b8920355f
+Revises: c1d316c60985
+Create Date: 2022-01-29 20:23:29.996133
+
+"""
+from alembic import op
+import click
+import sqlalchemy as sa
+
+
+# revision identifiers, used by Alembic.
+revision = "7f8b8920355f"
+down_revision = "c1d316c60985"
+branch_labels = None
+depends_on = None
+
+
+def upgrade():
+    create_annotation_table()
+    create_annotation_asset_relationship_table()
+    create_annotation_sensor_relationship_table()
+
+
+def downgrade():
+    click.confirm(
+        "This downgrade drops the tables 'annotations_assets', 'annotations_sensors' and 'annotation'. Continue?",
+        abort=True,
+    )
+    op.drop_table("annotations_assets")
+    op.drop_table("annotations_sensors")
+    op.drop_constraint(op.f("annotation_name_key"), "annotation", type_="unique")
+    op.drop_table("annotation")
+    op.execute("DROP TYPE annotation_type;")
+
+
+def create_annotation_sensor_relationship_table():
+    op.create_table(
+        "annotations_sensors",
+        sa.Column("sensor_id", sa.Integer(), nullable=False, primary_key=True),
+        sa.Column("annotation_id", sa.Integer(), nullable=False, primary_key=True),
+        sa.ForeignKeyConstraint(("sensor_id",), ["sensor.id"]),
+        sa.ForeignKeyConstraint(("annotation_id",), ["annotation.id"]),
+    )
+
+
+def create_annotation_asset_relationship_table():
+    op.create_table(
+        "annotations_assets",
+        sa.Column("generic_asset_id", sa.Integer(), nullable=False, primary_key=True),
+        sa.Column("annotation_id", sa.Integer(), nullable=False, primary_key=True),
+        sa.ForeignKeyConstraint(("generic_asset_id",), ["generic_asset.id"]),
+        sa.ForeignKeyConstraint(("annotation_id",), ["annotation.id"]),
+    )
+
+
+def create_annotation_table():
+    op.create_table(
+        "annotation",
+        sa.Column(
+            "id", sa.Integer(), nullable=False, autoincrement=True, primary_key=True
+        ),
+        sa.Column("name", sa.String(255), nullable=False),
+        sa.Column("start", sa.DateTime(timezone=True), nullable=False),
+        sa.Column("end", sa.DateTime(timezone=True), nullable=False),
+        sa.Column("source_id", sa.Integer(), nullable=False),
+        sa.Column(
+            "type",
+            sa.Enum("alert", "holiday", "label", name="annotation_type"),
+            nullable=False,
+        ),
+        sa.ForeignKeyConstraint(("source_id",), ["data_source.id"]),
+    )
+    op.create_unique_constraint(
+        op.f("annotation_name_key"),
+        "annotation",
+        ["name", "start", "source_id", "type"],
+    )

flexmeasures/data/models/data_sources.py

@@ -87,11 +87,16 @@ def __str__(self):
 
 
 def get_or_create_source(
-    source: Union[User, str], source_type: Optional[str] = None, flush: bool = True
+    source: Union[User, str],
+    source_type: Optional[str] = None,
+    model: Optional[str] = None,
+    flush: bool = True,
 ) -> DataSource:
     if is_user(source):
         source_type = "user"
     query = DataSource.query.filter(DataSource.type == source_type)
+    if model is not None:
+        query = query.filter(DataSource.model == model)
     if is_user(source):
         query = query.filter(DataSource.user == source)
     elif isinstance(source, str):
@@ -102,11 +107,11 @@ def get_or_create_source(
     if not _source:
         current_app.logger.info(f"Setting up '{source}' as new data source...")
         if is_user(source):
-            _source = DataSource(user=source)
+            _source = DataSource(user=source, model=model)
         else:
             if source_type is None:
                 raise TypeError("Please specify a source type")
-            _source = DataSource(name=source, type=source_type)
+            _source = DataSource(name=source, model=model, type=source_type)
         db.session.add(_source)
         if flush:
             # assigns id so that we can reference the new object in the current db session

flexmeasures/data/models/time_series.py

@@ -6,6 +6,7 @@
 from sqlalchemy.ext.declarative import declared_attr
 from sqlalchemy.ext.mutable import MutableDict
 from sqlalchemy.orm import Query, Session
+from sqlalchemy import UniqueConstraint
 import timely_beliefs as tb
 from timely_beliefs.beliefs.probabilistic_utils import get_median_belief
 import timely_beliefs.utils as tb_utils
@@ -682,3 +683,83 @@ def parse_source_arg(
         else:
             parsed_sources.append(source)
     return parsed_sources
+
+
+class Annotation(db.Model):
+    """An annotation is a nominal value that applies to a specific time or time span.
+
+    Examples of annotation types:
+        - user annotation: annotation.type == "label" and annotation.source.type == "user"
+        - unresolved alert: annotation.type == "alert"
+        - resolved alert: annotation.type == "label" and annotation.source.type == "alerting script"
+        - organisation holiday: annotation.type == "holiday" and annotation.source.type == "user"
+        - public holiday: annotation.type == "holiday" and annotation.source.name == "workalendar"
+    """
+
+    id = db.Column(db.Integer, nullable=False, autoincrement=True, primary_key=True)
+    name = db.Column(db.String(255), nullable=False)
+    start = db.Column(db.DateTime(timezone=True), nullable=False)
+    end = db.Column(db.DateTime(timezone=True), nullable=False)
+    source_id = db.Column(db.Integer, db.ForeignKey(DataSource.__tablename__ + ".id"))
+    source = db.relationship(
+        "DataSource",
+        foreign_keys=[source_id],
+        backref=db.backref("annotations", lazy=True),
+    )
+    type = db.Column(db.Enum("alert", "holiday", "label", name="annotation_type"))
+    UniqueConstraint("name", "start", "source_id", "type")
+
+    @property
+    def duration(self) -> timedelta:
+        return self.end - self.start
+
+    def __repr__(self) -> str:
+        return f"<Annotation {self.id}: {self.name} ({self.type}), start: {self.start} end: {self.end}, source: {self.source}>"
+
+
+class GenericAssetAnnotationRelationship(db.Model):
+    """Links annotations to generic assets."""
+
+    __tablename__ = "annotations_assets"
+
+    generic_asset_id = db.Column(
+        db.Integer, db.ForeignKey("generic_asset.id"), nullable=False, primary_key=True
+    )
+    asset = db.relationship(
+        "GenericAsset",
+        foreign_keys=[generic_asset_id],
+        backref=db.backref("annotations", lazy=True),
+    )
+
+    annotation_id = db.Column(
+        db.Integer, db.ForeignKey("annotation.id"), nullable=False
+    )
+    annotation = db.relationship(
+        "Annotation",
+        foreign_keys=[annotation_id],
+        backref=db.backref("generic_assets", lazy=True),
+    )
+
+
+class SensorAnnotationRelationship(db.Model):
+    """Links annotations to sensors."""
+
+    __tablename__ = "annotations_sensors"
+
+    sensor_id = db.Column(
+        db.Integer, db.ForeignKey("sensor.id"), nullable=False, primary_key=True
+    )
+    sensor = db.relationship(
+        "Sensor",
+        foreign_keys=[sensor_id],
+        backref=db.backref("annotations", lazy=True),
+    )
+
+    annotation_id = db.Column(
+        db.Integer, db.ForeignKey("annotation.id"), nullable=False
+    )
+    annotation = db.relationship(
+        "Annotation",
+        foreign_keys=[annotation_id],
+        backref=db.backref("sensors", lazy=True),
+    )

flexmeasures/utils/flexmeasures_inflection.py

@@ -1,7 +1,10 @@
+from __future__ import annotations
 import re
 
+import inflect
 import inflection
 
+p = inflect.engine()
 
 # Give the inflection module some help for our domain
 inflection.UNCOUNTABLES.add("solar")
@@ -30,10 +33,10 @@ def parameterize(word):
     return inflection.parameterize(word).replace("-", "_")
 
 
-def pluralize(word):
+def pluralize(word, count: str | int | None = None):
     if word.lower().split()[-1] in inflection.UNCOUNTABLES:
         return word
-    return inflection.pluralize(word)
+    return p.plural(word, count)
 
 
 def titleize(word):

requirements/app.in

@@ -10,6 +10,7 @@ pint
 py-moneyed
 iso8601
 xlrd
+workalendar
 inflection
 inflect
 humanize