feat: add control methods to chlorinators

Author: cryptk

pyomnilogic_local/chlorinator.py

@@ -1,7 +1,13 @@
 from pyomnilogic_local._base import OmniEquipment
+from pyomnilogic_local.decorators import dirties_state
 from pyomnilogic_local.models.mspconfig import MSPChlorinator
 from pyomnilogic_local.models.telemetry import TelemetryChlorinator
-from pyomnilogic_local.omnitypes import ChlorinatorOperatingMode, ChlorinatorStatus
+from pyomnilogic_local.omnitypes import (
+    ChlorinatorCellType,
+    ChlorinatorOperatingMode,
+    ChlorinatorStatus,
+)
+from pyomnilogic_local.util import OmniEquipmentNotInitializedError
 
 
 class Chlorinator(OmniEquipment[MSPChlorinator, TelemetryChlorinator]):
@@ -54,7 +60,7 @@ def dispenser_type(self) -> str:
         return self.mspconfig.dispenser_type
 
     @property
-    def cell_type(self) -> str:
+    def cell_type(self) -> ChlorinatorCellType:
         """Type of T-Cell installed (e.g., T3, T5, T9, T15)."""
         return self.mspconfig.cell_type
 
@@ -321,3 +327,71 @@ def is_ready(self) -> bool:
 
         # Then check chlorinator-specific readiness
         return self.is_authenticated and not self.has_error
+
+    # Control methods
+    @dirties_state()
+    async def turn_on(self) -> None:
+        """Turn the chlorinator on (enable it).
+
+        Raises:
+            OmniEquipmentNotInitializedError: If bow_id is None.
+        """
+        if self.bow_id is None:
+            raise OmniEquipmentNotInitializedError("Cannot turn on chlorinator: bow_id is None")
+        await self._api.async_set_chlorinator_enable(self.bow_id, True)
+
+    @dirties_state()
+    async def turn_off(self) -> None:
+        """Turn the chlorinator off (disable it).
+
+        Raises:
+            OmniEquipmentNotInitializedError: If bow_id is None.
+        """
+        if self.bow_id is None:
+            raise OmniEquipmentNotInitializedError("Cannot turn off chlorinator: bow_id is None")
+        await self._api.async_set_chlorinator_enable(self.bow_id, False)
+
+    @dirties_state()
+    async def set_timed_percent(self, percent: int) -> None:
+        """Set the timed percent for chlorine generation.
+
+        Args:
+            percent: The chlorine generation percentage (0-100)
+
+        Raises:
+            OmniEquipmentNotInitializedError: If bow_id or system_id is None.
+            ValueError: If percent is outside the valid range (0-100).
+
+        Note:
+            This method uses the async_set_chlorinator_params API which requires
+            all chlorinator configuration parameters. The current values from
+            mspconfig are used for unchanged parameters.
+        """
+        if self.bow_id is None or self.system_id is None:
+            raise OmniEquipmentNotInitializedError("Cannot set timed percent: bow_id or system_id is None")
+
+        if not 0 <= percent <= 100:
+            raise ValueError(f"Timed percent {percent} is outside valid range [0, 100]")
+
+        # Get the parent Bow to determine bow_type
+        # We need to find our bow in the backyard
+        if (bow := self._omni.backyard.bow.get(self.bow_id)) is None:
+            raise OmniEquipmentNotInitializedError(f"Cannot find bow with id {self.bow_id}")
+
+        # Map equipment type to numeric bow_type value
+        # BOW_POOL = 0, BOW_SPA = 1 (based on typical protocol values)
+        bow_type = 0 if bow.equip_type == "BOW_POOL" else 1
+
+        # Get operating mode from telemetry (it's already an int or enum with .value)
+        op_mode = self.telemetry.operating_mode if isinstance(self.telemetry.operating_mode, int) else self.telemetry.operating_mode.value
+
+        await self._api.async_set_chlorinator_params(
+            pool_id=self.bow_id,
+            equipment_id=self.system_id,
+            timed_percent=percent,
+            cell_type=self.mspconfig.cell_type.value,  # ChlorinatorCellType is now IntEnum, use .value
+            op_mode=op_mode,
+            sc_timeout=self.mspconfig.superchlor_timeout,
+            bow_type=bow_type,
+            orp_timeout=self.mspconfig.orp_timeout,
+        )

pyomnilogic_local/models/mspconfig.py

@@ -219,6 +219,21 @@ class MSPChlorinator(OmniBase):
     cell_type: ChlorinatorCellType = Field(alias="Cell-Type")
     chlorinator_equipment: list[MSPChlorinatorEquip] | None = None
 
+    @model_validator(mode="before")
+    @classmethod
+    def convert_cell_type(cls, data: Any) -> Any:
+        """Convert cell_type string to ChlorinatorCellType enum by name."""
+        if isinstance(data, dict) and "Cell-Type" in data:
+            cell_type_str = data["Cell-Type"]
+            if isinstance(cell_type_str, str):
+                # Parse by enum member name (e.g., "CELL_TYPE_T15" -> ChlorinatorCellType.CELL_TYPE_T15)
+                try:
+                    data["Cell-Type"] = ChlorinatorCellType[cell_type_str]
+                except KeyError:
+                    # If not found, try to parse as int or leave as-is for Pydantic to handle
+                    pass
+        return data
+
     def __init__(self, **data: Any) -> None:
         super().__init__(**data)
 

pyomnilogic_local/omnitypes.py

@@ -155,36 +155,18 @@ class ChlorinatorDispenserType(StrEnum, PrettyEnum):
     TABLET = "TABLET_DISPENSING"
 
 
-class ChlorinatorCellType(StrEnum, PrettyEnum):
-    UNKNOWN = "CELL_TYPE_UNKNOWN"
-    T3 = "CELL_TYPE_T3"
-    T5 = "CELL_TYPE_T5"
-    T9 = "CELL_TYPE_T9"
-    T15 = "CELL_TYPE_T15"
-    T15_LS = "CELL_TYPE_T15_LS"
-    TCELLS315 = "CELL_TYPE_TCELLS315"
-    TCELLS325 = "CELL_TYPE_TCELLS325"
-    TCELLS340 = "CELL_TYPE_TCELLS340"
-    LIQUID = "CELL_TYPE_LIQUID"
-    TABLET = "CELL_TYPE_TABLET"
-
-    # There is probably an easier way to do this
-    def __int__(self) -> int:
-        return ChlorinatorCellInt[self.name].value
-
-
-class ChlorinatorCellInt(IntEnum, PrettyEnum):
-    UNKNOWN = 0
-    T3 = 1
-    T5 = 2
-    T9 = 3
-    T15 = 4
-    T15_LS = 5
-    TCELLS315 = 6
-    TCELLS325 = 7
-    TCELLS340 = 8
-    LIQUID = 9
-    TABLET = 10
+class ChlorinatorCellType(IntEnum, PrettyEnum):
+    CELL_TYPE_UNKNOWN = 0
+    CELL_TYPE_T3 = 1
+    CELL_TYPE_T5 = 2
+    CELL_TYPE_T9 = 3
+    CELL_TYPE_T15 = 4
+    CELL_TYPE_T15_LS = 5
+    CELL_TYPE_TCELLS315 = 6
+    CELL_TYPE_TCELLS325 = 7
+    CELL_TYPE_TCELLS340 = 8
+    CELL_TYPE_LIQUID = 9
+    CELL_TYPE_TABLET = 10
 
 
 # Lights