From 0996296350f027187d01ddae5b26d415841e605c Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Jos=C3=A9=20S=C3=A1nchez-Gallego?= <gallegoj@uw.edu>
Date: Thu, 23 Nov 2023 11:51:30 -0800
Subject: [PATCH] Add on/off/cycle commands

---
 src/lvmnps/actor/commands/__init__.py |  1 +
 src/lvmnps/actor/commands/onoff.py    | 95 +++++++++++++++++++++++++++
 src/lvmnps/nps/core.py                | 28 +++++++-
 3 files changed, 121 insertions(+), 3 deletions(-)
 create mode 100644 src/lvmnps/actor/commands/onoff.py

diff --git a/src/lvmnps/actor/commands/__init__.py b/src/lvmnps/actor/commands/__init__.py
index b4db3fb..2e99251 100644
--- a/src/lvmnps/actor/commands/__init__.py
+++ b/src/lvmnps/actor/commands/__init__.py
@@ -10,5 +10,6 @@
 
 from clu.parsers.click import command_parser as lvmnps_command_parser
 
+from .onoff import cycle, off, on
 from .refresh import refresh
 from .status import status
diff --git a/src/lvmnps/actor/commands/onoff.py b/src/lvmnps/actor/commands/onoff.py
new file mode 100644
index 0000000..ea2881e
--- /dev/null
+++ b/src/lvmnps/actor/commands/onoff.py
@@ -0,0 +1,95 @@
+#!/usr/bin/env python
+# -*- coding: utf-8 -*-
+#
+# @Author: José Sánchez-Gallego (gallegoj@uw.edu)
+# @Date: 2023-11-23
+# @Filename: onoff.py
+# @License: BSD 3-clause (http://www.opensource.org/licenses/BSD-3-Clause)
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+import click
+
+from . import lvmnps_command_parser
+
+
+if TYPE_CHECKING:
+    from src.lvmnps.actor.actor import NPSCommand
+
+
+__all__ = ["on", "off", "cycle"]
+
+
+class OutletNameParamType(click.ParamType):
+    name = "outlet_name"
+
+    def convert(self, value, param, ctx):
+        if isinstance(value, int):
+            return value
+
+        try:
+            return int(value, 10)
+        except ValueError:
+            return value
+
+
+@lvmnps_command_parser.command()
+@click.argument("OUTLETS", type=OutletNameParamType(), nargs=-1)
+async def on(command: NPSCommand, outlets: tuple[str | int, ...]):
+    """Turns an outlet or outlets on.
+
+    OUTLETS can be a single outlet name or outlet ID, or a list of them that will
+    be switched on as quickly as possible. An argument that can be casted into
+    an integer will be considered an outlet ID.
+
+    """
+
+    nps = command.actor.nps
+
+    outlet_data = await nps.set_state(outlets, on=True)
+
+    command.finish(outlets=[outlet.model_dump() for outlet in outlet_data])
+
+
+@lvmnps_command_parser.command()
+@click.argument("OUTLETS", type=OutletNameParamType(), nargs=-1)
+async def off(command: NPSCommand, outlets: tuple[str | int, ...]):
+    """Turns an outlet or outlets off.
+
+    OUTLETS can be a single outlet name or outlet ID, or a list of them that will
+    be switched on as quickly as possible. An argument that can be casted into
+    an integer will be considered an outlet ID.
+
+    """
+
+    nps = command.actor.nps
+
+    outlet_data = await nps.set_state(outlets, on=False)
+
+    command.finish(outlets=[outlet.model_dump() for outlet in outlet_data])
+
+
+@lvmnps_command_parser.command()
+@click.argument("OUTLETS", type=OutletNameParamType(), nargs=-1)
+@click.option(
+    "--delay",
+    type=float,
+    default=3,
+    help="Delay between off an on in seconds",
+)
+async def cycle(command: NPSCommand, outlets: tuple[str | int, ...], delay: float = 3):
+    """Cycles an outlet or outlets.
+
+    OUTLETS can be a single outlet name or outlet ID, or a list of them that will
+    be switched on as quickly as possible. An argument that can be casted into
+    an integer will be considered an outlet ID.
+
+    """
+
+    nps = command.actor.nps
+
+    outlet_data = await nps.cycle(outlets, delay=delay)
+
+    command.finish(outlets=[outlet.model_dump() for outlet in outlet_data])
diff --git a/src/lvmnps/nps/core.py b/src/lvmnps/nps/core.py
index 7c98ad7..4e332ec 100644
--- a/src/lvmnps/nps/core.py
+++ b/src/lvmnps/nps/core.py
@@ -97,7 +97,11 @@ async def refresh(self):
 
         pass
 
-    async def set_state(self, outlets: OutletArgType, on: bool = False):
+    async def set_state(
+        self,
+        outlets: OutletArgType,
+        on: bool = False,
+    ) -> list[OutletModel]:
         """Sets the state of an outlet or list of outlets.
 
         Parameters
@@ -132,6 +136,18 @@ async def set_state(self, outlets: OutletArgType, on: bool = False):
 
         await self.refresh()
 
+        # Gets the updated outlets we modified.
+        switched_outlets: list[OutletModel] = []
+        for _outlet in _outlets:
+            switched_outlets.append(
+                get_outlet_by_name(
+                    self.outlets,
+                    _outlet.normalised_name,
+                )
+            )
+
+        return switched_outlets
+
     @abc.abstractmethod
     async def _set_state_internal(self, outlets: list[OutletModel], on: bool = False):
         """Internal method for setting the outlet state.
@@ -142,7 +158,11 @@ async def _set_state_internal(self, outlets: list[OutletModel], on: bool = False
 
         """
 
-    async def cycle(self, outlets: OutletArgType, delay: float = 3):
+    async def cycle(
+        self,
+        outlets: OutletArgType,
+        delay: float = 3,
+    ) -> list[OutletModel]:
         """Turns off the selected outlets and turns them on again after a delay.
 
         Parameters
@@ -158,4 +178,6 @@ async def cycle(self, outlets: OutletArgType, delay: float = 3):
 
         await self.set_state(outlets, on=False)
         await asyncio.sleep(delay)
-        await self.set_state(outlets, on=True)
+        switched_outlets = await self.set_state(outlets, on=True)
+
+        return switched_outlets