From b5f660398c509ca73dae1f44a821b591646e9f52 Mon Sep 17 00:00:00 2001
From: Oxan van Leeuwen <oxan@oxanvanleeuwen.nl>
Date: Thu, 25 Nov 2021 21:35:33 +0100
Subject: [PATCH] Add map filter for text sensors (#2761)

---
 esphome/components/text_sensor/__init__.py | 38 ++++++++++++----------
 esphome/components/text_sensor/filter.cpp  |  6 ++++
 esphome/components/text_sensor/filter.h    | 11 +++++++
 tests/test2.yaml                           | 10 ++++++
 4 files changed, 48 insertions(+), 17 deletions(-)

diff --git a/esphome/components/text_sensor/__init__.py b/esphome/components/text_sensor/__init__.py
index a070e6f86d..e0fc6af19c 100644
--- a/esphome/components/text_sensor/__init__.py
+++ b/esphome/components/text_sensor/__init__.py
@@ -49,6 +49,7 @@ ToLowerFilter = text_sensor_ns.class_("ToLowerFilter", Filter)
 AppendFilter = text_sensor_ns.class_("AppendFilter", Filter)
 PrependFilter = text_sensor_ns.class_("PrependFilter", Filter)
 SubstituteFilter = text_sensor_ns.class_("SubstituteFilter", Filter)
+MapFilter = text_sensor_ns.class_("MapFilter", Filter)
 
 
 @FILTER_REGISTRY.register("lambda", LambdaFilter, cv.returning_lambda)
@@ -79,26 +80,21 @@ async def prepend_filter_to_code(config, filter_id):
     return cg.new_Pvariable(filter_id, config)
 
 
-def validate_substitute(value):
-    if isinstance(value, dict):
-        return cv.Schema(
-            {
-                cv.Required(CONF_FROM): cv.string,
-                cv.Required(CONF_TO): cv.string,
-            }
-        )(value)
-    value = cv.string(value)
-    if "->" not in value:
-        raise cv.Invalid("Substitute mapping must contain '->'")
-    a, b = value.split("->", 1)
-    a, b = a.strip(), b.strip()
-    return validate_substitute({CONF_FROM: cv.string(a), CONF_TO: cv.string(b)})
+def validate_mapping(value):
+    if not isinstance(value, dict):
+        value = cv.string(value)
+        if "->" not in value:
+            raise cv.Invalid("Mapping must contain '->'")
+        a, b = value.split("->", 1)
+        value = {CONF_FROM: a.strip(), CONF_TO: b.strip()}
+
+    return cv.Schema(
+        {cv.Required(CONF_FROM): cv.string, cv.Required(CONF_TO): cv.string}
+    )(value)
 
 
 @FILTER_REGISTRY.register(
-    "substitute",
-    SubstituteFilter,
-    cv.All(cv.ensure_list(validate_substitute), cv.Length(min=2)),
+    "substitute", SubstituteFilter, cv.ensure_list(validate_mapping)
 )
 async def substitute_filter_to_code(config, filter_id):
     from_strings = [conf[CONF_FROM] for conf in config]
@@ -106,6 +102,14 @@ async def substitute_filter_to_code(config, filter_id):
     return cg.new_Pvariable(filter_id, from_strings, to_strings)
 
 
+@FILTER_REGISTRY.register("map", MapFilter, cv.ensure_list(validate_mapping))
+async def map_filter_to_code(config, filter_id):
+    map_ = cg.std_ns.class_("map").template(cg.std_string, cg.std_string)
+    return cg.new_Pvariable(
+        filter_id, map_([(item[CONF_FROM], item[CONF_TO]) for item in config])
+    )
+
+
 icon = cv.icon
 
 
diff --git a/esphome/components/text_sensor/filter.cpp b/esphome/components/text_sensor/filter.cpp
index 14df6238ff..a692378e8b 100644
--- a/esphome/components/text_sensor/filter.cpp
+++ b/esphome/components/text_sensor/filter.cpp
@@ -70,5 +70,11 @@ optional<std::string> SubstituteFilter::new_value(std::string value) {
   return value;
 }
 
+// Map
+optional<std::string> MapFilter::new_value(std::string value) {
+  auto item = mappings_.find(value);
+  return item == mappings_.end() ? value : item->second;
+}
+
 }  // namespace text_sensor
 }  // namespace esphome
diff --git a/esphome/components/text_sensor/filter.h b/esphome/components/text_sensor/filter.h
index 6a1d9ab04e..38f35e6172 100644
--- a/esphome/components/text_sensor/filter.h
+++ b/esphome/components/text_sensor/filter.h
@@ -4,6 +4,7 @@
 #include "esphome/core/helpers.h"
 #include <queue>
 #include <utility>
+#include <map>
 
 namespace esphome {
 namespace text_sensor {
@@ -108,5 +109,15 @@ class SubstituteFilter : public Filter {
   std::vector<std::string> to_strings_;
 };
 
+/// A filter that maps values from one set to another
+class MapFilter : public Filter {
+ public:
+  MapFilter(std::map<std::string, std::string> mappings) : mappings_(std::move(mappings)) {}
+  optional<std::string> new_value(std::string value) override;
+
+ protected:
+  std::map<std::string, std::string> mappings_;
+};
+
 }  // namespace text_sensor
 }  // namespace esphome
diff --git a/tests/test2.yaml b/tests/test2.yaml
index 3afef9501d..50743dc643 100644
--- a/tests/test2.yaml
+++ b/tests/test2.yaml
@@ -458,6 +458,16 @@ text_sensor:
     name: 'Template Text Sensor'
     lambda: |-
       return {"Hello World"};
+    filters:
+      - to_upper:
+      - to_lower:
+      - append: "xyz"
+      - prepend: "abcd"
+      - substitute:
+          - Hello -> Goodbye
+      - map:
+          - red -> green
+      - lambda: return {"1234"};
   - platform: homeassistant
     entity_id: sensor.hello_world2
     id: ha_hello_world2