Add helper to calculate statistic period start and end (#82493)

* Add helper to calculate statistic period start and end

* Don't parse values in resolve_period

* Add specific test for resolve_period

* Improve typing

* Move to recorder/util.py

* Extract period schema
This commit is contained in:
Erik Montnemery 2022-11-26 19:00:40 +01:00 committed by GitHub
parent 405c2ca82d
commit 2fe8e95309
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
4 changed files with 206 additions and 86 deletions

View File

@ -1,9 +1,9 @@
"""Models for Recorder."""
from __future__ import annotations
from datetime import datetime
from datetime import datetime, timedelta
import logging
from typing import Any, TypedDict, overload
from typing import Any, Literal, TypedDict, overload
from sqlalchemy.engine.row import Row
@ -284,3 +284,32 @@ def row_to_compressed_state(
row_changed_changed
)
return comp_state
class CalendarStatisticPeriod(TypedDict, total=False):
"""Statistic period definition."""
period: Literal["hour", "day", "week", "month", "year"]
offset: int
class FixedStatisticPeriod(TypedDict, total=False):
"""Statistic period definition."""
end_time: datetime
start_time: datetime
class RollingWindowStatisticPeriod(TypedDict, total=False):
"""Statistic period definition."""
duration: timedelta
offset: timedelta
class StatisticPeriod(TypedDict, total=False):
"""Statistic period definition."""
calendar: CalendarStatisticPeriod
fixed_period: FixedStatisticPeriod
rolling_window: RollingWindowStatisticPeriod

View File

@ -24,8 +24,10 @@ from sqlalchemy.orm.query import Query
from sqlalchemy.orm.session import Session
from sqlalchemy.sql.lambdas import StatementLambdaElement
from typing_extensions import Concatenate, ParamSpec
import voluptuous as vol
from homeassistant.core import HomeAssistant
from homeassistant.helpers import config_validation as cv
import homeassistant.util.dt as dt_util
from .const import DATA_INSTANCE, SQLITE_URL_PREFIX, SupportedDialect
@ -35,7 +37,7 @@ from .db_schema import (
TABLES_TO_CHECK,
RecorderRuns,
)
from .models import UnsupportedDialect, process_timestamp
from .models import StatisticPeriod, UnsupportedDialect, process_timestamp
if TYPE_CHECKING:
from . import Recorder
@ -604,3 +606,83 @@ def get_instance(hass: HomeAssistant) -> Recorder:
"""Get the recorder instance."""
instance: Recorder = hass.data[DATA_INSTANCE]
return instance
PERIOD_SCHEMA = vol.Schema(
{
vol.Exclusive("calendar", "period"): vol.Schema(
{
vol.Required("period"): vol.Any("hour", "day", "week", "month", "year"),
vol.Optional("offset"): int,
}
),
vol.Exclusive("fixed_period", "period"): vol.Schema(
{
vol.Optional("start_time"): vol.All(cv.datetime, dt_util.as_utc),
vol.Optional("end_time"): vol.All(cv.datetime, dt_util.as_utc),
}
),
vol.Exclusive("rolling_window", "period"): vol.Schema(
{
vol.Required("duration"): cv.time_period_dict,
vol.Optional("offset"): cv.time_period_dict,
}
),
}
)
def resolve_period(
period_def: StatisticPeriod,
) -> tuple[datetime | None, datetime | None]:
"""Return start and end datetimes for a statistic period definition."""
start_time = None
end_time = None
if "calendar" in period_def:
calendar_period = period_def["calendar"]["period"]
start_of_day = dt_util.start_of_local_day()
cal_offset = period_def["calendar"].get("offset", 0)
if calendar_period == "hour":
start_time = dt_util.now().replace(minute=0, second=0, microsecond=0)
start_time += timedelta(hours=cal_offset)
end_time = start_time + timedelta(hours=1)
elif calendar_period == "day":
start_time = start_of_day
start_time += timedelta(days=cal_offset)
end_time = start_time + timedelta(days=1)
elif calendar_period == "week":
start_time = start_of_day - timedelta(days=start_of_day.weekday())
start_time += timedelta(days=cal_offset * 7)
end_time = start_time + timedelta(weeks=1)
elif calendar_period == "month":
start_time = start_of_day.replace(day=28)
# This works for up to 48 months of offset
start_time = (start_time + timedelta(days=cal_offset * 31)).replace(day=1)
end_time = (start_time + timedelta(days=31)).replace(day=1)
else: # calendar_period = "year"
start_time = start_of_day.replace(month=12, day=31)
# This works for 100+ years of offset
start_time = (start_time + timedelta(days=cal_offset * 366)).replace(
month=1, day=1
)
end_time = (start_time + timedelta(days=365)).replace(day=1)
start_time = dt_util.as_utc(start_time)
end_time = dt_util.as_utc(end_time)
elif "fixed_period" in period_def:
start_time = period_def["fixed_period"].get("start_time")
end_time = period_def["fixed_period"].get("end_time")
elif "rolling_window" in period_def:
duration = period_def["rolling_window"]["duration"]
now = dt_util.utcnow()
start_time = now - duration
end_time = start_time + duration
if offset := period_def["rolling_window"].get("offset"):
start_time += offset
end_time += offset
return (start_time, end_time)

View File

@ -1,9 +1,9 @@
"""The Recorder websocket API."""
from __future__ import annotations
from datetime import datetime as dt, timedelta
from datetime import datetime as dt
import logging
from typing import Any, Literal
from typing import Any, Literal, cast
import voluptuous as vol
@ -26,6 +26,7 @@ from homeassistant.util.unit_conversion import (
)
from .const import MAX_QUEUE_BACKLOG
from .models import StatisticPeriod
from .statistics import (
STATISTIC_UNIT_TO_UNIT_CONVERTER,
async_add_external_statistics,
@ -36,7 +37,13 @@ from .statistics import (
statistics_during_period,
validate_statistics,
)
from .util import async_migration_in_progress, async_migration_is_live, get_instance
from .util import (
PERIOD_SCHEMA,
async_migration_in_progress,
async_migration_is_live,
get_instance,
resolve_period,
)
_LOGGER: logging.Logger = logging.getLogger(__package__)
@ -82,24 +89,6 @@ def _ws_get_statistic_during_period(
@websocket_api.websocket_command(
{
vol.Required("type"): "recorder/statistic_during_period",
vol.Exclusive("calendar", "period"): vol.Schema(
{
vol.Required("period"): vol.Any("hour", "day", "week", "month", "year"),
vol.Optional("offset"): int,
}
),
vol.Exclusive("fixed_period", "period"): vol.Schema(
{
vol.Optional("start_time"): str,
vol.Optional("end_time"): str,
}
),
vol.Exclusive("rolling_window", "period"): vol.Schema(
{
vol.Required("duration"): cv.time_period_dict,
vol.Optional("offset"): cv.time_period_dict,
}
),
vol.Optional("statistic_id"): str,
vol.Optional("types"): vol.All(
[vol.Any("max", "mean", "min", "change")], vol.Coerce(set)
@ -116,6 +105,7 @@ def _ws_get_statistic_during_period(
vol.Optional("volume"): vol.In(VolumeConverter.VALID_UNITS),
}
),
**PERIOD_SCHEMA.schema,
}
)
@websocket_api.async_response
@ -128,67 +118,7 @@ async def ws_get_statistic_during_period(
if "offset" in msg and "duration" not in msg:
raise HomeAssistantError
start_time = None
end_time = None
if "calendar" in msg:
calendar_period = msg["calendar"]["period"]
start_of_day = dt_util.start_of_local_day()
offset = msg["calendar"].get("offset", 0)
if calendar_period == "hour":
start_time = dt_util.now().replace(minute=0, second=0, microsecond=0)
start_time += timedelta(hours=offset)
end_time = start_time + timedelta(hours=1)
elif calendar_period == "day":
start_time = start_of_day
start_time += timedelta(days=offset)
end_time = start_time + timedelta(days=1)
elif calendar_period == "week":
start_time = start_of_day - timedelta(days=start_of_day.weekday())
start_time += timedelta(days=offset * 7)
end_time = start_time + timedelta(weeks=1)
elif calendar_period == "month":
start_time = start_of_day.replace(day=28)
# This works for up to 48 months of offset
start_time = (start_time + timedelta(days=offset * 31)).replace(day=1)
end_time = (start_time + timedelta(days=31)).replace(day=1)
else: # calendar_period = "year"
start_time = start_of_day.replace(month=12, day=31)
# This works for 100+ years of offset
start_time = (start_time + timedelta(days=offset * 366)).replace(
month=1, day=1
)
end_time = (start_time + timedelta(days=365)).replace(day=1)
start_time = dt_util.as_utc(start_time)
end_time = dt_util.as_utc(end_time)
elif "fixed_period" in msg:
if start_time_str := msg["fixed_period"].get("start_time"):
if start_time := dt_util.parse_datetime(start_time_str):
start_time = dt_util.as_utc(start_time)
else:
connection.send_error(
msg["id"], "invalid_start_time", "Invalid start_time"
)
return
if end_time_str := msg["fixed_period"].get("end_time"):
if end_time := dt_util.parse_datetime(end_time_str):
end_time = dt_util.as_utc(end_time)
else:
connection.send_error(msg["id"], "invalid_end_time", "Invalid end_time")
return
elif "rolling_window" in msg:
duration = msg["rolling_window"]["duration"]
now = dt_util.utcnow()
start_time = now - duration
end_time = start_time + duration
if offset := msg["rolling_window"].get("offset"):
start_time += offset
end_time += offset
start_time, end_time = resolve_period(cast(StatisticPeriod, msg))
connection.send_message(
await get_instance(hass).async_add_executor_job(

View File

@ -1,9 +1,10 @@
"""Test util methods."""
from datetime import datetime, timedelta
from datetime import datetime, timedelta, timezone
import os
import sqlite3
from unittest.mock import MagicMock, Mock, patch
from freezegun import freeze_time
import pytest
from sqlalchemy import text
from sqlalchemy.engine.result import ChunkedIteratorResult
@ -19,6 +20,7 @@ from homeassistant.components.recorder.models import UnsupportedDialect
from homeassistant.components.recorder.util import (
end_incomplete_runs,
is_second_sunday,
resolve_period,
session_scope,
)
from homeassistant.const import EVENT_HOMEASSISTANT_STOP
@ -776,3 +778,80 @@ def test_execute_stmt_lambda_element(hass_recorder):
with patch.object(session, "execute", MockExecutor):
rows = util.execute_stmt_lambda_element(session, stmt, now, tomorrow)
assert rows == ["mock_row"]
@freeze_time(datetime(2022, 10, 21, 7, 25, tzinfo=timezone.utc))
async def test_resolve_period(hass):
"""Test statistic_during_period."""
now = dt_util.utcnow()
start_t, end_t = resolve_period({"calendar": {"period": "hour"}})
assert start_t.isoformat() == "2022-10-21T07:00:00+00:00"
assert end_t.isoformat() == "2022-10-21T08:00:00+00:00"
start_t, end_t = resolve_period({"calendar": {"period": "hour"}})
assert start_t.isoformat() == "2022-10-21T07:00:00+00:00"
assert end_t.isoformat() == "2022-10-21T08:00:00+00:00"
start_t, end_t = resolve_period({"calendar": {"period": "hour", "offset": -1}})
assert start_t.isoformat() == "2022-10-21T06:00:00+00:00"
assert end_t.isoformat() == "2022-10-21T07:00:00+00:00"
start_t, end_t = resolve_period({"calendar": {"period": "day"}})
assert start_t.isoformat() == "2022-10-21T07:00:00+00:00"
assert end_t.isoformat() == "2022-10-22T07:00:00+00:00"
start_t, end_t = resolve_period({"calendar": {"period": "day", "offset": -1}})
assert start_t.isoformat() == "2022-10-20T07:00:00+00:00"
assert end_t.isoformat() == "2022-10-21T07:00:00+00:00"
start_t, end_t = resolve_period({"calendar": {"period": "week"}})
assert start_t.isoformat() == "2022-10-17T07:00:00+00:00"
assert end_t.isoformat() == "2022-10-24T07:00:00+00:00"
start_t, end_t = resolve_period({"calendar": {"period": "week", "offset": -1}})
assert start_t.isoformat() == "2022-10-10T07:00:00+00:00"
assert end_t.isoformat() == "2022-10-17T07:00:00+00:00"
start_t, end_t = resolve_period({"calendar": {"period": "month"}})
assert start_t.isoformat() == "2022-10-01T07:00:00+00:00"
assert end_t.isoformat() == "2022-11-01T07:00:00+00:00"
start_t, end_t = resolve_period({"calendar": {"period": "month", "offset": -1}})
assert start_t.isoformat() == "2022-09-01T07:00:00+00:00"
assert end_t.isoformat() == "2022-10-01T07:00:00+00:00"
start_t, end_t = resolve_period({"calendar": {"period": "year"}})
assert start_t.isoformat() == "2022-01-01T08:00:00+00:00"
assert end_t.isoformat() == "2023-01-01T08:00:00+00:00"
start_t, end_t = resolve_period({"calendar": {"period": "year", "offset": -1}})
assert start_t.isoformat() == "2021-01-01T08:00:00+00:00"
assert end_t.isoformat() == "2022-01-01T08:00:00+00:00"
# Fixed period
assert resolve_period({}) == (None, None)
assert resolve_period({"fixed_period": {"end_time": now}}) == (None, now)
assert resolve_period({"fixed_period": {"start_time": now}}) == (now, None)
assert resolve_period({"fixed_period": {"end_time": now, "start_time": now}}) == (
now,
now,
)
# Rolling window
assert resolve_period(
{"rolling_window": {"duration": timedelta(hours=1, minutes=25)}}
) == (now - timedelta(hours=1, minutes=25), now)
assert resolve_period(
{
"rolling_window": {
"duration": timedelta(hours=1),
"offset": timedelta(minutes=-25),
}
}
) == (now - timedelta(hours=1, minutes=25), now - timedelta(minutes=25))