Add module support & query their information during update cycle (#243)

* Add module support & modularize existing query

This creates a base to expose more features on the supported devices.
At the moment, the most visible change is that each update cycle gets information from all available modules:
* Basic system info
* Cloud (new)
* Countdown (new)
* Antitheft (new)
* Schedule (new)
* Time (existing, implements the time/timezone handling)
* Emeter (existing, partially separated from smartdevice)

* Fix imports

* Fix linting

* Use device host instead of alias in module repr

* Add property to list available modules, print them in cli state report

* usage: fix the get_realtime query

* separate usage from schedule to avoid multi-inheritance

* Fix module querying

* Add is_supported property to modules

kasa/cli.py

@@ -221,7 +221,7 @@ async def state(ctx, dev: SmartDevice):
         click.echo()
 
     click.echo(click.style("\t== Generic information ==", bold=True))
-    click.echo(f"\tTime:         {await dev.get_time()}")
+    click.echo(f"\tTime:         {dev.time} (tz: {dev.timezone}")
     click.echo(f"\tHardware:     {dev.hw_info['hw_ver']}")
     click.echo(f"\tSoftware:     {dev.hw_info['sw_ver']}")
     click.echo(f"\tMAC (rssi):   {dev.mac} ({dev.rssi})")
@@ -236,6 +236,13 @@ async def state(ctx, dev: SmartDevice):
         emeter_status = dev.emeter_realtime
         click.echo(f"\t{emeter_status}")
 
+    click.echo(click.style("\n\t== Modules ==", bold=True))
+    for module in dev.modules.values():
+        if module.is_supported:
+            click.echo(click.style(f"\t+ {module}", fg="green"))
+        else:
+            click.echo(click.style(f"\t- {module}", fg="red"))
+
 
 @cli.command()
 @pass_dev
@@ -430,7 +437,7 @@ async def led(dev, state):
 @pass_dev
 async def time(dev):
     """Get the device time."""
-    res = await dev.get_time()
+    res = dev.time
     click.echo(f"Current time: {res}")
     return res
 
@@ -488,5 +495,23 @@ async def reboot(plug, delay):
     return await plug.reboot(delay)
 
 
+@cli.group()
+@pass_dev
+async def schedule(dev):
+    """Scheduling commands."""
+
+
+@schedule.command(name="list")
+@pass_dev
+@click.argument("type", default="schedule")
+def _schedule_list(dev, type):
+    """Return the list of schedule actions for the given type."""
+    sched = dev.modules[type]
+    for rule in sched.rules:
+        print(rule)
+    else:
+        click.echo(f"No rules of type {type}")
+
+
 if __name__ == "__main__":
     cli()

kasa/modules/__init__.py

@@ -0,0 +1,10 @@
+# flake8: noqa
+from .antitheft import Antitheft
+from .cloud import Cloud
+from .countdown import Countdown
+from .emeter import Emeter
+from .module import Module
+from .rulemodule import Rule, RuleModule
+from .schedule import Schedule
+from .time import Time
+from .usage import Usage

kasa/modules/antitheft.py

@@ -0,0 +1,9 @@
+"""Implementation of the antitheft module."""
+from .rulemodule import RuleModule
+
+
+class Antitheft(RuleModule):
+    """Implementation of the antitheft module.
+
+    This shares the functionality among other rule-based modules.
+    """

kasa/modules/cloud.py

@@ -0,0 +1,50 @@
+"""Cloud module implementation."""
+from pydantic import BaseModel
+
+from .module import Module
+
+
+class CloudInfo(BaseModel):
+    """Container for cloud settings."""
+
+    binded: bool
+    cld_connection: int
+    fwDlPage: str
+    fwNotifyType: int
+    illegalType: int
+    server: str
+    stopConnect: int
+    tcspInfo: str
+    tcspStatus: int
+    username: str
+
+
+class Cloud(Module):
+    """Module implementing support for cloud services."""
+
+    def query(self):
+        """Request cloud connectivity info."""
+        return self.query_for_command("get_info")
+
+    @property
+    def info(self) -> CloudInfo:
+        """Return information about the cloud connectivity."""
+        return CloudInfo.parse_obj(self.data["get_info"])
+
+    def get_available_firmwares(self):
+        """Return list of available firmwares."""
+        return self.query_for_command("get_intl_fw_list")
+
+    def set_server(self, url: str):
+        """Set the update server URL."""
+        return self.query_for_command("set_server_url", {"server": url})
+
+    def connect(self, username: str, password: str):
+        """Login to the cloud using given information."""
+        return self.query_for_command(
+            "bind", {"username": username, "password": password}
+        )
+
+    def disconnect(self):
+        """Disconnect from the cloud."""
+        return self.query_for_command("unbind")

kasa/modules/countdown.py

@@ -0,0 +1,6 @@
+"""Implementation for the countdown timer."""
+from .rulemodule import RuleModule
+
+
+class Countdown(RuleModule):
+    """Implementation of countdown module."""

kasa/modules/emeter.py

@@ -0,0 +1,20 @@
+"""Implementation of the emeter module."""
+from ..emeterstatus import EmeterStatus
+from .usage import Usage
+
+
+class Emeter(Usage):
+    """Emeter module."""
+
+    def query(self):
+        """Prepare query for emeter data."""
+        return self._device._create_emeter_request()
+
+    @property  # type: ignore
+    def realtime(self) -> EmeterStatus:
+        """Return current energy readings."""
+        return EmeterStatus(self.data["get_realtime"])
+
+    async def erase_stats(self):
+        """Erase all stats."""
+        return await self.call("erase_emeter_stat")

kasa/modules/module.py

@@ -0,0 +1,59 @@
+"""Base class for all module implementations."""
+import collections
+from abc import ABC, abstractmethod
+from typing import TYPE_CHECKING
+
+if TYPE_CHECKING:
+    from kasa import SmartDevice
+
+
+# TODO: This is used for query construcing
+def merge(d, u):
+    """Update dict recursively."""
+    for k, v in u.items():
+        if isinstance(v, collections.abc.Mapping):
+            d[k] = merge(d.get(k, {}), v)
+        else:
+            d[k] = v
+    return d
+
+
+class Module(ABC):
+    """Base class implemention for all modules.
+
+    The base classes should implement `query` to return the query they want to be
+    executed during the regular update cycle.
+    """
+
+    def __init__(self, device: "SmartDevice", module: str):
+        self._device: "SmartDevice" = device
+        self._module = module
+
+    @abstractmethod
+    def query(self):
+        """Query to execute during the update cycle.
+
+        The inheriting modules implement this to include their wanted
+        queries to the query that gets executed when Device.update() gets called.
+        """
+
+    @property
+    def data(self):
+        """Return the module specific raw data from the last update."""
+        return self._device._last_update[self._module]
+
+    @property
+    def is_supported(self) -> bool:
+        """Return whether the module is supported by the device."""
+        return "err_code" not in self.data
+
+    def call(self, method, params=None):
+        """Call the given method with the given parameters."""
+        return self._device._query_helper(self._module, method, params)
+
+    def query_for_command(self, query, params=None):
+        """Create a request object for the given parameters."""
+        return self._device._create_request(self._module, query, params)
+
+    def __repr__(self) -> str:
+        return f"<Module {self.__class__.__name__} ({self._module}) for {self._device.host}>"

kasa/modules/rulemodule.py

@@ -0,0 +1,83 @@
+"""Base implementation for all rule-based modules."""
+import logging
+from enum import Enum
+from typing import Dict, List, Optional
+
+from pydantic import BaseModel
+
+from .module import Module, merge
+
+
+class Action(Enum):
+    """Action to perform."""
+
+    Disabled = -1
+    TurnOff = 0
+    TurnOn = 1
+    Unknown = 2
+
+
+class TimeOption(Enum):
+    """Time when the action is executed."""
+
+    Disabled = -1
+    Enabled = 0
+    AtSunrise = 1
+    AtSunset = 2
+
+
+class Rule(BaseModel):
+    """Representation of a rule."""
+
+    id: str
+    name: str
+    enable: bool
+    wday: List[int]
+    repeat: bool
+
+    # start action
+    sact: Optional[Action]
+    stime_opt: TimeOption
+    smin: int
+
+    eact: Optional[Action]
+    etime_opt: TimeOption
+    emin: int
+
+    # Only on bulbs
+    s_light: Optional[Dict]
+
+
+_LOGGER = logging.getLogger(__name__)
+
+
+class RuleModule(Module):
+    """Base class for rule-based modules, such as countdown and antitheft."""
+
+    def query(self):
+        """Prepare the query for rules."""
+        q = self.query_for_command("get_rules")
+        return merge(q, self.query_for_command("get_next_action"))
+
+    @property
+    def rules(self) -> List[Rule]:
+        """Return the list of rules for the service."""
+        try:
+            return [
+                Rule.parse_obj(rule) for rule in self.data["get_rules"]["rule_list"]
+            ]
+        except Exception as ex:
+            _LOGGER.error("Unable to read rule list: %s (data: %s)", ex, self.data)
+            return []
+
+    async def set_enabled(self, state: bool):
+        """Enable or disable the service."""
+        return await self.call("set_overall_enable", state)
+
+    async def delete_rule(self, rule: Rule):
+        """Delete the given rule."""
+        return await self.call("delete_rule", {"id": rule.id})
+
+    async def delete_all_rules(self):
+        """Delete all rules."""
+        return await self.call("delete_all_rules")

kasa/modules/schedule.py

@@ -0,0 +1,6 @@
+"""Schedule module implementation."""
+from .rulemodule import RuleModule
+
+
+class Schedule(RuleModule):
+    """Implements the scheduling interface."""

kasa/modules/time.py

@@ -0,0 +1,34 @@
+"""Provides the current time and timezone information."""
+from datetime import datetime
+
+from .module import Module, merge
+
+
+class Time(Module):
+    """Implements the timezone settings."""
+
+    def query(self):
+        """Request time and timezone."""
+        q = self.query_for_command("get_time")
+
+        merge(q, self.query_for_command("get_timezone"))
+        return q
+
+    @property
+    def time(self) -> datetime:
+        """Return current device time."""
+        res = self.data["get_time"]
+        return datetime(
+            res["year"],
+            res["month"],
+            res["mday"],
+            res["hour"],
+            res["min"],
+            res["sec"],
+        )
+
+    @property
+    def timezone(self):
+        """Return current timezone."""
+        res = self.data["get_timezone"]
+        return res

kasa/modules/usage.py

@@ -0,0 +1,40 @@
+"""Implementation of the usage interface."""
+from datetime import datetime
+
+from .module import Module, merge
+
+
+class Usage(Module):
+    """Baseclass for emeter/usage interfaces."""
+
+    def query(self):
+        """Return the base query."""
+        year = datetime.now().year
+        month = datetime.now().month
+
+        req = self.query_for_command("get_realtime")
+        req = merge(
+            req, self.query_for_command("get_daystat", {"year": year, "month": month})
+        )
+        req = merge(req, self.query_for_command("get_monthstat", {"year": year}))
+        req = merge(req, self.query_for_command("get_next_action"))
+
+        return req
+
+    async def get_daystat(self, year, month):
+        """Return stats for the current day."""
+        if year is None:
+            year = datetime.now().year
+        if month is None:
+            month = datetime.now().month
+        return await self.call("get_daystat", {"year": year, "month": month})
+
+    async def get_monthstat(self, year):
+        """Return stats for the current month."""
+        if year is None:
+            year = datetime.now().year
+        return await self.call("get_monthstat", {"year": year})
+
+    async def erase_stats(self):
+        """Erase all stats."""
+        return await self.call("erase_runtime_stat")

kasa/smartbulb.py

@@ -3,6 +3,7 @@ import logging
 import re
 from typing import Any, Dict, NamedTuple, cast
 
+from .modules import Antitheft, Cloud, Countdown, Emeter, Schedule, Time, Usage
 from .smartdevice import DeviceType, SmartDevice, SmartDeviceException, requires_update
 
 
@@ -109,13 +110,19 @@ class SmartBulb(SmartDevice):
     """
 
     LIGHT_SERVICE = "smartlife.iot.smartbulb.lightingservice"
-    TIME_SERVICE = "smartlife.iot.common.timesetting"
     SET_LIGHT_METHOD = "transition_light_state"
+    emeter_type = "smartlife.iot.common.emeter"
 
     def __init__(self, host: str) -> None:
         super().__init__(host=host)
-        self.emeter_type = "smartlife.iot.common.emeter"
         self._device_type = DeviceType.Bulb
+        self.add_module("schedule", Schedule(self, "smartlife.iot.common.schedule"))
+        self.add_module("usage", Usage(self, "smartlife.iot.common.schedule"))
+        self.add_module("antitheft", Antitheft(self, "smartlife.iot.common.anti_theft"))
+        self.add_module("time", Time(self, "smartlife.iot.common.timesetting"))
+        self.add_module("emeter", Emeter(self, self.emeter_type))
+        self.add_module("countdown", Countdown(self, "countdown"))
+        self.add_module("cloud", Cloud(self, "smartlife.iot.common.cloud"))
 
     @property  # type: ignore
     @requires_update

kasa/smartdevice.py

@@ -22,6 +22,7 @@ from typing import Any, Dict, List, Optional, Set
 
 from .emeterstatus import EmeterStatus
 from .exceptions import SmartDeviceException
+from .modules import Emeter, Module
 from .protocol import TPLinkSmartHomeProtocol
 
 _LOGGER = logging.getLogger(__name__)
@@ -186,6 +187,7 @@ class SmartDevice:
     """
 
     TIME_SERVICE = "time"
+    emeter_type = "emeter"
 
     def __init__(self, host: str) -> None:
         """Create a new SmartDevice instance.
@@ -195,7 +197,6 @@ class SmartDevice:
         self.host = host
 
         self.protocol = TPLinkSmartHomeProtocol(host)
-        self.emeter_type = "emeter"
         _LOGGER.debug("Initializing %s of type %s", self.host, type(self))
         self._device_type = DeviceType.Unknown
         # TODO: typing Any is just as using Optional[Dict] would require separate checks in
@@ -203,9 +204,21 @@ class SmartDevice:
         #       are not accessed incorrectly.
         self._last_update: Any = None
         self._sys_info: Any = None  # TODO: this is here to avoid changing tests
+        self.modules: Dict[str, Any] = {}
 
         self.children: List["SmartDevice"] = []
 
+    def add_module(self, name: str, module: Module):
+        """Register a module."""
+        if name in self.modules:
+            _LOGGER.debug("Module %s already registered, ignoring..." % name)
+            return
+
+        assert name not in self.modules
+
+        _LOGGER.debug("Adding module %s", module)
+        self.modules[name] = module
+
     def _create_request(
         self, target: str, cmd: str, arg: Optional[Dict] = None, child_ids=None
     ):
@@ -268,6 +281,14 @@ class SmartDevice:
             _LOGGER.debug("Device does not have feature information")
             return set()
 
+    @property  # type: ignore
+    @requires_update
+    def supported_modules(self) -> List[str]:
+        """Return a set of modules supported by the device."""
+        # TODO: this should rather be called `features`, but we don't want to break
+        #       the API now. Maybe just deprecate it and point the users to use this?
+        return list(self.modules.keys())
+
     @property  # type: ignore
     @requires_update
     def has_emeter(self) -> bool:
@@ -303,7 +324,12 @@ class SmartDevice:
             _LOGGER.debug(
                 "The device has emeter, querying its information along sysinfo"
             )
-            req.update(self._create_emeter_request())
+            self.add_module("emeter", Emeter(self, self.emeter_type))
+
+        for module in self.modules.values():
+            q = module.query()
+            _LOGGER.debug("Adding query for %s: %s", module, q)
+            req = merge(req, module.query())
 
         self._last_update = await self.protocol.query(req)
         self._sys_info = self._last_update["system"]["get_sysinfo"]
@@ -337,6 +363,18 @@ class SmartDevice:
         """Set the device name (alias)."""
         return await self._query_helper("system", "set_dev_alias", {"alias": alias})
 
+    @property  # type: ignore
+    @requires_update
+    def time(self) -> datetime:
+        """Return current time from the device."""
+        return self.modules["time"].time
+
+    @property  # type: ignore
+    @requires_update
+    def timezone(self) -> Dict:
+        """Return the current timezone."""
+        return self.modules["time"].timezone
+
     async def get_time(self) -> Optional[datetime]:
         """Return current time from the device, if available."""
         try:
@@ -435,7 +473,7 @@ class SmartDevice:
     def emeter_realtime(self) -> EmeterStatus:
         """Return current energy readings."""
         self._verify_emeter()
-        return EmeterStatus(self._last_update[self.emeter_type]["get_realtime"])
+        return EmeterStatus(self.modules["emeter"].realtime)
 
     async def get_emeter_realtime(self) -> EmeterStatus:
         """Retrieve current energy readings."""
@@ -555,7 +593,7 @@ class SmartDevice:
     async def erase_emeter_stats(self) -> Dict:
         """Erase energy meter statistics."""
         self._verify_emeter()
-        return await self._query_helper(self.emeter_type, "erase_emeter_stat", None)
+        return await self.modules["emeter"].erase_stats()
 
     @requires_update
     async def current_consumption(self) -> float:

kasa/smartplug.py

@@ -2,6 +2,7 @@
 import logging
 from typing import Any, Dict
 
+from kasa.modules import Antitheft, Cloud, Schedule, Time, Usage
 from kasa.smartdevice import DeviceType, SmartDevice, requires_update
 
 _LOGGER = logging.getLogger(__name__)
@@ -40,6 +41,11 @@ class SmartPlug(SmartDevice):
         super().__init__(host)
         self.emeter_type = "emeter"
         self._device_type = DeviceType.Plug
+        self.add_module("schedule", Schedule(self, "schedule"))
+        self.add_module("usage", Usage(self, "schedule"))
+        self.add_module("antitheft", Antitheft(self, "anti_theft"))
+        self.add_module("time", Time(self, "time"))
+        self.add_module("cloud", Cloud(self, "cnCloud"))
 
     @property  # type: ignore
     @requires_update

kasa/smartstrip.py

@@ -13,6 +13,8 @@ from kasa.smartdevice import (
 )
 from kasa.smartplug import SmartPlug
 
+from .modules import Antitheft, Countdown, Schedule, Time, Usage
+
 _LOGGER = logging.getLogger(__name__)
 
 
@@ -80,6 +82,11 @@ class SmartStrip(SmartDevice):
         super().__init__(host=host)
         self.emeter_type = "emeter"
         self._device_type = DeviceType.Strip
+        self.add_module("antitheft", Antitheft(self, "anti_theft"))
+        self.add_module("schedule", Schedule(self, "schedule"))
+        self.add_module("usage", Usage(self, "schedule"))
+        self.add_module("time", Time(self, "time"))
+        self.add_module("countdown", Countdown(self, "countdown"))
 
     @property  # type: ignore
     @requires_update