From e6a0c45aab4b0c99545b9112f11363d497dd399a Mon Sep 17 00:00:00 2001 From: Teemu R Date: Thu, 3 Nov 2022 00:08:07 +0100 Subject: [PATCH] Use rich for logging and cli print outs (#1568) If rich is installed on the system, use it for logging and pretty printing to make the output nicer. --- miio/cli.py | 22 +++++++++++++++++----- miio/click_common.py | 25 ++++++++++++------------- miio/devtools/pcapparser.py | 15 +++++++++++---- miio/miioprotocol.py | 5 +++-- 4 files changed, 43 insertions(+), 24 deletions(-) diff --git a/miio/cli.py b/miio/cli.py index 653b219f0..5d722d142 100644 --- a/miio/cli.py +++ b/miio/cli.py @@ -1,4 +1,5 @@ import logging +from typing import Any, Dict import click @@ -29,11 +30,22 @@ @click.version_option() @click.pass_context def cli(ctx, debug: int, output: str): - if debug: - logging.basicConfig(level=logging.DEBUG) - _LOGGER.info("Debug mode active") - else: - logging.basicConfig(level=logging.INFO) + logging_config: Dict[str, Any] = { + "level": logging.DEBUG if debug > 0 else logging.INFO + } + try: + from rich.logging import RichHandler + + rich_config = { + "show_time": False, + } + logging_config["handlers"] = [RichHandler(**rich_config)] + logging_config["format"] = "%(message)s" + except ImportError: + pass + + # The configuration should be converted to use dictConfig, but this keeps mypy happy for now + logging.basicConfig(**logging_config) # type: ignore if output in ("json", "json_pretty"): output_func = json_output(pretty=output == "json_pretty") diff --git a/miio/click_common.py b/miio/click_common.py index 081dbc044..46ab00e0f 100644 --- a/miio/click_common.py +++ b/miio/click_common.py @@ -12,10 +12,13 @@ import click -import miio - from .exceptions import DeviceError +try: + from rich import print as echo +except ImportError: + echo = click.echo + _LOGGER = logging.getLogger(__name__) @@ -49,9 +52,8 @@ class ExceptionHandlerGroup(click.Group): def __call__(self, *args, **kwargs): try: return self.main(*args, **kwargs) - except (ValueError, miio.DeviceException) as ex: - _LOGGER.debug("Exception: %s", ex, exc_info=True) - click.echo(click.style("Error: %s" % ex, fg="red", bold=True)) + except Exception as ex: + _LOGGER.exception("Exception: %s", ex) class EnumType(click.Choice): @@ -179,10 +181,7 @@ def _wrap(self, *args, **kwargs): and self._model is None and self._info is None ): - _LOGGER.debug( - "Unknown model, trying autodetection. %s %s" - % (self._model, self._info) - ) + _LOGGER.debug("Unknown model, trying autodetection") self._fetch_info() return func(self, *args, **kwargs) @@ -304,7 +303,7 @@ def wrap(*args, **kwargs): else: msg = msg_fmt.format(**kwargs) if msg: - click.echo(msg.strip()) + echo(msg.strip()) kwargs["result"] = func(*args, **kwargs) if result_msg_fmt: if callable(result_msg_fmt): @@ -312,7 +311,7 @@ def wrap(*args, **kwargs): else: result_msg = result_msg_fmt.format(**kwargs) if result_msg: - click.echo(result_msg.strip()) + echo(result_msg.strip()) return wrap @@ -328,7 +327,7 @@ def wrap(*args, **kwargs): try: result = func(*args, **kwargs) except DeviceError as ex: - click.echo(json.dumps(ex.args[0], indent=indent)) + echo(json.dumps(ex.args[0], indent=indent)) return get_json_data_func = getattr(result, "__json__", None) @@ -337,7 +336,7 @@ def wrap(*args, **kwargs): result = get_json_data_func() elif data_variable is not None: result = data_variable - click.echo(json.dumps(result, indent=indent)) + echo(json.dumps(result, indent=indent)) return wrap diff --git a/miio/devtools/pcapparser.py b/miio/devtools/pcapparser.py index 3cc4385f2..7def63bae 100644 --- a/miio/devtools/pcapparser.py +++ b/miio/devtools/pcapparser.py @@ -1,10 +1,17 @@ """Parse PCAP files for miio traffic.""" from collections import Counter, defaultdict from ipaddress import ip_address +from pprint import pformat as pf from typing import List import click +try: + from rich import print as echo +except ImportError: + echo = click.echo + + from miio import Message @@ -14,7 +21,7 @@ def read_payloads_from_file(file, tokens: List[str]): import dpkt from dpkt.ethernet import ETH_TYPE_IP, Ethernet except ImportError: - print("You need to install dpkt to use this tool") # noqa: T201 + echo("You need to install dpkt to use this tool") return pcap = dpkt.pcap.Reader(file) @@ -70,9 +77,9 @@ def read_payloads_from_file(file, tokens: List[str]): yield src_addr, dst_addr, payload for cat in stats: - print(f"\n== {cat} ==") # noqa: T201 + echo(f"\n== {cat} ==") for stat, value in stats[cat].items(): - print(f"\t{stat}: {value}") # noqa: T201 + echo(f"\t{stat}: {value}") @click.command() @@ -81,4 +88,4 @@ def read_payloads_from_file(file, tokens: List[str]): def parse_pcap(file, token: List[str]): """Read PCAP file and output decrypted miio communication.""" for src_addr, dst_addr, payload in read_payloads_from_file(file, token): - print(f"{src_addr:<15} -> {dst_addr:<15} {payload}") # noqa: T201 + echo(f"{src_addr:<15} -> {dst_addr:<15} {pf(payload)}") diff --git a/miio/miioprotocol.py b/miio/miioprotocol.py index 958a62423..90e4d21e5 100644 --- a/miio/miioprotocol.py +++ b/miio/miioprotocol.py @@ -8,6 +8,7 @@ import logging import socket from datetime import datetime, timedelta +from pprint import pformat as pf from typing import Any, Dict, List import construct @@ -172,7 +173,7 @@ def send( msg = {"data": {"value": request}, "header": {"value": header}, "checksum": 0} m = Message.build(msg, token=self.token) - _LOGGER.debug("%s:%s >>: %s", self.ip, self.port, request) + _LOGGER.debug("%s:%s >>: %s", self.ip, self.port, pf(request)) if self.debug > 1: _LOGGER.debug( "send (timeout %s): %s", @@ -208,7 +209,7 @@ def send( self.port, header["ts"], payload["id"], - payload, + pf(payload), ) if "error" in payload: self._handle_error(payload["error"])