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