Add support for JSON fragments (#107213)

This commit is contained in:
J. Nick Koston 2024-01-07 17:36:31 -10:00 committed by GitHub
parent 50edc334de
commit d04e2d56da
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
11 changed files with 289 additions and 103 deletions

View File

@ -1,11 +1,9 @@
"""Rest API for Home Assistant."""
import asyncio
from asyncio import shield, timeout
from collections.abc import Collection
from functools import lru_cache
from http import HTTPStatus
import logging
from typing import Any
from aiohttp import web
from aiohttp.web_exceptions import HTTPBadRequest
@ -42,11 +40,10 @@ from homeassistant.exceptions import (
)
from homeassistant.helpers import config_validation as cv, template
from homeassistant.helpers.event import EventStateChangedData
from homeassistant.helpers.json import json_dumps
from homeassistant.helpers.json import json_dumps, json_fragment
from homeassistant.helpers.service import async_get_all_descriptions
from homeassistant.helpers.typing import ConfigType, EventType
from homeassistant.util.json import json_loads
from homeassistant.util.read_only_dict import ReadOnlyDict
_LOGGER = logging.getLogger(__name__)
@ -377,14 +374,14 @@ class APIDomainServicesView(HomeAssistantView):
)
context = self.context(request)
changed_states: list[ReadOnlyDict[str, Collection[Any]]] = []
changed_states: list[json_fragment] = []
@ha.callback
def _async_save_changed_entities(
event: EventType[EventStateChangedData],
) -> None:
if event.context == context and (state := event.data["new_state"]):
changed_states.append(state.as_dict())
changed_states.append(state.json_fragment)
cancel_listen = hass.bus.async_listen(
EVENT_STATE_CHANGED, _async_save_changed_entities, run_immediately=True

View File

@ -115,11 +115,15 @@ async def async_attach_trigger(
if event_context_items:
# Fast path for simple items comparison
if not (event.context.as_dict().items() >= event_context_items):
# This is safe because we do not mutate the event context
# pylint: disable-next=protected-access
if not (event.context._as_dict.items() >= event_context_items):
return False
elif event_context_schema:
# Slow path for schema validation
event_context_schema(dict(event.context.as_dict()))
# This is safe because we make a copy of the event context
# pylint: disable-next=protected-access
event_context_schema(dict(event.context._as_dict))
except vol.Invalid:
# If event doesn't match, skip event
return False

View File

@ -116,7 +116,7 @@ def _partial_cached_event_message(event: Event) -> str:
in cached_event_message.
"""
return (
_message_to_json_or_none({"type": "event", "event": event.as_dict()})
_message_to_json_or_none({"type": "event", "event": event.json_fragment})
or INVALID_JSON_PARTIAL_MESSAGE
)

View File

@ -87,7 +87,7 @@ from .helpers.deprecation import (
check_if_deprecated_constant,
dir_with_deprecated_constants,
)
from .helpers.json import json_dumps
from .helpers.json import json_dumps, json_fragment
from .util import dt as dt_util, location
from .util.async_ import (
cancelling,
@ -996,8 +996,6 @@ class HomeAssistant:
class Context:
"""The context that triggered something."""
__slots__ = ("user_id", "parent_id", "id", "origin_event", "_as_dict")
def __init__(
self,
user_id: str | None = None,
@ -1009,23 +1007,37 @@ class Context:
self.user_id = user_id
self.parent_id = parent_id
self.origin_event: Event | None = None
self._as_dict: ReadOnlyDict[str, str | None] | None = None
def __eq__(self, other: Any) -> bool:
"""Compare contexts."""
return bool(self.__class__ == other.__class__ and self.id == other.id)
@cached_property
def _as_dict(self) -> dict[str, str | None]:
"""Return a dictionary representation of the context.
Callers should be careful to not mutate the returned dictionary
as it will mutate the cached version.
"""
return {
"id": self.id,
"parent_id": self.parent_id,
"user_id": self.user_id,
}
def as_dict(self) -> ReadOnlyDict[str, str | None]:
"""Return a dictionary representation of the context."""
if not self._as_dict:
self._as_dict = ReadOnlyDict(
{
"id": self.id,
"parent_id": self.parent_id,
"user_id": self.user_id,
}
)
return self._as_dict
"""Return a ReadOnlyDict representation of the context."""
return self._as_read_only_dict
@cached_property
def _as_read_only_dict(self) -> ReadOnlyDict[str, str | None]:
"""Return a ReadOnlyDict representation of the context."""
return ReadOnlyDict(self._as_dict)
@cached_property
def json_fragment(self) -> json_fragment:
"""Return a JSON fragment of the context."""
return json_fragment(json_dumps(self._as_dict))
class EventOrigin(enum.Enum):
@ -1042,8 +1054,6 @@ class EventOrigin(enum.Enum):
class Event:
"""Representation of an event within the bus."""
__slots__ = ("event_type", "data", "origin", "time_fired", "context", "_as_dict")
def __init__(
self,
event_type: str,
@ -1062,26 +1072,54 @@ class Event:
id=ulid_at_time(dt_util.utc_to_timestamp(self.time_fired))
)
self.context = context
self._as_dict: ReadOnlyDict[str, Any] | None = None
if not context.origin_event:
context.origin_event = self
def as_dict(self) -> ReadOnlyDict[str, Any]:
@cached_property
def _as_dict(self) -> dict[str, Any]:
"""Create a dict representation of this Event.
Callers should be careful to not mutate the returned dictionary
as it will mutate the cached version.
"""
return {
"event_type": self.event_type,
"data": self.data,
"origin": self.origin.value,
"time_fired": self.time_fired.isoformat(),
# _as_dict is marked as protected
# to avoid callers outside of this module
# from misusing it by mistake.
"context": self.context._as_dict, # pylint: disable=protected-access
}
def as_dict(self) -> ReadOnlyDict[str, Any]:
"""Create a ReadOnlyDict representation of this Event.
Async friendly.
"""
if not self._as_dict:
self._as_dict = ReadOnlyDict(
{
"event_type": self.event_type,
"data": ReadOnlyDict(self.data),
"origin": self.origin.value,
"time_fired": self.time_fired.isoformat(),
"context": self.context.as_dict(),
}
)
return self._as_dict
return self._as_read_only_dict
@cached_property
def _as_read_only_dict(self) -> ReadOnlyDict[str, Any]:
"""Create a ReadOnlyDict representation of this Event."""
as_dict = self._as_dict
data = as_dict["data"]
context = as_dict["context"]
# json_fragment will serialize data from a ReadOnlyDict
# or a normal dict so its ok to have either. We only
# mutate the cache if someone asks for the as_dict version
# to avoid storing multiple copies of the data in memory.
if type(data) is not ReadOnlyDict:
as_dict["data"] = ReadOnlyDict(data)
if type(context) is not ReadOnlyDict:
as_dict["context"] = ReadOnlyDict(context)
return ReadOnlyDict(as_dict)
@cached_property
def json_fragment(self) -> json_fragment:
"""Return an event as a JSON fragment."""
return json_fragment(json_dumps(self._as_dict))
def __repr__(self) -> str:
"""Return the representation."""
@ -1397,7 +1435,6 @@ class State:
self.context = context or Context()
self.state_info = state_info
self.domain, self.object_id = split_entity_id(self.entity_id)
self._as_dict: ReadOnlyDict[str, Collection[Any]] | None = None
@property
def name(self) -> str:
@ -1406,36 +1443,66 @@ class State:
"_", " "
)
def as_dict(self) -> ReadOnlyDict[str, Collection[Any]]:
@cached_property
def _as_dict(self) -> dict[str, Any]:
"""Return a dict representation of the State.
Callers should be careful to not mutate the returned dictionary
as it will mutate the cached version.
"""
last_changed_isoformat = self.last_changed.isoformat()
if self.last_changed == self.last_updated:
last_updated_isoformat = last_changed_isoformat
else:
last_updated_isoformat = self.last_updated.isoformat()
return {
"entity_id": self.entity_id,
"state": self.state,
"attributes": self.attributes,
"last_changed": last_changed_isoformat,
"last_updated": last_updated_isoformat,
# _as_dict is marked as protected
# to avoid callers outside of this module
# from misusing it by mistake.
"context": self.context._as_dict, # pylint: disable=protected-access
}
def as_dict(
self,
) -> ReadOnlyDict[str, datetime.datetime | Collection[Any]]:
"""Return a ReadOnlyDict representation of the State.
Async friendly.
To be used for JSON serialization.
Can be used for JSON serialization.
Ensures: state == State.from_dict(state.as_dict())
"""
if not self._as_dict:
last_changed_isoformat = self.last_changed.isoformat()
if self.last_changed == self.last_updated:
last_updated_isoformat = last_changed_isoformat
else:
last_updated_isoformat = self.last_updated.isoformat()
self._as_dict = ReadOnlyDict(
{
"entity_id": self.entity_id,
"state": self.state,
"attributes": self.attributes,
"last_changed": last_changed_isoformat,
"last_updated": last_updated_isoformat,
"context": self.context.as_dict(),
}
)
return self._as_dict
return self._as_read_only_dict
@cached_property
def _as_read_only_dict(
self,
) -> ReadOnlyDict[str, datetime.datetime | Collection[Any]]:
"""Return a ReadOnlyDict representation of the State."""
as_dict = self._as_dict
context = as_dict["context"]
# json_fragment will serialize data from a ReadOnlyDict
# or a normal dict so its ok to have either. We only
# mutate the cache if someone asks for the as_dict version
# to avoid storing multiple copies of the data in memory.
if type(context) is not ReadOnlyDict:
as_dict["context"] = ReadOnlyDict(context)
return ReadOnlyDict(as_dict)
@cached_property
def as_dict_json(self) -> str:
"""Return a JSON string of the State."""
return json_dumps(self.as_dict())
return json_dumps(self._as_dict)
@cached_property
def json_fragment(self) -> json_fragment:
"""Return a JSON fragment of the State."""
return json_fragment(self.as_dict_json)
@cached_property
def as_compressed_state(self) -> dict[str, Any]:
@ -1449,7 +1516,10 @@ class State:
if state_context.parent_id is None and state_context.user_id is None:
context: dict[str, Any] | str = state_context.id
else:
context = state_context.as_dict()
# _as_dict is marked as protected
# to avoid callers outside of this module
# from misusing it by mistake.
context = state_context._as_dict # pylint: disable=protected-access
compressed_state = {
COMPRESSED_STATE_STATE: self.state,
COMPRESSED_STATE_ATTRIBUTES: self.attributes,

View File

@ -45,6 +45,8 @@ def json_encoder_default(obj: Any) -> Any:
Hand other objects to the original method.
"""
if hasattr(obj, "json_fragment"):
return obj.json_fragment
if isinstance(obj, (set, tuple)):
return list(obj)
if isinstance(obj, float):
@ -114,6 +116,9 @@ def json_bytes_strip_null(data: Any) -> bytes:
return json_bytes(_strip_null(orjson.loads(result)))
json_fragment = orjson.Fragment
def json_dumps(data: Any) -> str:
r"""Dump json string.

View File

@ -10,6 +10,7 @@ from homeassistant.const import ATTR_RESTORED, EVENT_HOMEASSISTANT_STOP
from homeassistant.core import HomeAssistant, State, callback, valid_entity_id
from homeassistant.exceptions import HomeAssistantError
import homeassistant.util.dt as dt_util
from homeassistant.util.json import json_loads
from . import start
from .entity import Entity
@ -70,9 +71,9 @@ class StoredState:
self.state = state
def as_dict(self) -> dict[str, Any]:
"""Return a dict representation of the stored state."""
"""Return a dict representation of the stored state to be JSON serialized."""
result = {
"state": self.state.as_dict(),
"state": self.state.json_fragment,
"extra_data": self.extra_data.as_dict() if self.extra_data else None,
"last_seen": self.last_seen,
}
@ -270,7 +271,7 @@ class RestoreStateData:
# To fully mimic all the attribute data types when loaded from storage,
# we're going to serialize it to JSON and then re-load it.
if state is not None:
state = State.from_dict(_encode_complex(state.as_dict()))
state = State.from_dict(json_loads(state.as_dict_json)) # type: ignore[arg-type]
if state is not None:
self.last_states[entity_id] = StoredState(
state, extra_data, dt_util.utcnow()
@ -279,32 +280,6 @@ class RestoreStateData:
self.entities.pop(entity_id)
def _encode(value: Any) -> Any:
"""Little helper to JSON encode a value."""
try:
return JSONEncoder.default(
None, # type: ignore[arg-type]
value,
)
except TypeError:
return value
def _encode_complex(value: Any) -> Any:
"""Recursively encode all values with the JSONEncoder."""
if isinstance(value, dict):
return {_encode(key): _encode_complex(value) for key, value in value.items()}
if isinstance(value, list):
return [_encode_complex(val) for val in value]
new_value = _encode(value)
if isinstance(new_value, type(value)):
return new_value
return _encode_complex(new_value)
class RestoreEntity(Entity):
"""Mixin class for restoring previous entity state."""

View File

@ -5,7 +5,7 @@ from ast import literal_eval
import asyncio
import base64
import collections.abc
from collections.abc import Callable, Collection, Generator, Iterable
from collections.abc import Callable, Generator, Iterable
from contextlib import AbstractContextManager, suppress
from contextvars import ContextVar
from datetime import datetime, timedelta
@ -940,7 +940,6 @@ class TemplateStateBase(State):
self._hass = hass
self._collect = collect
self._entity_id = entity_id
self._as_dict: ReadOnlyDict[str, Collection[Any]] | None = None
def _collect_state(self) -> None:
if self._collect and (render_info := _render_info.get()):

View File

@ -74,7 +74,7 @@ from homeassistant.helpers.dispatcher import (
async_dispatcher_connect,
async_dispatcher_send,
)
from homeassistant.helpers.json import JSONEncoder, _orjson_default_encoder
from homeassistant.helpers.json import JSONEncoder, _orjson_default_encoder, json_dumps
from homeassistant.helpers.typing import ConfigType, StateType
from homeassistant.setup import setup_component
from homeassistant.util.async_ import run_callback_threadsafe
@ -507,6 +507,11 @@ def load_json_object_fixture(
return json_loads_object(load_fixture(filename, integration))
def json_round_trip(obj: Any) -> Any:
"""Round trip an object to JSON."""
return json_loads(json_dumps(obj))
def mock_state_change_event(
hass: HomeAssistant, new_state: State, old_state: State | None = None
) -> None:

View File

@ -19,12 +19,15 @@ from homeassistant.helpers.json import (
json_bytes_strip_null,
json_dumps,
json_dumps_sorted,
json_fragment,
save_json,
)
from homeassistant.util import dt as dt_util
from homeassistant.util.color import RGBColor
from homeassistant.util.json import SerializationError, load_json
from tests.common import json_round_trip
# Test data that can be saved as JSON
TEST_JSON_A = {"a": 1, "B": "two"}
TEST_JSON_B = {"a": "one", "B": 2}
@ -45,7 +48,8 @@ def test_json_encoder(hass: HomeAssistant, encoder: type[json.JSONEncoder]) -> N
assert sorted(ha_json_enc.default(data)) == sorted(data)
# Test serializing an object which implements as_dict
assert ha_json_enc.default(state) == state.as_dict()
default = ha_json_enc.default(state)
assert json_round_trip(default) == json_round_trip(state.as_dict())
def test_json_encoder_raises(hass: HomeAssistant) -> None:
@ -133,6 +137,35 @@ def test_json_dumps_rgb_color_subclass() -> None:
assert json_dumps(rgb) == "[4,2,1]"
def test_json_fragments() -> None:
"""Test the json dumps with a fragment."""
assert (
json_dumps(
[
json_fragment('{"inner":"fragment2"}'),
json_fragment('{"inner":"fragment2"}'),
]
)
== '[{"inner":"fragment2"},{"inner":"fragment2"}]'
)
class Fragment1:
@property
def json_fragment(self):
return json_fragment('{"inner":"fragment1"}')
class Fragment2:
@property
def json_fragment(self):
return json_fragment('{"inner":"fragment2"}')
assert (
json_dumps([Fragment1(), Fragment2()])
== '[{"inner":"fragment1"},{"inner":"fragment2"}]'
)
def test_json_bytes_strip_null() -> None:
"""Test stripping nul from strings."""

View File

@ -31,6 +31,7 @@ from tests.common import (
MockModule,
MockPlatform,
async_fire_time_changed,
json_round_trip,
mock_integration,
mock_platform,
)
@ -318,12 +319,15 @@ async def test_dump_data(hass: HomeAssistant) -> None:
# b4 should not be written, since it is now expired
# b5 should be written, since current state is restored by entity registry
assert len(written_states) == 3
assert written_states[0]["state"]["entity_id"] == "input_boolean.b1"
assert written_states[0]["state"]["state"] == "on"
assert written_states[1]["state"]["entity_id"] == "input_boolean.b3"
assert written_states[1]["state"]["state"] == "off"
assert written_states[2]["state"]["entity_id"] == "input_boolean.b5"
assert written_states[2]["state"]["state"] == "off"
state0 = json_round_trip(written_states[0])
state1 = json_round_trip(written_states[1])
state2 = json_round_trip(written_states[2])
assert state0["state"]["entity_id"] == "input_boolean.b1"
assert state0["state"]["state"] == "on"
assert state1["state"]["entity_id"] == "input_boolean.b3"
assert state1["state"]["state"] == "off"
assert state2["state"]["entity_id"] == "input_boolean.b5"
assert state2["state"]["state"] == "off"
# Test that removed entities are not persisted
await entity.async_remove()
@ -340,10 +344,12 @@ async def test_dump_data(hass: HomeAssistant) -> None:
args = mock_write_data.mock_calls[0][1]
written_states = args[0]
assert len(written_states) == 2
assert written_states[0]["state"]["entity_id"] == "input_boolean.b3"
assert written_states[0]["state"]["state"] == "off"
assert written_states[1]["state"]["entity_id"] == "input_boolean.b5"
assert written_states[1]["state"]["state"] == "off"
state0 = json_round_trip(written_states[0])
state1 = json_round_trip(written_states[1])
assert state0["state"]["entity_id"] == "input_boolean.b3"
assert state0["state"]["state"] == "off"
assert state1["state"]["entity_id"] == "input_boolean.b5"
assert state1["state"]["state"] == "off"
async def test_dump_error(hass: HomeAssistant) -> None:

View File

@ -54,6 +54,7 @@ from homeassistant.exceptions import (
MaxLengthExceeded,
ServiceNotFound,
)
from homeassistant.helpers.json import json_dumps
import homeassistant.util.dt as dt_util
from homeassistant.util.read_only_dict import ReadOnlyDict
from homeassistant.util.unit_system import METRIC_SYSTEM
@ -624,6 +625,38 @@ def test_event_eq() -> None:
assert event1.as_dict() == event2.as_dict()
def test_event_json_fragment() -> None:
"""Test event JSON fragments."""
now = dt_util.utcnow()
data = {"some": "attr"}
context = ha.Context()
event1, event2 = (
ha.Event("some_type", data, time_fired=now, context=context) for _ in range(2)
)
# We are testing that the JSON fragments are the same when as_dict is called
# after json_fragment or before.
json_fragment_1 = event1.json_fragment
as_dict_1 = event1.as_dict()
as_dict_2 = event2.as_dict()
json_fragment_2 = event2.json_fragment
assert json_dumps(json_fragment_1) == json_dumps(json_fragment_2)
# We also test that the as_dict is the same
assert as_dict_1 == as_dict_2
# Finally we verify that the as_dict is a ReadOnlyDict
# as is the data and context inside regardless of
# if the json fragment was called first or not
assert isinstance(as_dict_1, ReadOnlyDict)
assert isinstance(as_dict_1["data"], ReadOnlyDict)
assert isinstance(as_dict_1["context"], ReadOnlyDict)
assert isinstance(as_dict_2, ReadOnlyDict)
assert isinstance(as_dict_2["data"], ReadOnlyDict)
assert isinstance(as_dict_2["context"], ReadOnlyDict)
def test_event_repr() -> None:
"""Test that Event repr method works."""
assert str(ha.Event("TestEvent")) == "<Event TestEvent[L]>"
@ -712,6 +745,44 @@ def test_state_as_dict_json() -> None:
assert state.as_dict_json is as_dict_json_1
def test_state_json_fragment() -> None:
"""Test state JSON fragments."""
last_time = datetime(1984, 12, 8, 12, 0, 0)
state1, state2 = (
ha.State(
"happy.happy",
"on",
{"pig": "dog"},
last_updated=last_time,
last_changed=last_time,
context=ha.Context(id="01H0D6K3RFJAYAV2093ZW30PCW"),
)
for _ in range(2)
)
# We are testing that the JSON fragments are the same when as_dict is called
# after json_fragment or before.
json_fragment_1 = state1.json_fragment
as_dict_1 = state1.as_dict()
as_dict_2 = state2.as_dict()
json_fragment_2 = state2.json_fragment
assert json_dumps(json_fragment_1) == json_dumps(json_fragment_2)
# We also test that the as_dict is the same
assert as_dict_1 == as_dict_2
# Finally we verify that the as_dict is a ReadOnlyDict
# as is the attributes and context inside regardless of
# if the json fragment was called first or not
assert isinstance(as_dict_1, ReadOnlyDict)
assert isinstance(as_dict_1["attributes"], ReadOnlyDict)
assert isinstance(as_dict_1["context"], ReadOnlyDict)
assert isinstance(as_dict_2, ReadOnlyDict)
assert isinstance(as_dict_2["attributes"], ReadOnlyDict)
assert isinstance(as_dict_2["context"], ReadOnlyDict)
def test_state_as_compressed_state() -> None:
"""Test a State as compressed state."""
last_time = datetime(1984, 12, 8, 12, 0, 0, tzinfo=dt_util.UTC)
@ -1729,6 +1800,27 @@ def test_context() -> None:
assert c.id is not None
def test_context_json_fragment() -> None:
"""Test context JSON fragments."""
context1, context2 = (ha.Context(id="01H0D6K3RFJAYAV2093ZW30PCW") for _ in range(2))
# We are testing that the JSON fragments are the same when as_dict is called
# after json_fragment or before.
json_fragment_1 = context1.json_fragment
as_dict_1 = context1.as_dict()
as_dict_2 = context2.as_dict()
json_fragment_2 = context2.json_fragment
assert json_dumps(json_fragment_1) == json_dumps(json_fragment_2)
# We also test that the as_dict is the same
assert as_dict_1 == as_dict_2
# Finally we verify that the as_dict is a ReadOnlyDict
# regardless of if the json fragment was called first or not
assert isinstance(as_dict_1, ReadOnlyDict)
assert isinstance(as_dict_2, ReadOnlyDict)
async def test_async_functions_with_callback(hass: HomeAssistant) -> None:
"""Test we deal with async functions accidentally marked as callback."""
runs = []