Skip to content

Commit aca0e69

Browse files
epenetepenet
authored
Simplify service registration in recorder (home-assistant#146237)
1 parent f4e5036 commit aca0e69

File tree

2 files changed

+101
-119
lines changed

2 files changed

+101
-119
lines changed

‎homeassistant/components/recorder/__init__.py‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -45,7 +45,7 @@
4545
SupportedDialect,
4646
)
4747
from .core import Recorder
48-
from .services import async_register_services
48+
from .services import async_setup_services
4949
from .tasks import AddRecorderPlatformTask
5050
from .util import get_instance
5151

@@ -174,7 +174,7 @@ async def async_setup(hass: HomeAssistant, config: ConfigType) -> bool:
174174
instance.async_initialize()
175175
instance.async_register()
176176
instance.start()
177-
async_register_services(hass, instance)
177+
async_setup_services(hass)
178178
websocket_api.async_setup(hass)
179179

180180
await _async_setup_integration_platform(hass, instance)

‎homeassistant/components/recorder/services.py‎

Lines changed: 99 additions & 117 deletions
Original file line numberDiff line numberDiff line change
@@ -17,6 +17,7 @@
1717
)
1818
from homeassistant.helpers import config_validation as cv
1919
from homeassistant.helpers.entityfilter import generate_filter
20+
from homeassistant.helpers.recorder import DATA_INSTANCE
2021
from homeassistant.helpers.service import (
2122
async_extract_entity_ids,
2223
async_register_admin_service,
@@ -25,7 +26,6 @@
2526
from homeassistant.util.json import JsonArrayType, JsonObjectType
2627

2728
from .const import ATTR_APPLY_FILTER, ATTR_KEEP_DAYS, ATTR_REPACK, DOMAIN
28-
from .core import Recorder
2929
from .statistics import statistics_during_period
3030
from .tasks import PurgeEntitiesTask, PurgeTask
3131

@@ -87,155 +87,137 @@
8787
)
8888

8989

90-
@callback
91-
def _async_register_purge_service(hass: HomeAssistant, instance: Recorder) -> None:
92-
async def async_handle_purge_service(service: ServiceCall) -> None:
93-
"""Handle calls to the purge service."""
94-
kwargs = service.data
95-
keep_days = kwargs.get(ATTR_KEEP_DAYS, instance.keep_days)
96-
repack = cast(bool, kwargs[ATTR_REPACK])
97-
apply_filter = cast(bool, kwargs[ATTR_APPLY_FILTER])
98-
purge_before = dt_util.utcnow() - timedelta(days=keep_days)
99-
instance.queue_task(PurgeTask(purge_before, repack, apply_filter))
90+
async def _async_handle_purge_service(service: ServiceCall) -> None:
91+
"""Handle calls to the purge service."""
92+
hass = service.hass
93+
instance = hass.data[DATA_INSTANCE]
94+
kwargs = service.data
95+
keep_days = kwargs.get(ATTR_KEEP_DAYS, instance.keep_days)
96+
repack = cast(bool, kwargs[ATTR_REPACK])
97+
apply_filter = cast(bool, kwargs[ATTR_APPLY_FILTER])
98+
purge_before = dt_util.utcnow() - timedelta(days=keep_days)
99+
instance.queue_task(PurgeTask(purge_before, repack, apply_filter))
100+
101+
102+
async def _async_handle_purge_entities_service(service: ServiceCall) -> None:
103+
"""Handle calls to the purge entities service."""
104+
hass = service.hass
105+
entity_ids = await async_extract_entity_ids(hass, service)
106+
domains = service.data.get(ATTR_DOMAINS, [])
107+
keep_days = service.data.get(ATTR_KEEP_DAYS, 0)
108+
entity_globs = service.data.get(ATTR_ENTITY_GLOBS, [])
109+
entity_filter = generate_filter(domains, list(entity_ids), [], [], entity_globs)
110+
purge_before = dt_util.utcnow() - timedelta(days=keep_days)
111+
hass.data[DATA_INSTANCE].queue_task(PurgeEntitiesTask(entity_filter, purge_before))
112+
113+
114+
async def _async_handle_enable_service(service: ServiceCall) -> None:
115+
service.hass.data[DATA_INSTANCE].set_enable(True)
100116

117+
118+
async def _async_handle_disable_service(service: ServiceCall) -> None:
119+
service.hass.data[DATA_INSTANCE].set_enable(False)
120+
121+
122+
async def _async_handle_get_statistics_service(
123+
service: ServiceCall,
124+
) -> ServiceResponse:
125+
"""Handle calls to the get_statistics service."""
126+
hass = service.hass
127+
start_time = dt_util.as_utc(service.data["start_time"])
128+
end_time = (
129+
dt_util.as_utc(service.data["end_time"]) if "end_time" in service.data else None
130+
)
131+
132+
statistic_ids = service.data["statistic_ids"]
133+
types = service.data["types"]
134+
period = service.data["period"]
135+
units = service.data.get("units")
136+
137+
result = await hass.data[DATA_INSTANCE].async_add_executor_job(
138+
statistics_during_period,
139+
hass,
140+
start_time,
141+
end_time,
142+
statistic_ids,
143+
period,
144+
units,
145+
types,
146+
)
147+
148+
formatted_result: JsonObjectType = {}
149+
for statistic_id, statistic_rows in result.items():
150+
formatted_statistic_rows: JsonArrayType = []
151+
152+
for row in statistic_rows:
153+
formatted_row: JsonObjectType = {
154+
"start": dt_util.utc_from_timestamp(row["start"]).isoformat(),
155+
"end": dt_util.utc_from_timestamp(row["end"]).isoformat(),
156+
}
157+
if (last_reset := row.get("last_reset")) is not None:
158+
formatted_row["last_reset"] = dt_util.utc_from_timestamp(
159+
last_reset
160+
).isoformat()
161+
if (state := row.get("state")) is not None:
162+
formatted_row["state"] = state
163+
if (sum_value := row.get("sum")) is not None:
164+
formatted_row["sum"] = sum_value
165+
if (min_value := row.get("min")) is not None:
166+
formatted_row["min"] = min_value
167+
if (max_value := row.get("max")) is not None:
168+
formatted_row["max"] = max_value
169+
if (mean := row.get("mean")) is not None:
170+
formatted_row["mean"] = mean
171+
if (change := row.get("change")) is not None:
172+
formatted_row["change"] = change
173+
174+
formatted_statistic_rows.append(formatted_row)
175+
176+
formatted_result[statistic_id] = formatted_statistic_rows
177+
178+
return {"statistics": formatted_result}
179+
180+
181+
@callback
182+
def async_setup_services(hass: HomeAssistant) -> None:
183+
"""Register recorder services."""
101184
async_register_admin_service(
102185
hass,
103186
DOMAIN,
104187
SERVICE_PURGE,
105-
async_handle_purge_service,
188+
_async_handle_purge_service,
106189
schema=SERVICE_PURGE_SCHEMA,
107190
)
108191

109-
110-
@callback
111-
def _async_register_purge_entities_service(
112-
hass: HomeAssistant, instance: Recorder
113-
) -> None:
114-
async def async_handle_purge_entities_service(service: ServiceCall) -> None:
115-
"""Handle calls to the purge entities service."""
116-
entity_ids = await async_extract_entity_ids(hass, service)
117-
domains = service.data.get(ATTR_DOMAINS, [])
118-
keep_days = service.data.get(ATTR_KEEP_DAYS, 0)
119-
entity_globs = service.data.get(ATTR_ENTITY_GLOBS, [])
120-
entity_filter = generate_filter(domains, list(entity_ids), [], [], entity_globs)
121-
purge_before = dt_util.utcnow() - timedelta(days=keep_days)
122-
instance.queue_task(PurgeEntitiesTask(entity_filter, purge_before))
123-
124192
async_register_admin_service(
125193
hass,
126194
DOMAIN,
127195
SERVICE_PURGE_ENTITIES,
128-
async_handle_purge_entities_service,
196+
_async_handle_purge_entities_service,
129197
schema=SERVICE_PURGE_ENTITIES_SCHEMA,
130198
)
131199

132-
133-
@callback
134-
def _async_register_enable_service(hass: HomeAssistant, instance: Recorder) -> None:
135-
async def async_handle_enable_service(service: ServiceCall) -> None:
136-
instance.set_enable(True)
137-
138200
async_register_admin_service(
139201
hass,
140202
DOMAIN,
141203
SERVICE_ENABLE,
142-
async_handle_enable_service,
204+
_async_handle_enable_service,
143205
schema=SERVICE_ENABLE_SCHEMA,
144206
)
145207

146-
147-
@callback
148-
def _async_register_disable_service(hass: HomeAssistant, instance: Recorder) -> None:
149-
async def async_handle_disable_service(service: ServiceCall) -> None:
150-
instance.set_enable(False)
151-
152208
async_register_admin_service(
153209
hass,
154210
DOMAIN,
155211
SERVICE_DISABLE,
156-
async_handle_disable_service,
212+
_async_handle_disable_service,
157213
schema=SERVICE_DISABLE_SCHEMA,
158214
)
159215

160-
161-
@callback
162-
def _async_register_get_statistics_service(
163-
hass: HomeAssistant, instance: Recorder
164-
) -> None:
165-
async def async_handle_get_statistics_service(
166-
service: ServiceCall,
167-
) -> ServiceResponse:
168-
"""Handle calls to the get_statistics service."""
169-
start_time = dt_util.as_utc(service.data["start_time"])
170-
end_time = (
171-
dt_util.as_utc(service.data["end_time"])
172-
if "end_time" in service.data
173-
else None
174-
)
175-
176-
statistic_ids = service.data["statistic_ids"]
177-
types = service.data["types"]
178-
period = service.data["period"]
179-
units = service.data.get("units")
180-
181-
result = await instance.async_add_executor_job(
182-
statistics_during_period,
183-
hass,
184-
start_time,
185-
end_time,
186-
statistic_ids,
187-
period,
188-
units,
189-
types,
190-
)
191-
192-
formatted_result: JsonObjectType = {}
193-
for statistic_id, statistic_rows in result.items():
194-
formatted_statistic_rows: JsonArrayType = []
195-
196-
for row in statistic_rows:
197-
formatted_row: JsonObjectType = {
198-
"start": dt_util.utc_from_timestamp(row["start"]).isoformat(),
199-
"end": dt_util.utc_from_timestamp(row["end"]).isoformat(),
200-
}
201-
if (last_reset := row.get("last_reset")) is not None:
202-
formatted_row["last_reset"] = dt_util.utc_from_timestamp(
203-
last_reset
204-
).isoformat()
205-
if (state := row.get("state")) is not None:
206-
formatted_row["state"] = state
207-
if (sum_value := row.get("sum")) is not None:
208-
formatted_row["sum"] = sum_value
209-
if (min_value := row.get("min")) is not None:
210-
formatted_row["min"] = min_value
211-
if (max_value := row.get("max")) is not None:
212-
formatted_row["max"] = max_value
213-
if (mean := row.get("mean")) is not None:
214-
formatted_row["mean"] = mean
215-
if (change := row.get("change")) is not None:
216-
formatted_row["change"] = change
217-
218-
formatted_statistic_rows.append(formatted_row)
219-
220-
formatted_result[statistic_id] = formatted_statistic_rows
221-
222-
return {"statistics": formatted_result}
223-
224216
async_register_admin_service(
225217
hass,
226218
DOMAIN,
227219
SERVICE_GET_STATISTICS,
228-
async_handle_get_statistics_service,
220+
_async_handle_get_statistics_service,
229221
schema=SERVICE_GET_STATISTICS_SCHEMA,
230222
supports_response=SupportsResponse.ONLY,
231223
)
232-
233-
234-
@callback
235-
def async_register_services(hass: HomeAssistant, instance: Recorder) -> None:
236-
"""Register recorder services."""
237-
_async_register_purge_service(hass, instance)
238-
_async_register_purge_entities_service(hass, instance)
239-
_async_register_enable_service(hass, instance)
240-
_async_register_disable_service(hass, instance)
241-
_async_register_get_statistics_service(hass, instance)

0 commit comments

Comments
 (0)