do_not_disturb.py

import datetime

from roborock.data import DnDTimer
from roborock.devices.traits.v1 import common
from roborock.roborock_typing import RoborockCommand

_ENABLED_PARAM = "enabled"


class DoNotDisturbTrait(DnDTimer, common.V1TraitMixin, common.RoborockSwitchBase):
    """Trait for managing Do Not Disturb (DND) settings on Roborock devices."""

    command = RoborockCommand.GET_DND_TIMER
    converter = common.DefaultConverter(DnDTimer)

    @property
    def is_on(self) -> bool:
        """Return whether the Do Not Disturb (DND) timer is enabled."""
        return self.enabled == 1

    async def set_dnd_timer(self, dnd_timer: DnDTimer) -> None:
        """Set the Do Not Disturb (DND) timer settings of the device."""
        await self.rpc_channel.send_command(RoborockCommand.SET_DND_TIMER, params=dnd_timer.as_list())
        await self.refresh()

    async def set_start_time(self, start_time: datetime.time) -> None:
        """Set the start time of the Do Not Disturb (DND) timer."""
        timer = DnDTimer(
            start_hour=start_time.hour,
            start_minute=start_time.minute,
            end_hour=self.end_hour,
            end_minute=self.end_minute,
            enabled=self.enabled,
        )
        await self.rpc_channel.send_command(RoborockCommand.SET_DND_TIMER, params=timer.as_list())
        await self.refresh()

    async def set_end_time(self, end_time: datetime.time) -> None:
        """Set the end time of the Do Not Disturb (DND) timer."""
        timer = DnDTimer(
            start_hour=self.start_hour,
            start_minute=self.start_minute,
            end_hour=end_time.hour,
            end_minute=end_time.minute,
            enabled=self.enabled,
        )
        await self.rpc_channel.send_command(RoborockCommand.SET_DND_TIMER, params=timer.as_list())
        await self.refresh()

    async def clear_dnd_timer(self) -> None:
        """Clear the Do Not Disturb (DND) timer settings of the device."""
        await self.rpc_channel.send_command(RoborockCommand.CLOSE_DND_TIMER)
        await self.refresh()

    async def enable(self) -> None:
        """Set the Do Not Disturb (DND) timer settings of the device."""
        await self.rpc_channel.send_command(
            RoborockCommand.SET_DND_TIMER,
            params=self.as_list(),
        )
        # Optimistic update to avoid an extra refresh
        self.enabled = 1

    async def disable(self) -> None:
        """Disable the Do Not Disturb (DND) timer settings of the device."""
        await self.rpc_channel.send_command(RoborockCommand.CLOSE_DND_TIMER)
        # Optimistic update to avoid an extra refresh
        self.enabled = 0