mirror of
https://github.com/python-kasa/python-kasa.git
synced 2024-12-23 03:33:35 +00:00
232 lines
7.2 KiB
Python
232 lines
7.2 KiB
Python
"""Common cli module."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import json
|
|
import re
|
|
import sys
|
|
from contextlib import contextmanager
|
|
from functools import singledispatch, update_wrapper, wraps
|
|
from typing import Final
|
|
|
|
import asyncclick as click
|
|
|
|
from kasa import (
|
|
Device,
|
|
)
|
|
|
|
# Value for optional options if passed without a value
|
|
OPTIONAL_VALUE_FLAG: Final = "_FLAG_"
|
|
|
|
# Block list of commands which require no update
|
|
SKIP_UPDATE_COMMANDS = ["raw-command", "command"]
|
|
|
|
pass_dev = click.make_pass_decorator(Device) # type: ignore[type-abstract]
|
|
|
|
|
|
try:
|
|
from rich import print as _echo
|
|
except ImportError:
|
|
# Strip out rich formatting if rich is not installed
|
|
# but only lower case tags to avoid stripping out
|
|
# raw data from the device that is printed from
|
|
# the device state.
|
|
rich_formatting = re.compile(r"\[/?[a-z]+]")
|
|
|
|
def _strip_rich_formatting(echo_func):
|
|
"""Strip rich formatting from messages."""
|
|
|
|
@wraps(echo_func)
|
|
def wrapper(message=None, *args, **kwargs):
|
|
if message is not None:
|
|
message = rich_formatting.sub("", message)
|
|
echo_func(message, *args, **kwargs)
|
|
|
|
return wrapper
|
|
|
|
_echo = _strip_rich_formatting(click.echo)
|
|
|
|
|
|
def echo(*args, **kwargs):
|
|
"""Print a message."""
|
|
ctx = click.get_current_context().find_root()
|
|
if "json" not in ctx.params or ctx.params["json"] is False:
|
|
_echo(*args, **kwargs)
|
|
|
|
|
|
def error(msg: str):
|
|
"""Print an error and exit."""
|
|
echo(f"[bold red]{msg}[/bold red]")
|
|
sys.exit(1)
|
|
|
|
|
|
def json_formatter_cb(result, **kwargs):
|
|
"""Format and output the result as JSON, if requested."""
|
|
if not kwargs.get("json"):
|
|
return
|
|
|
|
@singledispatch
|
|
def to_serializable(val):
|
|
"""Regular obj-to-string for json serialization.
|
|
|
|
The singledispatch trick is from hynek: https://hynek.me/articles/serialization/
|
|
"""
|
|
return str(val)
|
|
|
|
@to_serializable.register(Device)
|
|
def _device_to_serializable(val: Device):
|
|
"""Serialize smart device data, just using the last update raw payload."""
|
|
return val.internal_state
|
|
|
|
json_content = json.dumps(result, indent=4, default=to_serializable)
|
|
print(json_content)
|
|
|
|
|
|
def pass_dev_or_child(wrapped_function):
|
|
"""Pass the device or child to the click command based on the child options."""
|
|
child_help = (
|
|
"Child ID or alias for controlling sub-devices. "
|
|
"If no value provided will show an interactive prompt allowing you to "
|
|
"select a child."
|
|
)
|
|
child_index_help = "Child index controlling sub-devices"
|
|
|
|
@contextmanager
|
|
def patched_device_update(parent: Device, child: Device):
|
|
try:
|
|
orig_update = child.update
|
|
# patch child update method. Can be removed once update can be called
|
|
# directly on child devices
|
|
child.update = parent.update # type: ignore[method-assign]
|
|
yield child
|
|
finally:
|
|
child.update = orig_update # type: ignore[method-assign]
|
|
|
|
@click.pass_obj
|
|
@click.pass_context
|
|
@click.option(
|
|
"--child",
|
|
"--name",
|
|
is_flag=False,
|
|
flag_value=OPTIONAL_VALUE_FLAG,
|
|
default=None,
|
|
required=False,
|
|
type=click.STRING,
|
|
help=child_help,
|
|
)
|
|
@click.option(
|
|
"--child-index",
|
|
"--index",
|
|
required=False,
|
|
default=None,
|
|
type=click.INT,
|
|
help=child_index_help,
|
|
)
|
|
async def wrapper(ctx: click.Context, dev, *args, child, child_index, **kwargs):
|
|
if child := await _get_child_device(dev, child, child_index, ctx.info_name):
|
|
ctx.obj = ctx.with_resource(patched_device_update(dev, child))
|
|
dev = child
|
|
return await ctx.invoke(wrapped_function, dev, *args, **kwargs)
|
|
|
|
# Update wrapper function to look like wrapped function
|
|
return update_wrapper(wrapper, wrapped_function)
|
|
|
|
|
|
async def _get_child_device(
|
|
device: Device, child_option, child_index_option, info_command
|
|
) -> Device | None:
|
|
def _list_children():
|
|
return "\n".join(
|
|
[
|
|
f"{idx}: {child.device_id} ({child.alias})"
|
|
for idx, child in enumerate(device.children)
|
|
]
|
|
)
|
|
|
|
if child_option is None and child_index_option is None:
|
|
return None
|
|
|
|
if info_command in SKIP_UPDATE_COMMANDS:
|
|
# The device hasn't had update called (e.g. for cmd_command)
|
|
# The way child devices are accessed requires a ChildDevice to
|
|
# wrap the communications. Doing this properly would require creating
|
|
# a common interfaces for both IOT and SMART child devices.
|
|
# As a stop-gap solution, we perform an update instead.
|
|
await device.update()
|
|
|
|
if not device.children:
|
|
error(f"Device: {device.host} does not have children")
|
|
|
|
if child_option is not None and child_index_option is not None:
|
|
raise click.BadOptionUsage(
|
|
"child", "Use either --child or --child-index, not both."
|
|
)
|
|
|
|
if child_option is not None:
|
|
if child_option is OPTIONAL_VALUE_FLAG:
|
|
msg = _list_children()
|
|
child_index_option = click.prompt(
|
|
f"\n{msg}\nEnter the index number of the child device",
|
|
type=click.IntRange(0, len(device.children) - 1),
|
|
)
|
|
elif child := device.get_child_device(child_option):
|
|
echo(f"Targeting child device {child.alias}")
|
|
return child
|
|
else:
|
|
error(
|
|
"No child device found with device_id or name: "
|
|
f"{child_option} children are:\n{_list_children()}"
|
|
)
|
|
|
|
if child_index_option + 1 > len(device.children) or child_index_option < 0:
|
|
error(
|
|
f"Invalid index {child_index_option}, "
|
|
f"device has {len(device.children)} children"
|
|
)
|
|
child_by_index = device.children[child_index_option]
|
|
echo(f"Targeting child device {child_by_index.alias}")
|
|
return child_by_index
|
|
|
|
|
|
def CatchAllExceptions(cls):
|
|
"""Capture all exceptions and prints them nicely.
|
|
|
|
Idea from https://stackoverflow.com/a/44347763 and
|
|
https://stackoverflow.com/questions/52213375
|
|
"""
|
|
|
|
def _handle_exception(debug, exc):
|
|
if isinstance(exc, click.ClickException):
|
|
raise
|
|
# Handle exit request from click.
|
|
if isinstance(exc, click.exceptions.Exit):
|
|
sys.exit(exc.exit_code)
|
|
|
|
echo(f"Raised error: {exc}")
|
|
if debug:
|
|
raise
|
|
echo("Run with --debug enabled to see stacktrace")
|
|
sys.exit(1)
|
|
|
|
class _CommandCls(cls):
|
|
_debug = False
|
|
|
|
async def make_context(self, info_name, args, parent=None, **extra):
|
|
self._debug = any(
|
|
[arg for arg in args if arg in ["--debug", "-d", "--verbose", "-v"]]
|
|
)
|
|
try:
|
|
return await super().make_context(
|
|
info_name, args, parent=parent, **extra
|
|
)
|
|
except Exception as exc:
|
|
_handle_exception(self._debug, exc)
|
|
|
|
async def invoke(self, ctx):
|
|
try:
|
|
return await super().invoke(ctx)
|
|
except Exception as exc:
|
|
_handle_exception(self._debug, exc)
|
|
|
|
return _CommandCls
|