From fa21f7e87942693c066856d0532e5bfb7d39620a Mon Sep 17 00:00:00 2001
From: Adam Goldsmith <adam@adamgoldsmith.name>
Date: Thu, 20 Jun 2024 01:08:37 -0400
Subject: [PATCH] Add filtering by area and label

---
 check_home_assistant_state.py | 16 ++++++++++++++++
 1 file changed, 16 insertions(+)

diff --git a/check_home_assistant_state.py b/check_home_assistant_state.py
index 7ee1002..8635a1d 100755
--- a/check_home_assistant_state.py
+++ b/check_home_assistant_state.py
@@ -68,6 +68,8 @@ class Entities(nagiosplugin.Resource):
     token: str
     device_class: str
     numeric: bool
+    area: str
+    label: str
     filters: list[AttributeFilter]
     attribute: str | None
     friendly_name: bool
@@ -152,6 +154,16 @@ class Entities(nagiosplugin.Resource):
             template_filter += (
                 f"|selectattr('attributes.{ f.attribute }', '{ op }', '{ f.value }')"
             )
+
+        if self.area:
+            template_filter += (
+                f"|selectattr('entity_id', 'in', area_entities('{ self.area }'))"
+            )
+        if self.label:
+            template_filter += (
+                f"|selectattr('entity_id', 'in', label_entities('{ self.label }'))"
+            )
+
         if self.exclude:
             template_filter += (
                 f"|rejectattr('entity_id', 'in', { json.dumps(self.exclude) })"
@@ -362,6 +374,8 @@ def main():
         action="extend",
         help="explicitly include entities by id. Other entities will not be considered if this is specified. Listed entities must also match the specified device class",
     )
+    common_args.add_argument("--area", type=str, help="area to filter by")
+    common_args.add_argument("--label", type=str, help="label to filter by")
     common_args.add_argument(
         "-e",
         "--exclude",
@@ -438,6 +452,8 @@ def main():
         "device_class": args.device_class,
         "attribute": args.attribute,
         "filters": args.filter,
+        "area": args.area,
+        "label": args.label,
         "include": args.include,
         "exclude": args.exclude,
         "friendly_name": args.friendly,