diff --git a/homeassistant/helpers/selector.py b/homeassistant/helpers/selector.py index 87574949f4e..ccb7ac67dfb 100644 --- a/homeassistant/helpers/selector.py +++ b/homeassistant/helpers/selector.py @@ -5,13 +5,12 @@ from collections.abc import Callable, Sequence from typing import Any, TypedDict, cast import voluptuous as vol -import yaml from homeassistant.backports.enum import StrEnum from homeassistant.const import CONF_MODE, CONF_UNIT_OF_MEASUREMENT from homeassistant.core import split_entity_id, valid_entity_id from homeassistant.util import decorator -from homeassistant.util.yaml.dumper import represent_odict +from homeassistant.util.yaml.dumper import add_representer, represent_odict from . import config_validation as cv @@ -889,7 +888,7 @@ class TimeSelector(Selector): return cast(str, data) -yaml.SafeDumper.add_representer( +add_representer( Selector, lambda dumper, value: represent_odict( dumper, "tag:yaml.org,2002:map", value.serialize() diff --git a/homeassistant/util/yaml/dumper.py b/homeassistant/util/yaml/dumper.py index 3eafc8abdd7..9f69c6c346e 100644 --- a/homeassistant/util/yaml/dumper.py +++ b/homeassistant/util/yaml/dumper.py @@ -1,5 +1,6 @@ """Custom dumper and representers.""" from collections import OrderedDict +from typing import Any import yaml @@ -8,10 +9,20 @@ from .objects import Input, NodeListClass # mypy: allow-untyped-calls, no-warn-return-any +try: + from yaml import CSafeDumper as FastestAvailableSafeDumper +except ImportError: + from yaml import SafeDumper as FastestAvailableSafeDumper # type: ignore[misc] + + def dump(_dict: dict) -> str: """Dump YAML to a string and remove null.""" - return yaml.safe_dump( - _dict, default_flow_style=False, allow_unicode=True, sort_keys=False + return yaml.dump( + _dict, + default_flow_style=False, + allow_unicode=True, + sort_keys=False, + Dumper=FastestAvailableSafeDumper, ).replace(": null\n", ":\n") @@ -51,17 +62,22 @@ def represent_odict( # type: ignore[no-untyped-def] return node -yaml.SafeDumper.add_representer( +def add_representer(klass: Any, representer: Any) -> None: + """Add to representer to the dumper.""" + FastestAvailableSafeDumper.add_representer(klass, representer) + + +add_representer( OrderedDict, lambda dumper, value: represent_odict(dumper, "tag:yaml.org,2002:map", value), ) -yaml.SafeDumper.add_representer( +add_representer( NodeListClass, lambda dumper, value: dumper.represent_sequence("tag:yaml.org,2002:seq", value), ) -yaml.SafeDumper.add_representer( +add_representer( Input, lambda dumper, value: dumper.represent_scalar("!input", value.name), ) diff --git a/tests/components/blueprint/test_websocket_api.py b/tests/components/blueprint/test_websocket_api.py index 9376710abee..eb2d12f5081 100644 --- a/tests/components/blueprint/test_websocket_api.py +++ b/tests/components/blueprint/test_websocket_api.py @@ -5,6 +5,7 @@ from unittest.mock import Mock, patch import pytest from homeassistant.setup import async_setup_component +from homeassistant.util.yaml import parse_yaml @pytest.fixture(autouse=True) @@ -130,9 +131,18 @@ async def test_save_blueprint(hass, aioclient_mock, hass_ws_client): assert msg["id"] == 6 assert msg["success"] assert write_mock.mock_calls - assert write_mock.call_args[0] == ( - "blueprint:\n name: Call service based on event\n domain: automation\n input:\n trigger_event:\n selector:\n text: {}\n service_to_call:\n a_number:\n selector:\n number:\n mode: box\n step: 1.0\n source_url: https://github.com/balloob/home-assistant-config/blob/main/blueprints/automation/motion_light.yaml\ntrigger:\n platform: event\n event_type: !input 'trigger_event'\naction:\n service: !input 'service_to_call'\n entity_id: light.kitchen\n", + # There are subtle differences in the dumper quoting + # behavior when quoting is not required as both produce + # valid yaml + output_yaml = write_mock.call_args[0][0] + assert output_yaml in ( + # pure python dumper will quote the value after !input + "blueprint:\n name: Call service based on event\n domain: automation\n input:\n trigger_event:\n selector:\n text: {}\n service_to_call:\n a_number:\n selector:\n number:\n mode: box\n step: 1.0\n source_url: https://github.com/balloob/home-assistant-config/blob/main/blueprints/automation/motion_light.yaml\ntrigger:\n platform: event\n event_type: !input 'trigger_event'\naction:\n service: !input 'service_to_call'\n entity_id: light.kitchen\n" + # c dumper will not quote the value after !input + "blueprint:\n name: Call service based on event\n domain: automation\n input:\n trigger_event:\n selector:\n text: {}\n service_to_call:\n a_number:\n selector:\n number:\n mode: box\n step: 1.0\n source_url: https://github.com/balloob/home-assistant-config/blob/main/blueprints/automation/motion_light.yaml\ntrigger:\n platform: event\n event_type: !input trigger_event\naction:\n service: !input service_to_call\n entity_id: light.kitchen\n" ) + # Make sure ita parsable and does not raise + assert len(parse_yaml(output_yaml)) > 1 async def test_save_existing_file(hass, aioclient_mock, hass_ws_client): diff --git a/tests/util/yaml/test_init.py b/tests/util/yaml/test_init.py index 1bdadf87a2d..11dc40233dc 100644 --- a/tests/util/yaml/test_init.py +++ b/tests/util/yaml/test_init.py @@ -33,6 +33,23 @@ def try_both_loaders(request): importlib.reload(yaml_loader) +@pytest.fixture(params=["enable_c_dumper", "disable_c_dumper"]) +def try_both_dumpers(request): + """Disable the yaml c dumper.""" + if not request.param == "disable_c_dumper": + yield + return + try: + cdumper = pyyaml.CSafeDumper + except ImportError: + return + del pyyaml.CSafeDumper + importlib.reload(yaml_loader) + yield + pyyaml.CSafeDumper = cdumper + importlib.reload(yaml_loader) + + def test_simple_list(try_both_loaders): """Test simple list.""" conf = "config:\n - simple\n - list" @@ -283,12 +300,12 @@ def test_load_yaml_encoding_error(mock_open, try_both_loaders): yaml_loader.load_yaml("test") -def test_dump(): +def test_dump(try_both_dumpers): """The that the dump method returns empty None values.""" assert yaml.dump({"a": None, "b": "b"}) == "a:\nb: b\n" -def test_dump_unicode(): +def test_dump_unicode(try_both_dumpers): """The that the dump method returns empty None values.""" assert yaml.dump({"a": None, "b": "привет"}) == "a:\nb: привет\n" @@ -424,7 +441,7 @@ class TestSecrets(unittest.TestCase): ) -def test_representing_yaml_loaded_data(): +def test_representing_yaml_loaded_data(try_both_dumpers): """Test we can represent YAML loaded data.""" files = {YAML_CONFIG_FILE: 'key: [1, "2", 3]'} with patch_yaml_files(files): @@ -460,7 +477,7 @@ def test_input_class(): assert len({input, input2}) == 1 -def test_input(try_both_loaders): +def test_input(try_both_loaders, try_both_dumpers): """Test loading inputs.""" data = {"hello": yaml.Input("test_name")} assert yaml.parse_yaml(yaml.dump(data)) == data