From ebeaa4454fbf5a35cbd2fe1192c7113d0e411ea6 Mon Sep 17 00:00:00 2001
From: Andrew Jackson <andrew@codechimp.org>
Date: Sat, 6 Jan 2024 14:06:13 +0000
Subject: [PATCH] Attribs (#389)

* Add attributes to battery_type sensor

* Got some debugging working
---
 .gitignore                                |  1 -
 .vscode/launch.json                       | 34 +++++++++++++++++++++++
 .vscode/tasks.json                        | 16 +++++++++++
 custom_components/battery_notes/const.py  |  3 ++
 custom_components/battery_notes/sensor.py | 29 +++++++++++++++++++
 5 files changed, 82 insertions(+), 1 deletion(-)
 create mode 100644 .vscode/launch.json

diff --git a/.gitignore b/.gitignore
index 70e75e132..99f374232 100644
--- a/.gitignore
+++ b/.gitignore
@@ -8,7 +8,6 @@ __pycache__
 
 # misc
 .coverage
-.vscode
 coverage.xml
 
 
diff --git a/.vscode/launch.json b/.vscode/launch.json
new file mode 100644
index 000000000..237d30bd2
--- /dev/null
+++ b/.vscode/launch.json
@@ -0,0 +1,34 @@
+{
+    // Use IntelliSense to learn about possible attributes.
+    // Hover to view descriptions of existing attributes.
+    // For more information, visit: https://go.microsoft.com/fwlink/?linkid=830387
+    "version": "0.2.0",
+    "configurations": [
+        {
+            "name": "Home Assistant",
+            "type": "python",
+            "request": "launch",
+            "module": "homeassistant",
+            "justMyCode": false,
+            "args": [
+                "--debug",
+                "-c",
+                "config"
+            ]
+        },
+        {
+            // Example of attaching to my production server
+            "name": "Python: Attach Remote",
+            "type": "python",
+            "request": "attach",
+            "port": 5678,
+            "host": "localhost",
+            "pathMappings": [
+                {
+                    "localRoot": "${workspaceFolder}",
+                    "remoteRoot": "/usr/src/homeassistant"
+                }
+            ],
+        }
+    ]
+}
\ No newline at end of file
diff --git a/.vscode/tasks.json b/.vscode/tasks.json
index 54ba9c419..52099def7 100644
--- a/.vscode/tasks.json
+++ b/.vscode/tasks.json
@@ -7,5 +7,21 @@
             "command": "scripts/develop",
             "problemMatcher": []
         }
+    ],
+    "configurations": [
+        {
+            // Example of attaching to local debug server
+            "name": "Python: Attach Local",
+            "type": "python",
+            "request": "attach",
+            "port": 5678,
+            "host": "localhost",
+            "pathMappings": [
+                {
+                    "localRoot": "${workspaceFolder}",
+                    "remoteRoot": "."
+                }
+            ],
+        }
     ]
 }
\ No newline at end of file
diff --git a/custom_components/battery_notes/const.py b/custom_components/battery_notes/const.py
index bf8b48bc2..225b1a0c4 100644
--- a/custom_components/battery_notes/const.py
+++ b/custom_components/battery_notes/const.py
@@ -49,6 +49,9 @@
 ATTR_DEVICE_ID = "device_id"
 ATTR_DATE_TIME_REPLACED = "datetime_replaced"
 ATTR_REMOVE = "remove"
+ATTR_BATTERY_QUANTITY = "battery_quantity"
+ATTR_BATTERY_TYPE = "battery_type"
+
 
 SERVICE_BATTERY_REPLACED_SCHEMA = vol.Schema(
     {
diff --git a/custom_components/battery_notes/sensor.py b/custom_components/battery_notes/sensor.py
index 3a49152c9..63a564180 100644
--- a/custom_components/battery_notes/sensor.py
+++ b/custom_components/battery_notes/sensor.py
@@ -4,6 +4,7 @@
 from datetime import datetime
 from dataclasses import dataclass
 import voluptuous as vol
+import re
 
 from homeassistant.components.sensor import (
     PLATFORM_SCHEMA,
@@ -43,6 +44,8 @@
     LAST_REPLACED,
     DOMAIN_CONFIG,
     CONF_ENABLE_REPLACED,
+    ATTR_BATTERY_QUANTITY,
+    ATTR_BATTERY_TYPE,
 )
 
 from .coordinator import BatteryNotesCoordinator
@@ -243,6 +246,32 @@ def native_value(self) -> str:
 
         return self._battery_type
 
+    @property
+    def extra_state_attributes(self) -> dict[str, str] | None:
+        """Return the state attributes of the battery type."""
+
+        matches: re.Match = re.search(
+            "^(\d+)(?=x)(?:x\s)(\w+$)|([\s\S]+)", self._battery_type
+        )
+        if matches:
+            _qty = matches.group(1) if matches.group(1) is not None else "1"
+            _type = (
+                matches.group(2) if matches.group(2) is not None else matches.group(3)
+            )
+        else:
+            _qty = 1
+            _type = self._battery_type
+
+        attrs = {
+            ATTR_BATTERY_QUANTITY: _qty,
+            ATTR_BATTERY_TYPE: _type,
+        }
+
+        super_attrs = super().extra_state_attributes
+        if super_attrs:
+            attrs.update(super_attrs)
+        return attrs
+
 
 class BatteryNotesLastReplacedSensor(SensorEntity, CoordinatorEntity):
     """Represents a battery note sensor."""