• Home
  • Line#
  • Scopes#
  • Navigate#
  • Raw
  • Download
1from __future__ import annotations
2
3import re
4from dataclasses import dataclass, field
5from datetime import datetime, timedelta
6from typing import TYPE_CHECKING, Optional
7
8from lava.utils.console_format import CONSOLE_LOG
9
10if TYPE_CHECKING:
11    from lava.utils.log_section import LogSectionType
12
13
14# TODO: Add section final status to assist with monitoring
15@dataclass
16class GitlabSection:
17    id: str
18    header: str
19    type: LogSectionType
20    start_collapsed: bool = False
21    escape: str = "\x1b[0K"
22    colour: str = f"{CONSOLE_LOG['BOLD']}{CONSOLE_LOG['FG_GREEN']}"
23    __start_time: Optional[datetime] = field(default=None, init=False)
24    __end_time: Optional[datetime] = field(default=None, init=False)
25
26    @classmethod
27    def section_id_filter(cls, value) -> str:
28        return str(re.sub(r"[^\w_-]+", "-", value))
29
30    def __post_init__(self):
31        self.id = self.section_id_filter(self.id)
32
33    @property
34    def has_started(self) -> bool:
35        return self.__start_time is not None
36
37    @property
38    def has_finished(self) -> bool:
39        return self.__end_time is not None
40
41    @property
42    def start_time(self) -> datetime:
43        return self.__start_time
44
45    @property
46    def end_time(self) -> Optional[datetime]:
47        return self.__end_time
48
49    def get_timestamp(self, time: datetime) -> str:
50        unix_ts = datetime.timestamp(time)
51        return str(int(unix_ts))
52
53    def section(self, marker: str, header: str, time: datetime) -> str:
54        preamble = f"{self.escape}section_{marker}"
55        collapse = marker == "start" and self.start_collapsed
56        collapsed = "[collapsed=true]" if collapse else ""
57        section_id = f"{self.id}{collapsed}"
58
59        timestamp = self.get_timestamp(time)
60        before_header = ":".join([preamble, timestamp, section_id])
61        colored_header = f"{self.colour}{header}\x1b[0m" if header else ""
62        header_wrapper = "\r" + f"{self.escape}{colored_header}"
63
64        return f"{before_header}{header_wrapper}"
65
66    def __str__(self) -> str:
67        status = "NS" if not self.has_started else "F" if self.has_finished else "IP"
68        delta = self.delta_time()
69        elapsed_time = "N/A" if delta is None else str(delta)
70        return (
71            f"GitlabSection({self.id}, {self.header}, {self.type}, "
72            f"SC={self.start_collapsed}, S={status}, ST={self.start_time}, "
73            f"ET={self.end_time}, ET={elapsed_time})"
74        )
75
76    def __enter__(self):
77        print(self.start())
78        return self
79
80    def __exit__(self, exc_type, exc_val, exc_tb):
81        print(self.end())
82
83    def start(self) -> str:
84        assert not self.has_finished, "Starting an already finished section"
85        self.__start_time = datetime.now()
86        return self.section(marker="start", header=self.header, time=self.__start_time)
87
88    def end(self) -> str:
89        assert self.has_started, "Ending an uninitialized section"
90        self.__end_time = datetime.now()
91        assert (
92            self.__end_time >= self.__start_time
93        ), "Section execution time will be negative"
94        return self.section(marker="end", header="", time=self.__end_time)
95
96    def delta_time(self) -> Optional[timedelta]:
97        if self.__start_time and self.__end_time:
98            return self.__end_time - self.__start_time
99
100        if self.has_started:
101            return datetime.now() - self.__start_time
102
103        return None
104