Set up development environment

.envrc

@@ -0,0 +1 @@
+use_flake

.github/workflows/checks.yml

@@ -0,0 +1,16 @@
+name: "Checks"
+
+on:
+  push:
+  pull_request:
+
+jobs:
+  checks:
+    runs-on: ubuntu-latest
+    steps:
+    - uses: actions/checkout@v4
+    - uses: DeterminateSystems/nix-installer-action@main
+    - uses: DeterminateSystems/magic-nix-cache-action@main
+    - uses: DeterminateSystems/flake-checker-action@main
+    - run: nix flake check -L
+

.github/workflows/update.yml

@@ -0,0 +1,16 @@
+name: update-flake-lock
+on:
+  workflow_dispatch: # allows manual triggering
+  schedule:
+    - cron: '0 17 * * 5' # Fridays at 17:00
+
+jobs:
+  lockfile:
+    runs-on: ubuntu-latest
+    steps:
+      - name: Checkout repository
+        uses: actions/checkout@v4
+      - name: Install Nix
+        uses: DeterminateSystems/nix-installer-action@main
+      - name: Update flake.lock
+        uses: DeterminateSystems/update-flake-lock@main

.gitignore

@@ -0,0 +1,2 @@
+/.pre-commit-config.yaml
+.direnv

custom_components/prometheus_sensor/sensor.py

@@ -1,83 +1,98 @@
 """Prometheus Sensor component."""
+
 from datetime import timedelta
 import logging
-from typing import Union
+from typing import Final
 from urllib.parse import urljoin
 
+import aiohttp
 import voluptuous as vol
 
-from homeassistant.components.sensor import (
-    DEVICE_CLASSES_SCHEMA,
-    PLATFORM_SCHEMA as SENSOR_PLATFORM_SCHEMA,
-    STATE_CLASSES_SCHEMA,
-    SensorEntity,
-)
+from homeassistant.components.sensor import SensorEntity
+from homeassistant.components.sensor.const import SensorDeviceClass, SensorStateClass
 from homeassistant.const import (
     CONF_DEVICE_CLASS,
     CONF_NAME,
+    CONF_UNIQUE_ID,
     CONF_UNIT_OF_MEASUREMENT,
     CONF_URL,
-    CONF_UNIQUE_ID,
     STATE_PROBLEM,
     STATE_UNKNOWN,
 )
+from homeassistant.core import HomeAssistant
 from homeassistant.helpers.aiohttp_client import async_get_clientsession
 import homeassistant.helpers.config_validation as cv
+from homeassistant.helpers.config_validation import (
+    PLATFORM_SCHEMA as SENSOR_PLATFORM_SCHEMA,
+)
+from homeassistant.helpers.entity_platform import AddEntitiesCallback
+from homeassistant.helpers.typing import ConfigType, DiscoveryInfoType, StateType
 from homeassistant.util import Throttle
 
-_LOGGER = logging.getLogger(__name__)
+_LOGGER: Final = logging.getLogger(__name__)
 
-DEFAULT_URL = "http://localhost:9090"
-CONF_QUERIES = "queries"
-CONF_EXPR = "expr"
-CONF_STATE_CLASS = "state_class"
-MIN_TIME_BETWEEN_UPDATES = timedelta(seconds=60)
+DEFAULT_URL: Final = "http://localhost:9090"
+CONF_QUERIES: Final = "queries"
+CONF_EXPR: Final = "expr"
+CONF_STATE_CLASS: Final = "state_class"
+MIN_TIME_BETWEEN_UPDATES: Final = timedelta(seconds=60)
 
-_LOGGER = logging.getLogger(__name__)
-
-_QUERY_SCHEMA = vol.Schema(
+_QUERY_SCHEMA: Final = vol.Schema(
     {
         vol.Required(CONF_NAME): cv.string,
         vol.Optional(CONF_UNIQUE_ID): cv.string,
         vol.Optional(CONF_UNIT_OF_MEASUREMENT): cv.string,
         vol.Required(CONF_EXPR): cv.string,
-        vol.Optional(CONF_DEVICE_CLASS): DEVICE_CLASSES_SCHEMA,
-        vol.Optional(CONF_STATE_CLASS): STATE_CLASSES_SCHEMA,
+        vol.Optional(CONF_DEVICE_CLASS): vol.Coerce(SensorDeviceClass),
+        vol.Optional(CONF_STATE_CLASS): vol.Coerce(SensorStateClass),
     }
 )
 
-PLATFORM_SCHEMA = SENSOR_PLATFORM_SCHEMA.extend(
+PLATFORM_SCHEMA: Final = SENSOR_PLATFORM_SCHEMA.extend(
     {
-        vol.Optional(CONF_URL, default=DEFAULT_URL): cv.string,
+        vol.Optional(CONF_URL, default=DEFAULT_URL): cv.string,  # type: ignore
         vol.Required(CONF_QUERIES): [_QUERY_SCHEMA],
     }
 )
 
 
-async def async_setup_platform(hass, config, async_add_entities, discovery_info=None):
+async def async_setup_platform(
+    hass: HomeAssistant,
+    config: ConfigType,
+    async_add_entities: AddEntitiesCallback,
+    discovery_info: DiscoveryInfoType | None = None,
+):
     """Set up the sensor platform."""
     session = async_get_clientsession(hass)
-
-    prometheus = Prometheus(config.get(CONF_URL), session)
-
-    sensors = []
-    for query in config.get("queries", {}):
-        sensors.append(PrometheusSensor(prometheus, query))
-
-    _LOGGER.debug("Setting up %d prometheus sensors", len(sensors))
-
-    async_add_entities(sensors, update_before_add=True)
+    url = config[CONF_URL]
+    prometheus = Prometheus(url, session)
+
+    async_add_entities(
+        new_entities=[
+            PrometheusSensor(
+                prometheus,
+                query[CONF_EXPR],
+                query[CONF_UNIQUE_ID],
+                query[CONF_NAME],
+                query.get(CONF_DEVICE_CLASS),
+                query.get(CONF_STATE_CLASS),
+                query.get(CONF_UNIT_OF_MEASUREMENT),
+            )
+            for query in config[CONF_QUERIES]
+        ],
+        update_before_add=True,
+    )
 
 
 class Prometheus:
     """Wrapper for Prometheus API Requests."""
 
-    def __init__(self, url: str, session) -> None:
+    def __init__(self, url: str, session: aiohttp.ClientSession) -> None:
         """Initialize the Prometheus API wrapper."""
         self._session = session
         self._url = urljoin(f"{url}/", "api/v1/query")
 
-    async def query(self, expr) -> Union[str, float]:
+    async def query(self, expr: str) -> float | StateType:
         """Query expression response."""
         response = await self._session.get(self._url, params={"query": expr})
         if response.status != 200:
@@ -111,18 +126,27 @@ async def query(self, expr) -> Union[str, float]:
 class PrometheusSensor(SensorEntity):
     """Sensor entity representing the result of a PromQL expression."""
 
-    def __init__(self, prometheus: Prometheus, query: dict) -> None:
+    def __init__(
+        self,
+        prometheus: Prometheus,
+        expression: str,
+        unique_id: str,
+        device_name: str,
+        device_class: SensorDeviceClass | None,
+        state_class: SensorStateClass | None,
+        unit_of_measurement: str | None,
+    ):
         """Initialize the sensor."""
-        self._expr = query.get(CONF_EXPR)
         self._prometheus: Prometheus = prometheus
+        self._expression = expression
 
-        self._attr_name = query.get(CONF_NAME)
-        self._attr_unique_id = query.get(CONF_UNIQUE_ID)
-        self._attr_native_unit_of_measurement = query.get(CONF_UNIT_OF_MEASUREMENT)
-        self._attr_state_class = query.get(CONF_STATE_CLASS)
-        self._attr_device_class = query.get(CONF_DEVICE_CLASS)
+        self._attr_device_class = device_class
+        self._attr_name = device_name
+        self._attr_native_unit_of_measurement = unit_of_measurement
+        self._attr_state_class = state_class
+        self._attr_unique_id = unique_id
 
     @Throttle(MIN_TIME_BETWEEN_UPDATES)
     async def async_update(self) -> None:
         """Update state by executing query."""
-        self._attr_native_value = await self._prometheus.query(self._expr)
+        self._attr_native_value = await self._prometheus.query(self._expression)