diff --git a/pyomnilogic_local/decorators.py b/pyomnilogic_local/decorators.py index 9252ad9..d135366 100644 --- a/pyomnilogic_local/decorators.py +++ b/pyomnilogic_local/decorators.py @@ -5,14 +5,18 @@ import functools import logging from collections.abc import Callable -from typing import Any, cast +from typing import Any, cast, overload from pyomnilogic_local.util import OmniEquipmentNotReadyError _LOGGER = logging.getLogger(__name__) -def control_method[F: Callable[..., Any]](func: F) -> F: +@overload +def control_method[FUNC: Callable[..., Any]](func: FUNC, *, check_ready: bool = ...) -> FUNC: ... +@overload +def control_method[FUNC: Callable[..., Any]](func: None = ..., *, check_ready: bool = ...) -> Callable[[FUNC], FUNC]: ... +def control_method[FUNC: Callable[..., Any]](func: FUNC | None = None, *, check_ready: bool = True) -> FUNC | Callable[[FUNC], FUNC]: """Check readiness and mark state as dirty. This decorator ensures equipment is ready before executing control methods and @@ -20,11 +24,16 @@ def control_method[F: Callable[..., Any]](func: F) -> F: pattern of checking is_ready and using @dirties_state() separately. The decorator: - 1. Checks if equipment is ready (via is_ready property) + 1. Optionally checks if equipment is ready (via is_ready property) 2. Raises OmniEquipmentNotReadyError with descriptive message if not ready 3. Executes the control method 4. Marks telemetry as dirty + Args: + func: The function being decorated. Supplied automatically when used without parentheses. + check_ready: If False, skip the is_ready check and execute unconditionally. + Defaults to True. + Raises: OmniEquipmentNotReadyError: If equipment is not ready to accept commands @@ -32,27 +41,38 @@ def control_method[F: Callable[..., Any]](func: F) -> F: @control_method async def turn_on(self) -> None: await self._api.async_set_equipment(...) + + @control_method(check_ready=False) + async def turn_off(self) -> None: + await self._api.async_set_equipment(...) """ - # Import here to avoid circular dependency - @functools.wraps(func) - async def wrapper(self: Any, *args: Any, **kwargs: Any) -> Any: - # Check if equipment is ready - if not self.is_ready: - # Generate descriptive error message from function name - action = func.__name__.replace("_", " ") - msg = f"Cannot {action}: equipment is not ready to accept commands" - raise OmniEquipmentNotReadyError(msg) + def decorator(f: FUNC) -> FUNC: + @functools.wraps(f) + async def wrapper(self: Any, *args: Any, **kwargs: Any) -> Any: + # Check if equipment is ready + if check_ready and not self.is_ready: + # Generate descriptive error message from function name + action = f.__name__.replace("_", " ") + msg = f"Cannot {action}: equipment is not ready to accept commands" + raise OmniEquipmentNotReadyError(msg) + + # Execute the original function + result = await f(self, *args, **kwargs) + + # Mark telemetry as dirty + if hasattr(self, "_omni"): + self._omni._telemetry_dirty = True + else: + _LOGGER.warning("%s does not have _omni reference, cannot mark state as dirty", self.__class__.__name__) - # Execute the original function - result = await func(self, *args, **kwargs) + return result - # Mark telemetry as dirty - if hasattr(self, "_omni"): - self._omni._telemetry_dirty = True - else: - _LOGGER.warning("%s does not have _omni reference, cannot mark state as dirty", self.__class__.__name__) + return cast("FUNC", wrapper) - return result + if func is not None: + # Used as @control_method without parentheses + return decorator(func) - return cast("F", wrapper) + # Used as @control_method(...) with parentheses + return decorator diff --git a/pyomnilogic_local/filter.py b/pyomnilogic_local/filter.py index 1ced2de..0326f89 100644 --- a/pyomnilogic_local/filter.py +++ b/pyomnilogic_local/filter.py @@ -236,7 +236,7 @@ async def turn_on(self) -> None: is_on=target_speed, ) - @control_method + @control_method(check_ready=False) async def turn_off(self) -> None: """Turn the filter off. diff --git a/pyomnilogic_local/pump.py b/pyomnilogic_local/pump.py index 2881b87..2169803 100644 --- a/pyomnilogic_local/pump.py +++ b/pyomnilogic_local/pump.py @@ -216,7 +216,7 @@ async def turn_on(self) -> None: is_on=target_speed, ) - @control_method + @control_method(check_ready=False) async def turn_off(self) -> None: """Turn the pump off.