• 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@dataclass
15class GitlabSection:
16    id: str
17    header: str
18    type: LogSectionType
19    start_collapsed: bool = False
20    escape: str = "\x1b[0K"
21    colour: str = f"{CONSOLE_LOG['BOLD']}{CONSOLE_LOG['FG_GREEN']}"
22    __start_time: Optional[datetime] = field(default=None, init=False)
23    __end_time: Optional[datetime] = field(default=None, init=False)
24
25    @classmethod
26    def section_id_filter(cls, value) -> str:
27        return str(re.sub(r"[^\w_-]+", "-", value))
28
29    def __post_init__(self):
30        self.id = self.section_id_filter(self.id)
31
32    @property
33    def has_started(self) -> bool:
34        return self.__start_time is not None
35
36    @property
37    def has_finished(self) -> bool:
38        return self.__end_time is not None
39
40    def get_timestamp(self, time: datetime) -> str:
41        unix_ts = datetime.timestamp(time)
42        return str(int(unix_ts))
43
44    def section(self, marker: str, header: str, time: datetime) -> str:
45        preamble = f"{self.escape}section_{marker}"
46        collapse = marker == "start" and self.start_collapsed
47        collapsed = "[collapsed=true]" if collapse else ""
48        section_id = f"{self.id}{collapsed}"
49
50        timestamp = self.get_timestamp(time)
51        before_header = ":".join([preamble, timestamp, section_id])
52        colored_header = f"{self.colour}{header}\x1b[0m" if header else ""
53        header_wrapper = "\r" + f"{self.escape}{colored_header}"
54
55        return f"{before_header}{header_wrapper}"
56
57    def __enter__(self):
58        print(self.start())
59        return self
60
61    def __exit__(self, exc_type, exc_val, exc_tb):
62        print(self.end())
63
64    def start(self) -> str:
65        assert not self.has_finished, "Starting an already finished section"
66        self.__start_time = datetime.now()
67        return self.section(marker="start", header=self.header, time=self.__start_time)
68
69    def end(self) -> str:
70        assert self.has_started, "Ending an uninitialized section"
71        self.__end_time = datetime.now()
72        assert (
73            self.__end_time >= self.__start_time
74        ), "Section execution time will be negative"
75        return self.section(marker="end", header="", time=self.__end_time)
76
77    def delta_time(self) -> Optional[timedelta]:
78        if self.__start_time and self.__end_time:
79            return self.__end_time - self.__start_time
80
81        if self.has_started:
82            return datetime.now() - self.__start_time
83
84        return None
85