feat: begin major refactor of codebase

- provide abstraction layer within library
- allows interacting with omnilogic in a more pythonic way

Author: cryptk

pyomnilogic_local/__init__.py

@@ -0,0 +1,3 @@
+from .omnilogic import OmniLogic
+
+__all__ = ["OmniLogic"]

pyomnilogic_local/_base.py

@@ -0,0 +1,61 @@
+from typing import Any
+
+from pyomnilogic_local.models import MSPEquipmentType, Telemetry
+
+
+class OmniEquipment:
+    """Base class for OmniLogic equipment."""
+
+    def __init__(self, mspconfig: MSPEquipmentType, telemetry: Telemetry | None = None) -> None:
+        """Initialize the equipment with configuration and telemetry data."""
+        # If the Equipment has subdevices, we don't store those as part of this device's config
+        # They will get parsed and stored as their own equipment instances
+        try:
+            self.mspconfig = mspconfig.without_subdevices()
+        except AttributeError:
+            self.mspconfig = mspconfig
+
+        if hasattr(self, "telemetry") and telemetry is not None:
+            self.telemetry = telemetry.get_telem_by_systemid(self.mspconfig.system_id)
+
+        # Populate fields from MSP configuration and telemetry
+        # This is some moderate magic to avoid having to manually set each field
+        # The TL;DR is that we loop over all fields defined in the MSPConfig and Telemetry models
+        # and set the corresponding attributes on this equipment instance.
+        for field in self.mspconfig.__class__.model_fields:
+            if getattr(self.mspconfig, field, None) is not None:
+                setattr(self, field, self._from_mspconfig(field))
+        for field in self.mspconfig.__class__.model_computed_fields:
+            if getattr(self.mspconfig, field, None) is not None:
+                setattr(self, field, self._from_mspconfig(field))
+        if hasattr(self, "telemetry") and self.telemetry is not None:
+            for field in self.telemetry.__class__.model_fields:
+                if getattr(self.telemetry, field, None) is not None:
+                    setattr(self, field, self._from_telemetry(field))
+            for field in self.telemetry.__class__.model_computed_fields:
+                if getattr(self.telemetry, field, None) is not None:
+                    setattr(self, field, self._from_telemetry(field))
+
+    def update_config(self, mspconfig: MSPEquipmentType) -> None:
+        """Update the configuration data for the equipment."""
+        if hasattr(self, "mspconfig"):
+            self.mspconfig = mspconfig.without_subdevices()
+        else:
+            raise NotImplementedError("This equipment does not have MSP configuration.")
+
+    def update_telemetry(self, telemetry: Telemetry) -> None:
+        """Update the telemetry data for the equipment."""
+        if hasattr(self, "telemetry"):
+            self.telemetry = telemetry.get_telem_by_systemid(self.mspconfig.system_id)
+        else:
+            raise NotImplementedError("This equipment does not have telemetry data.")
+
+    def _from_mspconfig(self, attribute: str) -> Any:
+        """Helper method to get a value from the MSP configuration."""
+        return getattr(self.mspconfig, attribute, None)
+
+    def _from_telemetry(self, attribute: str) -> Any:
+        """Helper method to get a value from the telemetry data."""
+        if hasattr(self, "telemetry"):
+            return getattr(self.telemetry, attribute, None)
+        return None

pyomnilogic_local/api/__init__.py

@@ -0,0 +1,5 @@
+from .api import OmniLogicAPI
+
+__all__ = [
+    "OmniLogicAPI",
+]

pyomnilogic_local/api.py pyomnilogic_local/api/api.pypyomnilogic_local/api.py renamed to pyomnilogic_local/api/api.py

@@ -6,13 +6,16 @@
 import xml.etree.ElementTree as ET
 from typing import Literal, overload
 
-from .models.filter_diagnostics import FilterDiagnostics
-from .models.mspconfig import MSPConfig
-from .models.telemetry import Telemetry
-from .models.util import to_pydantic
-from .omnitypes import (
+from pyomnilogic_local.models.filter_diagnostics import FilterDiagnostics
+from pyomnilogic_local.models.mspconfig import MSPConfig
+from pyomnilogic_local.models.telemetry import Telemetry
+
+from ..omnitypes import (
     ColorLogicBrightness,
-    ColorLogicShow,
+    ColorLogicShow25,
+    ColorLogicShow40,
+    ColorLogicShowUCL,
+    ColorLogicShowUCLV2,
     ColorLogicSpeed,
     HeaterMode,
     MessageType,
@@ -23,12 +26,12 @@
 
 
 class OmniLogicAPI:
-    def __init__(self, controller_ip: str, controller_port: int, response_timeout: float) -> None:
+    def __init__(self, controller_ip: str, controller_port: int, response_timeout: float = 5.0) -> None:
         self.controller_ip = controller_ip
         self.controller_port = controller_port
         self.response_timeout = response_timeout
-        self._loop = asyncio.get_running_loop()
-        self._protocol_factory = OmniLogicProtocol
+        # self._loop = asyncio.get_running_loop()
+        # self._protocol_factory = OmniLogicProtocol
 
     @overload
     async def async_send_message(self, message_type: MessageType, message: str | None, need_response: Literal[True]) -> str: ...
@@ -61,8 +64,13 @@ async def async_send_message(self, message_type: MessageType, message: str | Non
 
         return resp
 
-    @to_pydantic(pydantic_type=MSPConfig)
-    async def async_get_config(self) -> str:
+    @overload
+    async def async_get_mspconfig(self, raw: Literal[True]) -> str: ...
+    @overload
+    async def async_get_mspconfig(self, raw: Literal[False]) -> MSPConfig: ...
+    @overload
+    async def async_get_mspconfig(self) -> MSPConfig: ...
+    async def async_get_mspconfig(self, raw: bool = False) -> MSPConfig | str:
         """Retrieve the MSPConfig from the Omni, optionally parse it into a pydantic model.
 
         Args:
@@ -78,14 +86,21 @@ async def async_get_config(self) -> str:
 
         req_body = ET.tostring(body_element, xml_declaration=True, encoding="unicode")
 
-        return await self.async_send_message(MessageType.REQUEST_CONFIGURATION, req_body, True)
+        resp = await self.async_send_message(MessageType.REQUEST_CONFIGURATION, req_body, True)
 
-    @to_pydantic(pydantic_type=FilterDiagnostics)
-    async def async_get_filter_diagnostics(
-        self,
-        pool_id: int,
-        equipment_id: int,
-    ) -> str:
+        if raw:
+            return resp
+        return MSPConfig.load_xml(resp)
+
+    @overload
+    async def async_get_filter_diagnostics(self, pool_id: int, equipment_id: int, raw: Literal[True]) -> str: ...
+    @overload
+    async def async_get_filter_diagnostics(self, pool_id: int, equipment_id: int, raw: Literal[False]) -> FilterDiagnostics: ...
+    @overload
+    async def async_get_filter_diagnostics(self, pool_id: int, equipment_id: int) -> FilterDiagnostics: ...
+    @overload
+    async def async_get_filter_diagnostics(self, pool_id: int, equipment_id: int, raw: bool) -> FilterDiagnostics | str: ...
+    async def async_get_filter_diagnostics(self, pool_id: int, equipment_id: int, raw: bool = False) -> FilterDiagnostics | str:
         """Retrieve filter diagnostics from the Omni, optionally parse it into a pydantic model.
 
         Args:
@@ -108,10 +123,19 @@ async def async_get_filter_diagnostics(
 
         req_body = ET.tostring(body_element, xml_declaration=True, encoding="unicode")
 
-        return await self.async_send_message(MessageType.GET_FILTER_DIAGNOSTIC_INFO, req_body, True)
+        resp = await self.async_send_message(MessageType.GET_FILTER_DIAGNOSTIC_INFO, req_body, True)
+
+        if raw:
+            return resp
+        return FilterDiagnostics.load_xml(resp)
 
-    @to_pydantic(pydantic_type=Telemetry)
-    async def async_get_telemetry(self) -> str:
+    @overload
+    async def async_get_telemetry(self, raw: Literal[True]) -> str: ...
+    @overload
+    async def async_get_telemetry(self, raw: Literal[False]) -> Telemetry: ...
+    @overload
+    async def async_get_telemetry(self) -> Telemetry: ...
+    async def async_get_telemetry(self, raw: bool = False) -> Telemetry | str:
         """Retrieve the current telemetry data from the Omni, optionally parse it into a pydantic model.
 
         Returns:
@@ -124,7 +148,11 @@ async def async_get_telemetry(self) -> str:
 
         req_body = ET.tostring(body_element, xml_declaration=True, encoding="unicode")
 
-        return await self.async_send_message(MessageType.GET_TELEMETRY, req_body, True)
+        resp = await self.async_send_message(MessageType.GET_TELEMETRY, req_body, True)
+
+        if raw:
+            return resp
+        return Telemetry.load_xml(resp)
 
     async def async_set_heater(
         self,
@@ -352,7 +380,7 @@ async def async_set_light_show(
         self,
         pool_id: int,
         equipment_id: int,
-        show: ColorLogicShow,
+        show: ColorLogicShow25 | ColorLogicShow40 | ColorLogicShowUCL | ColorLogicShowUCLV2,
         speed: ColorLogicSpeed = ColorLogicSpeed.ONE_TIMES,
         brightness: ColorLogicBrightness = ColorLogicBrightness.ONE_HUNDRED_PERCENT,
         reserved: int = 0,

pyomnilogic_local/exceptions.py pyomnilogic_local/api/exceptions.pypyomnilogic_local/exceptions.py renamed to pyomnilogic_local/api/exceptions.py

@@ -4,7 +4,3 @@ class OmniLogicException(Exception):
 
 class OmniTimeoutException(OmniLogicException):
     pass
-
-
-class OmniParsingException(OmniLogicException):
-    pass

pyomnilogic_local/protocol.py pyomnilogic_local/api/protocol.pypyomnilogic_local/protocol.py renamed to pyomnilogic_local/api/protocol.py

@@ -9,9 +9,9 @@
 
 from typing_extensions import Self
 
+from ..models.leadmessage import LeadMessage
+from ..omnitypes import ClientType, MessageType
 from .exceptions import OmniTimeoutException
-from .models.leadmessage import LeadMessage
-from .omnitypes import ClientType, MessageType
 
 _LOGGER = logging.getLogger(__name__)
 
@@ -186,7 +186,7 @@ async def _wait_for_ack(self, ack_id: int) -> None:
                 exc = error_task.result()
                 if isinstance(exc, Exception):
                     raise exc
-                _LOGGER.error("Unknown error occurred during communication with Omnilogic: %s", exc)
+                _LOGGER.error("Unknown error occurred during communication with OmniLogic: %s", exc)
             if data_task in done:
                 message = data_task.result()
                 if message.id == ack_id:

pyomnilogic_local/backyard.py

@@ -0,0 +1,48 @@
+from pyomnilogic_local.models.mspconfig import MSPBackyard
+from pyomnilogic_local.models.telemetry import Telemetry
+
+from ._base import OmniEquipment
+from .bow import Bow
+from .colorlogiclight import ColorLogicLight
+from .relay import Relay
+from .sensor import Sensor
+
+
+class Backyard(OmniEquipment):
+    """Represents the backyard equipment in the OmniLogic system."""
+
+    bow: list[Bow] = []
+    lights: list[ColorLogicLight] = []
+    relays: list[Relay] = []
+    sensors: list[Sensor] = []
+
+    def __init__(self, mspconfig: MSPBackyard, telemetry: Telemetry) -> None:
+        super().__init__(mspconfig, telemetry)
+
+        self._update_bows(mspconfig, telemetry)
+        self._update_relays(mspconfig, telemetry)
+        self._update_sensors(mspconfig, telemetry)
+
+    def _update_bows(self, mspconfig: MSPBackyard, telemetry: Telemetry) -> None:
+        """Update the bows based on the MSP configuration."""
+        if mspconfig.bow is None:
+            self.bow = []
+            return
+
+        self.bow = [Bow(bow, telemetry) for bow in mspconfig.bow]
+
+    def _update_relays(self, mspconfig: MSPBackyard, telemetry: Telemetry) -> None:
+        """Update the relays based on the MSP configuration."""
+        if mspconfig.relay is None:
+            self.relays = []
+            return
+
+        self.relays = [Relay(relay, telemetry) for relay in mspconfig.relay]
+
+    def _update_sensors(self, mspconfig: MSPBackyard, telemetry: Telemetry) -> None:
+        """Update the sensors, bows, lights, and relays based on the MSP configuration."""
+        if mspconfig.sensor is None:
+            self.sensors = []
+            return
+
+        self.sensors = [Sensor(sensor, telemetry) for sensor in mspconfig.sensor]

pyomnilogic_local/bow.py

@@ -0,0 +1,5 @@
+from pyomnilogic_local._base import OmniEquipment
+
+
+class Bow(OmniEquipment):
+    """Represents a bow in the OmniLogic system."""

pyomnilogic_local/cli/debug/commands.py

@@ -5,7 +5,7 @@
 
 import click
 
-from pyomnilogic_local.api import OmniLogicAPI
+from pyomnilogic_local.api.api import OmniLogicAPI
 from pyomnilogic_local.cli import ensure_connection
 from pyomnilogic_local.cli.pcap_utils import parse_pcap_file, process_pcap_messages
 from pyomnilogic_local.cli.utils import async_get_filter_diagnostics
@@ -39,7 +39,7 @@ def get_mspconfig(ctx: click.Context) -> None:
     """
     ensure_connection(ctx)
     omni: OmniLogicAPI = ctx.obj["OMNI"]
-    mspconfig = asyncio.run(omni.async_get_config(raw=ctx.obj["RAW"]))
+    mspconfig = asyncio.run(omni.async_get_mspconfig(raw=ctx.obj["RAW"]))
     click.echo(mspconfig)
 
 
@@ -104,7 +104,7 @@ def get_filter_diagnostics(ctx: click.Context, pool_id: int, filter_id: int) ->
 @click.argument("pcap_file", type=click.Path(exists=True, path_type=Path))
 @click.pass_context
 def parse_pcap(ctx: click.Context, pcap_file: Path) -> None:
-    """Parse a PCAP file and reconstruct Omnilogic protocol communication.
+    """Parse a PCAP file and reconstruct OmniLogic protocol communication.
 
     Analyzes network packet captures to decode OmniLogic protocol messages.
     Automatically reassembles multi-part messages (LeadMessage + BlockMessages)

pyomnilogic_local/cli/get/lights.py

@@ -1,7 +1,7 @@
 # Need to figure out how to resolve the 'Untyped decorator makes function "..." untyped' errors in mypy when using click decorators
 # mypy: disable-error-code="misc"
 
-from typing import Any
+from typing import Any, cast
 
 import click
 
@@ -11,12 +11,11 @@
 )
 from pyomnilogic_local.models.telemetry import (
     Telemetry,
-    TelemetryType,
+    TelemetryColorLogicLight,
 )
 from pyomnilogic_local.omnitypes import (
     ColorLogicBrightness,
     ColorLogicPowerState,
-    ColorLogicShow,
     ColorLogicSpeed,
 )
 
@@ -41,21 +40,21 @@ def lights(ctx: click.Context) -> None:
     if mspconfig.backyard.colorlogic_light:
         for light in mspconfig.backyard.colorlogic_light:
             lights_found = True
-            _print_light_info(light, telemetry.get_telem_by_systemid(light.system_id))
+            _print_light_info(light, cast(TelemetryColorLogicLight, telemetry.get_telem_by_systemid(light.system_id)))
 
     # Check for lights in Bodies of Water
     if mspconfig.backyard.bow:
         for bow in mspconfig.backyard.bow:
             if bow.colorlogic_light:
                 for cl_light in bow.colorlogic_light:
                     lights_found = True
-                    _print_light_info(cl_light, telemetry.get_telem_by_systemid(cl_light.system_id))
+                    _print_light_info(cl_light, cast(TelemetryColorLogicLight, telemetry.get_telem_by_systemid(cl_light.system_id)))
 
     if not lights_found:
         click.echo("No ColorLogic lights found in the system configuration.")
 
 
-def _print_light_info(light: MSPColorLogicLight, telemetry: TelemetryType | None) -> None:
+def _print_light_info(light: MSPColorLogicLight, telemetry: TelemetryColorLogicLight | None) -> None:
     """Format and print light information in a nice table format.
 
     Args:
@@ -67,14 +66,15 @@ def _print_light_info(light: MSPColorLogicLight, telemetry: TelemetryType | None
     click.echo("=" * 60)
 
     light_data: dict[Any, Any] = {**dict(light), **dict(telemetry)} if telemetry else dict(light)
+
     for attr_name, value in light_data.items():
         if attr_name == "brightness":
             value = ColorLogicBrightness(value).pretty()
         elif attr_name == "effects" and isinstance(value, list):
             show_names = [show.pretty() if hasattr(show, "pretty") else str(show) for show in value]
             value = ", ".join(show_names) if show_names else "None"
         elif attr_name == "show" and value is not None:
-            value = ColorLogicShow(value).pretty()
+            value = telemetry.show_name(light.type, light.v2_active, True) if telemetry else str(value)
         elif attr_name == "speed":
             value = ColorLogicSpeed(value).pretty()
         elif attr_name == "state":