Add NodeStrClass.__voluptuous_compile__ (#104808)

homeassistant/util/yaml/objects.py

@@ -2,7 +2,10 @@
 from __future__ import annotations
 
 from dataclasses import dataclass
+from typing import Any
 
+import voluptuous as vol
+from voluptuous.schema_builder import _compile_scalar
 import yaml
 
 
@@ -13,6 +16,10 @@ class NodeListClass(list):
 class NodeStrClass(str):
     """Wrapper class to be able to add attributes on a string."""
 
+    def __voluptuous_compile__(self, schema: vol.Schema) -> Any:
+        """Needed because vol.Schema.compile does not handle str subclasses."""
+        return _compile_scalar(self)
+
 
 class NodeDictClass(dict):
     """Wrapper class to be able to add attributes on a dict."""

tests/util/yaml/test_init.py

@@ -8,6 +8,7 @@ import unittest
 from unittest.mock import patch
 
 import pytest
+import voluptuous as vol
 import yaml as pyyaml
 
 from homeassistant.config import YAML_CONFIG_FILE, load_yaml_config_file
@@ -615,3 +616,20 @@ def test_string_annotated(try_both_loaders) -> None:
             getattr(value, "__config_file__", None) == expected_annotations[key][1][0]
         )
         assert getattr(value, "__line__", None) == expected_annotations[key][1][1]
+
+
+def test_string_used_as_vol_schema(try_both_loaders) -> None:
+    """Test the subclassed strings can be used in voluptuous schemas."""
+    conf = "wanted_data:\n  key_1: value_1\n  key_2: value_2\n"
+    with io.StringIO(conf) as file:
+        doc = yaml_loader.parse_yaml(file)
+
+    # Test using the subclassed strings in a schema
+    schema = vol.Schema(
+        {vol.Required(key): value for key, value in doc["wanted_data"].items()},
+    )
+    # Test using the subclassed strings when validating a schema
+    schema(doc["wanted_data"])
+    schema({"key_1": "value_1", "key_2": "value_2"})
+    with pytest.raises(vol.Invalid):
+        schema({"key_1": "value_2", "key_2": "value_1"})