# mypy: disable-error-code="no-untyped-call" # drop the above line when mypy is upgraded to include # https://github.com/python/typeshed/commit/49b717ca52bf0781a538b04c0d76a5513f7119b8 import codecs import os import sys import time import unicodedata from contextlib import contextmanager from queue import Empty, Queue from typing import Any, Dict, Iterator from xml.sax.saxutils import XMLGenerator from colorama import Fore, Style class Logger: def __init__(self) -> None: self.logfile = os.environ.get("LOGFILE", "/dev/null") self.logfile_handle = codecs.open(self.logfile, "wb") self.xml = XMLGenerator(self.logfile_handle, encoding="utf-8") self.queue: "Queue[Dict[str, str]]" = Queue() self.xml.startDocument() self.xml.startElement("logfile", attrs={}) self._print_serial_logs = True @staticmethod def _eprint(*args: object, **kwargs: Any) -> None: print(*args, file=sys.stderr, **kwargs) def close(self) -> None: self.xml.endElement("logfile") self.xml.endDocument() self.logfile_handle.close() def sanitise(self, message: str) -> str: return "".join(ch for ch in message if unicodedata.category(ch)[0] != "C") def maybe_prefix(self, message: str, attributes: Dict[str, str]) -> str: if "machine" in attributes: return f"{attributes['machine']}: {message}" return message def log_line(self, message: str, attributes: Dict[str, str]) -> None: self.xml.startElement("line", attributes) self.xml.characters(message) self.xml.endElement("line") def info(self, *args, **kwargs) -> None: # type: ignore self.log(*args, **kwargs) def warning(self, *args, **kwargs) -> None: # type: ignore self.log(*args, **kwargs) def error(self, *args, **kwargs) -> None: # type: ignore self.log(*args, **kwargs) sys.exit(1) def log(self, message: str, attributes: Dict[str, str] = {}) -> None: self._eprint(self.maybe_prefix(message, attributes)) self.drain_log_queue() self.log_line(message, attributes) def log_serial(self, message: str, machine: str) -> None: self.enqueue({"msg": message, "machine": machine, "type": "serial"}) if self._print_serial_logs: self._eprint(Style.DIM + f"{machine} # {message}" + Style.RESET_ALL) def enqueue(self, item: Dict[str, str]) -> None: self.queue.put(item) def drain_log_queue(self) -> None: try: while True: item = self.queue.get_nowait() msg = self.sanitise(item["msg"]) del item["msg"] self.log_line(msg, item) except Empty: pass @contextmanager def nested(self, message: str, attributes: Dict[str, str] = {}) -> Iterator[None]: self._eprint( self.maybe_prefix( Style.BRIGHT + Fore.GREEN + message + Style.RESET_ALL, attributes ) ) self.xml.startElement("nest", attrs={}) self.xml.startElement("head", attributes) self.xml.characters(message) self.xml.endElement("head") tic = time.time() self.drain_log_queue() yield self.drain_log_queue() toc = time.time() self.log(f"(finished: {message}, in {toc - tic:.2f} seconds)") self.xml.endElement("nest") rootlog = Logger()