From 12902b8a6828134308b60cacca63bf4e695f7f47 Mon Sep 17 00:00:00 2001
From: Erik Montnemery <erik@montnemery.com>
Date: Thu, 30 Nov 2023 17:45:27 +0100
Subject: [PATCH] Add NodeStrClass.__voluptuous_compile__ (#104808)

---
 homeassistant/util/yaml/objects.py |  7 +++++++
 tests/util/yaml/test_init.py       | 18 ++++++++++++++++++
 2 files changed, 25 insertions(+)

diff --git a/homeassistant/util/yaml/objects.py b/homeassistant/util/yaml/objects.py
index b2320a74d2c..6aedc85cf60 100644
--- a/homeassistant/util/yaml/objects.py
+++ b/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."""
diff --git a/tests/util/yaml/test_init.py b/tests/util/yaml/test_init.py
index 990956ec908..a4b243a4b4a 100644
--- a/tests/util/yaml/test_init.py
+++ b/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"})