pyrox.dev
1# mypy: disable-error-code="no-untyped-call"
2# drop the above line when mypy is upgraded to include
3# https://github.com/python/typeshed/commit/49b717ca52bf0781a538b04c0d76a5513f7119b8
4import codecs
5import os
6import sys
7import time
8import unicodedata
9from contextlib import contextmanager
10from queue import Empty, Queue
11from typing import Any, Dict, Iterator
12from xml.sax.saxutils import XMLGenerator
13
14from colorama import Fore, Style
15
16
17class Logger:
18 def __init__(self) -> None:
19 self.logfile = os.environ.get("LOGFILE", "/dev/null")
20 self.logfile_handle = codecs.open(self.logfile, "wb")
21 self.xml = XMLGenerator(self.logfile_handle, encoding="utf-8")
22 self.queue: "Queue[Dict[str, str]]" = Queue()
23
24 self.xml.startDocument()
25 self.xml.startElement("logfile", attrs={})
26
27 self._print_serial_logs = True
28
29 @staticmethod
30 def _eprint(*args: object, **kwargs: Any) -> None:
31 print(*args, file=sys.stderr, **kwargs)
32
33 def close(self) -> None:
34 self.xml.endElement("logfile")
35 self.xml.endDocument()
36 self.logfile_handle.close()
37
38 def sanitise(self, message: str) -> str:
39 return "".join(ch for ch in message if unicodedata.category(ch)[0] != "C")
40
41 def maybe_prefix(self, message: str, attributes: Dict[str, str]) -> str:
42 if "machine" in attributes:
43 return f"{attributes['machine']}: {message}"
44 return message
45
46 def log_line(self, message: str, attributes: Dict[str, str]) -> None:
47 self.xml.startElement("line", attributes)
48 self.xml.characters(message)
49 self.xml.endElement("line")
50
51 def info(self, *args, **kwargs) -> None: # type: ignore
52 self.log(*args, **kwargs)
53
54 def warning(self, *args, **kwargs) -> None: # type: ignore
55 self.log(*args, **kwargs)
56
57 def error(self, *args, **kwargs) -> None: # type: ignore
58 self.log(*args, **kwargs)
59 sys.exit(1)
60
61 def log(self, message: str, attributes: Dict[str, str] = {}) -> None:
62 self._eprint(self.maybe_prefix(message, attributes))
63 self.drain_log_queue()
64 self.log_line(message, attributes)
65
66 def log_serial(self, message: str, machine: str) -> None:
67 self.enqueue({"msg": message, "machine": machine, "type": "serial"})
68 if self._print_serial_logs:
69 self._eprint(Style.DIM + f"{machine} # {message}" + Style.RESET_ALL)
70
71 def enqueue(self, item: Dict[str, str]) -> None:
72 self.queue.put(item)
73
74 def drain_log_queue(self) -> None:
75 try:
76 while True:
77 item = self.queue.get_nowait()
78 msg = self.sanitise(item["msg"])
79 del item["msg"]
80 self.log_line(msg, item)
81 except Empty:
82 pass
83
84 @contextmanager
85 def nested(self, message: str, attributes: Dict[str, str] = {}) -> Iterator[None]:
86 self._eprint(
87 self.maybe_prefix(
88 Style.BRIGHT + Fore.GREEN + message + Style.RESET_ALL, attributes
89 )
90 )
91
92 self.xml.startElement("nest", attrs={})
93 self.xml.startElement("head", attributes)
94 self.xml.characters(message)
95 self.xml.endElement("head")
96
97 tic = time.time()
98 self.drain_log_queue()
99 yield
100 self.drain_log_queue()
101 toc = time.time()
102 self.log(f"(finished: {message}, in {toc - tic:.2f} seconds)")
103
104 self.xml.endElement("nest")
105
106
107rootlog = Logger()