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

Co-authored-by: J. Nick Koston <nick@koston.org>
Co-authored-by: epenet <6771947+epenet@users.noreply.github.com>

homeassistant/helpers/template.py

@@ -42,6 +42,7 @@ from jinja2.runtime import AsyncLoopContext, LoopContext
 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 (
@@ -150,6 +151,10 @@ CACHED_TEMPLATE_NO_COLLECT_LRU: MutableMapping[State, TemplateState] = LRU(
 )
 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."""
@@ -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

tests/helpers/test_template.py

@@ -3,6 +3,7 @@ from __future__ import annotations
 
 from collections.abc import Iterable
 from datetime import datetime, timedelta
+import json
 import logging
 import math
 import random
@@ -10,6 +11,7 @@ from typing import Any
 from unittest.mock import patch
 
 from freezegun import freeze_time
+import orjson
 import pytest
 import voluptuous as vol
 
@@ -1047,14 +1049,31 @@ 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
 
-def test_to_json_string(hass: HomeAssistant) -> None:
+    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_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(
@@ -1062,6 +1081,19 @@ def test_to_json_string(hass: HomeAssistant) -> None:
     ).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."""