feat: add @auto_refresh decorator with smart timestamp-based refresh system

cryptk · cryptk · commit 69b9e51532f5 · 2025-10-30T15:56:56.000-05:00
- Implement @auto_refresh decorator that automatically refreshes equipment state after control methods
- Add update_if_older_than() with asyncio.Lock to prevent redundant API calls
- Refactor equipment classes to pass OmniLogic instance instead of OmniLogicAPI for explicit dependency injection
diff --git a/pyomnilogic_local/_base.py b/pyomnilogic_local/_base.py
@@ -1,10 +1,13 @@
 import logging
-from typing import Generic, TypeVar, cast
+from typing import TYPE_CHECKING, Generic, TypeVar, cast
 
 from pyomnilogic_local.api.api import OmniLogicAPI
 from pyomnilogic_local.models import MSPEquipmentType, Telemetry
 from pyomnilogic_local.models.telemetry import TelemetryType
 
+if TYPE_CHECKING:
+    from pyomnilogic_local.omnilogic import OmniLogic
+
 # Define type variables for generic equipment types
 MSPConfigT = TypeVar("MSPConfigT", bound=MSPEquipmentType)
 TelemetryT = TypeVar("TelemetryT", bound=TelemetryType | None)
@@ -27,18 +30,23 @@ class OmniEquipment(Generic[MSPConfigT, TelemetryT]):
     # Use a forward reference for the type hint to avoid issues with self-referential generics
     child_equipment: dict[int, "OmniEquipment[MSPConfigT, TelemetryT]"]
 
-    def __init__(self, _api: OmniLogicAPI, mspconfig: MSPConfigT, telemetry: Telemetry | None) -> None:
+    def __init__(self, omni: "OmniLogic", mspconfig: MSPConfigT, telemetry: Telemetry | None) -> None:
         """Initialize the equipment with configuration and telemetry data.
 
         Args:
-            _api: The OmniLogic API instance
+            omni: The OmniLogic instance (parent controller)
             mspconfig: The MSP configuration for this specific equipment
             telemetry: The full Telemetry object containing all equipment telemetry
         """
-        self._api = _api
+        self._omni = omni
 
         self.update(mspconfig, telemetry)
 
+    @property
+    def _api(self) -> "OmniLogicAPI":
+        """Access the OmniLogic API through the parent controller."""
+        return self._omni._api  # pylint: disable=protected-access
+
     @property
     def bow_id(self) -> int | None:
         """The bow ID of the equipment."""
@@ -85,7 +93,7 @@ def update_telemetry(self, telemetry: Telemetry) -> None:
         # if hasattr(self, "telemetry"):
         # Extract the specific telemetry for this equipment from the full telemetry object
         # Note: Some equipment (like sensors) don't have their own telemetry, so this may be None
-        if specific_telemetry := telemetry.get_telem_by_systemid(self.mspconfig.system_id) is not None:
+        if (specific_telemetry := telemetry.get_telem_by_systemid(self.mspconfig.system_id)) is not None:
             self.telemetry = cast(TelemetryT, specific_telemetry)
         else:
             self.telemetry = cast(TelemetryT, None)
diff --git a/pyomnilogic_local/backyard.py b/pyomnilogic_local/backyard.py
@@ -1,6 +1,6 @@
 import logging
+from typing import TYPE_CHECKING
 
-from pyomnilogic_local.api.api import OmniLogicAPI
 from pyomnilogic_local.collections import EquipmentDict
 from pyomnilogic_local.models.mspconfig import MSPBackyard
 from pyomnilogic_local.models.telemetry import Telemetry, TelemetryBackyard
@@ -11,6 +11,9 @@
 from .relay import Relay
 from .sensor import Sensor
 
+if TYPE_CHECKING:
+    from pyomnilogic_local.omnilogic import OmniLogic
+
 _LOGGER = logging.getLogger(__name__)
 
 
@@ -22,8 +25,8 @@ class Backyard(OmniEquipment[MSPBackyard, TelemetryBackyard]):
     relays: EquipmentDict[Relay] = EquipmentDict()
     sensors: EquipmentDict[Sensor] = EquipmentDict()
 
-    def __init__(self, _api: OmniLogicAPI, mspconfig: MSPBackyard, telemetry: Telemetry) -> None:
-        super().__init__(_api, mspconfig, telemetry)
+    def __init__(self, omni: "OmniLogic", mspconfig: MSPBackyard, telemetry: Telemetry) -> None:
+        super().__init__(omni, mspconfig, telemetry)
 
     def _update_equipment(self, mspconfig: MSPBackyard, telemetry: Telemetry | None) -> None:
         """Update both the configuration and telemetry data for the equipment."""
@@ -41,28 +44,28 @@ def _update_bows(self, mspconfig: MSPBackyard, telemetry: Telemetry) -> None:
             self.bow = EquipmentDict()
             return
 
-        self.bow = EquipmentDict([Bow(self._api, bow, telemetry) for bow in mspconfig.bow])
+        self.bow = EquipmentDict([Bow(self._omni, bow, telemetry) for bow in mspconfig.bow])
 
     def _update_lights(self, mspconfig: MSPBackyard, telemetry: Telemetry) -> None:
         """Update the lights based on the MSP configuration."""
         if mspconfig.colorlogic_light is None:
             self.lights = EquipmentDict()
             return
 
-        self.lights = EquipmentDict([ColorLogicLight(self._api, light, telemetry) for light in mspconfig.colorlogic_light])
+        self.lights = EquipmentDict([ColorLogicLight(self._omni, light, telemetry) for light in mspconfig.colorlogic_light])
 
     def _update_relays(self, mspconfig: MSPBackyard, telemetry: Telemetry) -> None:
         """Update the relays based on the MSP configuration."""
         if mspconfig.relay is None:
             self.relays = EquipmentDict()
             return
 
-        self.relays = EquipmentDict([Relay(self._api, relay, telemetry) for relay in mspconfig.relay])
+        self.relays = EquipmentDict([Relay(self._omni, relay, telemetry) for relay in mspconfig.relay])
 
     def _update_sensors(self, mspconfig: MSPBackyard, telemetry: Telemetry) -> None:
         """Update the sensors based on the MSP configuration."""
         if mspconfig.sensor is None:
             self.sensors = EquipmentDict()
             return
 
-        self.sensors = EquipmentDict([Sensor(self._api, sensor, telemetry) for sensor in mspconfig.sensor])
+        self.sensors = EquipmentDict([Sensor(self._omni, sensor, telemetry) for sensor in mspconfig.sensor])
diff --git a/pyomnilogic_local/bow.py b/pyomnilogic_local/bow.py
@@ -1,5 +1,6 @@
+from typing import TYPE_CHECKING
+
 from pyomnilogic_local._base import OmniEquipment
-from pyomnilogic_local.api.api import OmniLogicAPI
 from pyomnilogic_local.chlorinator import Chlorinator
 from pyomnilogic_local.collections import EquipmentDict
 from pyomnilogic_local.colorlogiclight import _LOGGER, ColorLogicLight
@@ -12,6 +13,9 @@
 from pyomnilogic_local.relay import Relay
 from pyomnilogic_local.sensor import Sensor
 
+if TYPE_CHECKING:
+    from pyomnilogic_local.omnilogic import OmniLogic
+
 
 class Bow(OmniEquipment[MSPBoW, TelemetryBoW]):
     """Represents a bow in the OmniLogic system."""
@@ -25,8 +29,8 @@ class Bow(OmniEquipment[MSPBoW, TelemetryBoW]):
     chlorinator: Chlorinator | None = None
     csads: EquipmentDict[CSAD] = EquipmentDict()
 
-    def __init__(self, _api: OmniLogicAPI, mspconfig: MSPBoW, telemetry: Telemetry) -> None:
-        super().__init__(_api, mspconfig, telemetry)
+    def __init__(self, omni: "OmniLogic", mspconfig: MSPBoW, telemetry: Telemetry) -> None:
+        super().__init__(omni, mspconfig, telemetry)
 
     @property
     def equip_type(self) -> str:
@@ -52,60 +56,60 @@ def _update_filters(self, mspconfig: MSPBoW, telemetry: Telemetry) -> None:
             self.filters = EquipmentDict()
             return
 
-        self.filters = EquipmentDict([Filter(self._api, filter_, telemetry) for filter_ in mspconfig.filter])
+        self.filters = EquipmentDict([Filter(self._omni, filter_, telemetry) for filter_ in mspconfig.filter])
 
     def _update_heater(self, mspconfig: MSPBoW, telemetry: Telemetry) -> None:
         """Update the heater based on the MSP configuration."""
         if mspconfig.heater is None:
             self.heater = None
             return
 
-        self.heater = Heater(self._api, mspconfig.heater, telemetry)
+        self.heater = Heater(self._omni, mspconfig.heater, telemetry)
 
     def _update_relays(self, mspconfig: MSPBoW, telemetry: Telemetry) -> None:
         """Update the relays based on the MSP configuration."""
         if mspconfig.relay is None:
             self.relays = EquipmentDict()
             return
 
-        self.relays = EquipmentDict([Relay(self._api, relay, telemetry) for relay in mspconfig.relay])
+        self.relays = EquipmentDict([Relay(self._omni, relay, telemetry) for relay in mspconfig.relay])
 
     def _update_sensors(self, mspconfig: MSPBoW, telemetry: Telemetry) -> None:
         """Update the sensors based on the MSP configuration."""
         if mspconfig.sensor is None:
             self.sensors = EquipmentDict()
             return
 
-        self.sensors = EquipmentDict([Sensor(self._api, sensor, telemetry) for sensor in mspconfig.sensor])
+        self.sensors = EquipmentDict([Sensor(self._omni, sensor, telemetry) for sensor in mspconfig.sensor])
 
     def _update_lights(self, mspconfig: MSPBoW, telemetry: Telemetry) -> None:
         """Update the lights based on the MSP configuration."""
         if mspconfig.colorlogic_light is None:
             self.lights = EquipmentDict()
             return
 
-        self.lights = EquipmentDict([ColorLogicLight(self._api, light, telemetry) for light in mspconfig.colorlogic_light])
+        self.lights = EquipmentDict([ColorLogicLight(self._omni, light, telemetry) for light in mspconfig.colorlogic_light])
 
     def _update_pumps(self, mspconfig: MSPBoW, telemetry: Telemetry) -> None:
         """Update the pumps based on the MSP configuration."""
         if mspconfig.pump is None:
             self.pumps = EquipmentDict()
             return
 
-        self.pumps = EquipmentDict([Pump(self._api, pump, telemetry) for pump in mspconfig.pump])
+        self.pumps = EquipmentDict([Pump(self._omni, pump, telemetry) for pump in mspconfig.pump])
 
     def _update_chlorinators(self, mspconfig: MSPBoW, telemetry: Telemetry) -> None:
         """Update the chlorinators based on the MSP configuration."""
         if mspconfig.chlorinator is None:
             self.chlorinator = None
             return
 
-        self.chlorinator = Chlorinator(self._api, mspconfig.chlorinator, telemetry)
+        self.chlorinator = Chlorinator(self._omni, mspconfig.chlorinator, telemetry)
 
     def _update_csads(self, mspconfig: MSPBoW, telemetry: Telemetry) -> None:
         """Update the CSADs based on the MSP configuration."""
         if mspconfig.csad is None:
             self.csads = EquipmentDict()
             return
 
-        self.csads = EquipmentDict([CSAD(self._api, csad, telemetry) for csad in mspconfig.csad])
+        self.csads = EquipmentDict([CSAD(self._omni, csad, telemetry) for csad in mspconfig.csad])
diff --git a/pyomnilogic_local/colorlogiclight.py b/pyomnilogic_local/colorlogiclight.py
@@ -1,7 +1,8 @@
 import logging
+from typing import TYPE_CHECKING
 
 from pyomnilogic_local._base import OmniEquipment
-from pyomnilogic_local.api.api import OmniLogicAPI
+from pyomnilogic_local.decorators import auto_refresh
 from pyomnilogic_local.models.mspconfig import MSPColorLogicLight
 from pyomnilogic_local.models.telemetry import Telemetry, TelemetryColorLogicLight
 from pyomnilogic_local.omnitypes import (
@@ -13,14 +14,17 @@
 )
 from pyomnilogic_local.util import OmniEquipmentNotInitializedError
 
+if TYPE_CHECKING:
+    from pyomnilogic_local.omnilogic import OmniLogic
+
 _LOGGER = logging.getLogger(__name__)
 
 
 class ColorLogicLight(OmniEquipment[MSPColorLogicLight, TelemetryColorLogicLight]):
     """Represents a color logic light."""
 
-    def __init__(self, _api: OmniLogicAPI, mspconfig: MSPColorLogicLight, telemetry: Telemetry) -> None:
-        super().__init__(_api, mspconfig, telemetry)
+    def __init__(self, omni: "OmniLogic", mspconfig: MSPColorLogicLight, telemetry: Telemetry) -> None:
+        super().__init__(omni, mspconfig, telemetry)
 
     @property
     def model(self) -> ColorLogicLightType:
@@ -68,23 +72,23 @@ def special_effect(self) -> int:
         """Returns the current special effect."""
         return self.telemetry.special_effect
 
+    @auto_refresh()
     async def turn_on(self) -> None:
         """Turns the light on."""
         if self.bow_id is None or self.system_id is None:
             raise OmniEquipmentNotInitializedError("Cannot turn on light: bow_id or system_id is None")
         await self._api.async_set_equipment(self.bow_id, self.system_id, True)
 
+    @auto_refresh()
     async def turn_off(self) -> None:
         """Turns the light off."""
         if self.bow_id is None or self.system_id is None:
             raise OmniEquipmentNotInitializedError("Cannot turn off light: bow_id or system_id is None")
         await self._api.async_set_equipment(self.bow_id, self.system_id, False)
 
+    @auto_refresh()
     async def set_show(
-        self,
-        show: LightShows | None = None,
-        speed: ColorLogicSpeed | None = None,
-        brightness: ColorLogicBrightness | None = None,
+        self, show: LightShows | None = None, speed: ColorLogicSpeed | None = None, brightness: ColorLogicBrightness | None = None
     ) -> None:
         """Sets the light show, speed, and brightness."""
 
diff --git a/pyomnilogic_local/decorators.py b/pyomnilogic_local/decorators.py
@@ -0,0 +1,96 @@
+"""Decorators for automatic state management in pyomnilogic_local."""
+
+import asyncio
+import functools
+import logging
+import time
+from collections.abc import Callable
+from typing import Any, TypeVar, cast
+
+_LOGGER = logging.getLogger(__name__)
+
+F = TypeVar("F", bound=Callable[..., Any])
+
+
+def auto_refresh(
+    update_mspconfig: bool = False,
+    update_telemetry: bool = True,
+    delay: float = 1.25,
+) -> Callable[[F], F]:
+    """Decorator to automatically refresh OmniLogic state after method execution.
+
+    This decorator will:
+    1. Execute the decorated method
+    2. Wait for the specified delay (to allow controller to update)
+    3. Refresh telemetry/mspconfig if they're older than the post-delay time
+
+    The decorator is lock-safe: if multiple decorated methods are called concurrently,
+    only one refresh will occur thanks to the update_if_older_than mechanism.
+
+    Args:
+        update_mspconfig: Whether to refresh MSPConfig after method execution
+        update_telemetry: Whether to refresh Telemetry after method execution
+        delay: Time in seconds to wait after method completes before refreshing
+
+    Usage:
+        @auto_refresh()  # Default: telemetry only, 0.25s delay
+        async def turn_on(self, auto_refresh: bool | None = None):
+            ...
+
+        @auto_refresh(update_mspconfig=True, delay=0.5)
+        async def configure(self, auto_refresh: bool | None = None):
+            ...
+
+    The decorated method can accept an optional `auto_refresh` parameter:
+    - If None (default): Uses the OmniLogic instance's auto_refresh_enabled setting
+    - If True: Forces auto-refresh regardless of instance setting
+    - If False: Disables auto-refresh for this call
+    """
+
+    def decorator(func: F) -> F:
+        @functools.wraps(func)
+        async def wrapper(*args: Any, **kwargs: Any) -> Any:
+            # Extract the 'auto_refresh' parameter if provided
+            auto_refresh_param = kwargs.pop("auto_refresh", None)
+
+            # First arg should be 'self' (equipment instance)
+            if not args:
+                raise RuntimeError("@auto_refresh decorator requires a method with 'self' parameter")
+
+            self_obj = args[0]
+
+            # Get the OmniLogic instance
+            # Equipment classes should have _omni attribute pointing to parent OmniLogic
+            if hasattr(self_obj, "_omni") and self_obj._omni is not None:  # pylint: disable=protected-access
+                omni = self_obj._omni  # pylint: disable=protected-access
+            elif hasattr(self_obj, "auto_refresh_enabled"):
+                # This IS the OmniLogic instance
+                omni = self_obj
+            else:
+                raise RuntimeError("@auto_refresh decorator requires equipment to have '_omni' attribute or be used on OmniLogic methods")
+
+            # Determine if we should auto-refresh
+            should_refresh = auto_refresh_param if auto_refresh_param is not None else omni.auto_refresh_enabled
+
+            # Execute the original method
+            result = await func(*args, **kwargs)
+
+            # Perform auto-refresh if enabled
+            if should_refresh:
+                # Wait for the controller to process the change
+                await asyncio.sleep(delay)
+
+                # Calculate the target time (after delay)
+                target_time = time.time()
+
+                # Update only if data is older than target time
+                await omni.update_if_older_than(
+                    telemetry_min_time=target_time if update_telemetry else None,
+                    mspconfig_min_time=target_time if update_mspconfig else None,
+                )
+
+            return result
+
+        return cast(F, wrapper)
+
+    return decorator
diff --git a/pyomnilogic_local/omnilogic.py b/pyomnilogic_local/omnilogic.py
