Skip to content

Commit

Permalink
Use rich for logging and cli print outs (#1568)
Browse files Browse the repository at this point in the history
If rich is installed on the system, use it for logging and pretty
printing to make the output nicer.
  • Loading branch information
rytilahti committed Nov 2, 2022
1 parent 34d3f0e commit e6a0c45
Show file tree
Hide file tree
Showing 4 changed files with 43 additions and 24 deletions.
22 changes: 17 additions & 5 deletions miio/cli.py
@@ -1,4 +1,5 @@
import logging
from typing import Any, Dict

import click

Expand Down Expand Up @@ -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")
Expand Down
25 changes: 12 additions & 13 deletions miio/click_common.py
Expand Up @@ -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__)


Expand Down Expand Up @@ -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):
Expand Down Expand Up @@ -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)

Expand Down Expand Up @@ -304,15 +303,15 @@ 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):
result_msg = result_msg_fmt(**kwargs)
else:
result_msg = result_msg_fmt.format(**kwargs)
if result_msg:
click.echo(result_msg.strip())
echo(result_msg.strip())

return wrap

Expand All @@ -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)
Expand All @@ -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

Expand Down
15 changes: 11 additions & 4 deletions 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


Expand All @@ -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)
Expand Down Expand Up @@ -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()
Expand All @@ -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)}")
5 changes: 3 additions & 2 deletions miio/miioprotocol.py
Expand Up @@ -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
Expand Down Expand Up @@ -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",
Expand Down Expand Up @@ -208,7 +209,7 @@ def send(
self.port,
header["ts"],
payload["id"],
payload,
pf(payload),
)
if "error" in payload:
self._handle_error(payload["error"])
Expand Down

0 comments on commit e6a0c45

Please sign in to comment.