Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Add pretty printing, key sorting, and better performance to to_json in Jinja #91253

Merged
merged 11 commits into from Apr 12, 2023
38 changes: 36 additions & 2 deletions homeassistant/helpers/template.py
Expand Up @@ -42,6 +42,7 @@
from jinja2.sandbox import ImmutableSandboxedEnvironment
from jinja2.utils import Namespace
from lru import LRU # pylint: disable=no-name-in-module
import orjson
import voluptuous as vol

from homeassistant.const import (
Expand Down Expand Up @@ -150,6 +151,10 @@
)
ENTITY_COUNT_GROWTH_FACTOR = 1.2

ORJSON_PASSTHROUGH_OPTIONS = (
orjson.OPT_PASSTHROUGH_DATACLASS | orjson.OPT_PASSTHROUGH_DATETIME
)


def _template_state_no_collect(hass: HomeAssistant, state: State) -> TemplateState:
"""Return a TemplateState for a state without collecting."""
Expand Down Expand Up @@ -2029,9 +2034,38 @@ def from_json(value):
return json_loads(value)


def to_json(value, ensure_ascii=True):
def _to_json_default(obj: Any) -> None:
"""Disable custom types in json serialization."""
raise TypeError(f"Object of type {type(obj).__name__} is not JSON serializable")


def to_json(
value: Any,
ensure_ascii: bool = False,
pretty_print: bool = False,
sort_keys: bool = False,
) -> str:
"""Convert an object to a JSON string."""
return json.dumps(value, ensure_ascii=ensure_ascii)
if ensure_ascii:
# For those who need ascii, we can't use orjson, so we fall back to the json library.
return json.dumps(
value,
ensure_ascii=ensure_ascii,
indent=2 if pretty_print else None,
sort_keys=sort_keys,
)

option = (
ORJSON_PASSTHROUGH_OPTIONS
| (orjson.OPT_INDENT_2 if pretty_print else 0)
| (orjson.OPT_SORT_KEYS if sort_keys else 0)
)

return orjson.dumps(
value,
option=option,
default=_to_json_default,
).decode("utf-8")


@pass_context
Expand Down
36 changes: 34 additions & 2 deletions tests/helpers/test_template.py
Expand Up @@ -3,13 +3,15 @@

from collections.abc import Iterable
from datetime import datetime, timedelta
import json
import logging
import math
import random
from typing import Any
from unittest.mock import patch

from freezegun import freeze_time
import orjson
import pytest
import voluptuous as vol

Expand Down Expand Up @@ -1047,21 +1049,51 @@ def test_to_json(hass: HomeAssistant) -> None:
).async_render()
assert actual_result == expected_result

expected_result = orjson.dumps({"Foo": "Bar"}, option=orjson.OPT_INDENT_2).decode()
actual_result = template.Template(
"{{ {'Foo': 'Bar'} | to_json(pretty_print=True) }}", hass
).async_render(parse_result=False)
assert actual_result == expected_result

expected_result = orjson.dumps(
{"Z": 26, "A": 1, "M": 13}, option=orjson.OPT_SORT_KEYS
).decode()
actual_result = template.Template(
"{{ {'Z': 26, 'A': 1, 'M': 13} | to_json(sort_keys=True) }}", hass
).async_render(parse_result=False)
assert actual_result == expected_result

with pytest.raises(TemplateError):
template.Template("{{ {'Foo': now()} | to_json }}", hass).async_render()


def test_to_json_string(hass: HomeAssistant) -> None:
def test_to_json_ensure_ascii(hass: HomeAssistant) -> None:
"""Test the object to JSON string filter."""

# Note that we're not testing the actual json.loads and json.dumps methods,
# only the filters, so we don't need to be exhaustive with our sample JSON.
actual_value_ascii = template.Template(
"{{ 'Bar ҝ éèà' | to_json }}", hass
"{{ 'Bar ҝ éèà' | to_json(ensure_ascii=True) }}", hass
).async_render()
assert actual_value_ascii == '"Bar \\u049d \\u00e9\\u00e8\\u00e0"'
actual_value = template.Template(
"{{ 'Bar ҝ éèà' | to_json(ensure_ascii=False) }}", hass
).async_render()
assert actual_value == '"Bar ҝ éèà"'

expected_result = json.dumps({"Foo": "Bar"}, indent=2)
actual_result = template.Template(
"{{ {'Foo': 'Bar'} | to_json(pretty_print=True, ensure_ascii=True) }}", hass
).async_render(parse_result=False)
assert actual_result == expected_result

expected_result = json.dumps({"Z": 26, "A": 1, "M": 13}, sort_keys=True)
actual_result = template.Template(
"{{ {'Z': 26, 'A': 1, 'M': 13} | to_json(sort_keys=True, ensure_ascii=True) }}",
hass,
).async_render(parse_result=False)
assert actual_result == expected_result


def test_from_json(hass: HomeAssistant) -> None:
"""Test the JSON string to object filter."""
Expand Down