Implement changing the bulb turn-on behavior (#381)

* Implement changing the bulb turn-on behavior

* Improve docstrings a bit

* Improve docs and expose TurnOnBehavior(s)

* fix typing

docs/source/smartbulb.rst

@@ -66,3 +66,10 @@ API documentation
 .. autoclass:: kasa.SmartBulbPreset
     :members:
     :undoc-members:
+
+.. autoclass:: kasa.TurnOnBehaviors
+    :undoc-members:
+
+
+.. autoclass:: kasa.TurnOnBehavior
+    :undoc-members:

kasa/__init__.py

@@ -17,7 +17,7 @@ from kasa.discover import Discover
 from kasa.emeterstatus import EmeterStatus
 from kasa.exceptions import SmartDeviceException
 from kasa.protocol import TPLinkSmartHomeProtocol
-from kasa.smartbulb import SmartBulb, SmartBulbPreset
+from kasa.smartbulb import SmartBulb, SmartBulbPreset, TurnOnBehavior, TurnOnBehaviors
 from kasa.smartdevice import DeviceType, SmartDevice
 from kasa.smartdimmer import SmartDimmer
 from kasa.smartlightstrip import SmartLightStrip
@@ -32,6 +32,8 @@ __all__ = [
     "TPLinkSmartHomeProtocol",
     "SmartBulb",
     "SmartBulbPreset",
+    "TurnOnBehaviors",
+    "TurnOnBehavior",
     "DeviceType",
     "EmeterStatus",
     "SmartDevice",

kasa/cli.py

@@ -595,5 +595,36 @@ async def presets_modify(
     await dev.save_preset(preset)
 
 
+@cli.command()
+@pass_dev
+@click.option("--type", type=click.Choice(["soft", "hard"], case_sensitive=False))
+@click.option("--last", is_flag=True)
+@click.option("--preset", type=int)
+async def turn_on_behavior(dev: SmartBulb, type, last, preset):
+    """Modify bulb turn-on behavior."""
+    settings = await dev.get_turn_on_behavior()
+    click.echo(f"Current turn on behavior: {settings}")
+
+    # Return if we are not setting the value
+    if not type and not last and not preset:
+        return
+
+    # If we are setting the value, the type has to be specified
+    if (last or preset) and type is None:
+        click.echo("To set the behavior, you need to define --type")
+        return
+
+    behavior = getattr(settings, type)
+
+    if last:
+        click.echo(f"Going to set {type} to last")
+        behavior.preset = None
+    elif preset is not None:
+        click.echo(f"Going to set {type} to preset {preset}")
+        behavior.preset = preset
+
+    await dev.set_turn_on_behavior(settings)
+
+
 if __name__ == "__main__":
     cli()

kasa/smartbulb.py

@@ -1,9 +1,10 @@
 """Module for bulbs (LB*, KL*, KB*)."""
 import logging
 import re
-from typing import Any, Dict, List, NamedTuple, cast
+from enum import Enum
+from typing import Any, Dict, List, NamedTuple, Optional, cast
 
-from pydantic import BaseModel
+from pydantic import BaseModel, Field, root_validator
 
 from .modules import Antitheft, Cloud, Countdown, Emeter, Schedule, Time, Usage
 from .smartdevice import DeviceType, SmartDevice, SmartDeviceException, requires_update
@@ -34,6 +35,53 @@ class SmartBulbPreset(BaseModel):
     color_temp: int
 
 
+class BehaviorMode(str, Enum):
+    """Enum to present type of turn on behavior."""
+
+    Last = "last_status"
+    Preset = "customize_preset"
+
+
+class TurnOnBehavior(BaseModel):
+    """Model to present a single turn on behavior.
+
+    :param int preset: the index number of wanted preset.
+    :param BehaviorMode mode: last status or preset mode. If you are changing existing settings, you should not set this manually.
+
+    To change the behavior, it is only necessary to change the :ref:`preset` field
+    to contain either the preset index, or ``None`` for the last known state.
+    """
+
+    preset: Optional[int] = Field(alias="index", default=None)
+    mode: BehaviorMode
+
+    @root_validator
+    def mode_based_on_preset(cls, values):
+        """Set the mode based on the preset value."""
+        if values["preset"] is not None:
+            values["mode"] = BehaviorMode.Preset
+        else:
+            values["mode"] = BehaviorMode.Last
+
+        return values
+
+    class Config:
+        """Configuration to make the validator run when changing the values."""
+
+        validate_assignment = True
+
+
+class TurnOnBehaviors(BaseModel):
+    """Model to contain turn on behaviors.
+
+    :param TurnOnBehavior soft: the default setting to turn the bulb programmatically on
+    :param TurnOnBehavior hard: default setting when the bulb has been off from mains power.
+    """
+
+    soft: TurnOnBehavior = Field(alias="soft_on")
+    hard: TurnOnBehavior = Field(alias="hard_on")
+
+
 TPLINK_KELVIN = {
     "LB130": ColorTempRange(2500, 9000),
     "LB120": ColorTempRange(2700, 6500),
@@ -226,14 +274,21 @@ class SmartBulb(SmartDevice):
         """
         return await self._query_helper(self.LIGHT_SERVICE, "get_light_details")
 
-    async def get_turn_on_behavior(self) -> Dict:
-        """Return the behavior for turning the bulb on.
+    async def get_turn_on_behavior(self) -> TurnOnBehaviors:
+        """Return the behavior for turning the bulb on."""
+        return TurnOnBehaviors.parse_obj(
+            await self._query_helper(self.LIGHT_SERVICE, "get_default_behavior")
+        )
 
-        Example:
-            {'soft_on': {'mode': 'last_status'},
-            'hard_on': {'mode': 'last_status'}}
+    async def set_turn_on_behavior(self, behavior: TurnOnBehaviors):
+        """Set the behavior for turning the bulb on.
+
+        If you do not want to manually construct the behavior object,
+        you should use :func:`get_turn_on_behavior` to get the current settings.
         """
-        return await self._query_helper(self.LIGHT_SERVICE, "get_default_behavior")
+        return await self._query_helper(
+            self.LIGHT_SERVICE, "set_default_behavior", behavior.dict(by_alias=True)
+        )
 
     async def get_light_state(self) -> Dict[str, Dict]:
         """Query the light state."""