Line data    Source code

       1             : # SPDX-FileCopyrightText: 2025 PairInteraction Developers
       2             : # SPDX-License-Identifier: LGPL-3.0-or-later
       3           1 : from __future__ import annotations
       4             : 
       5           1 : import datetime
       6           1 : import inspect
       7           1 : import logging
       8           1 : import re
       9           1 : from functools import wraps
      10           1 : from typing import TYPE_CHECKING, ClassVar, TypeVar
      11             : 
      12           1 : from colorama import Fore, Style, just_fix_windows_console
      13             : 
      14           1 : from pairinteraction._backend import get_pending_logs
      15             : 
      16             : if TYPE_CHECKING:
      17             :     from collections.abc import Callable
      18             : 
      19             :     from typing_extensions import ParamSpec
      20             : 
      21             :     P = ParamSpec("P")
      22             :     R = TypeVar("R")
      23             : 
      24             : 
      25           1 : def _extract_cpp_backend_log_fields(message: str) -> dict[str, str]:
      26           1 :     pattern = (
      27             :         r"^\[(?P<timestamp>\d{4}-\d{2}-\d{2} \d{2}:\d{2}:\d{2}\.\d+)\s+(?P<thread>\d+)\]\s*"
      28             :         r".*"
      29             :         r"\[(?P<filename>[^:\]]+):(?P<lineno>\d+)\]\s*"
      30             :         r"(?P<message>.*)$"
      31             :     )
      32           1 :     match = re.match(pattern, message.strip(), re.DOTALL)
      33           1 :     if not match:
      34           0 :         raise RuntimeError(f"Could not parse log message: {message}")
      35           1 :     return match.groupdict()
      36             : 
      37             : 
      38           1 : def _log_cpp_backend_record(level: int, message: str) -> None:
      39           1 :     logger = logging.getLogger("cpp")
      40           1 :     fields = _extract_cpp_backend_log_fields(message)
      41           1 :     record = logging.LogRecord(
      42             :         name=logger.name,
      43             :         level=level,
      44             :         pathname=fields["filename"],
      45             :         lineno=int(fields["lineno"]),
      46             :         msg=fields["message"],
      47             :         args=(),
      48             :         exc_info=None,
      49             :     )
      50           1 :     record.created = datetime.datetime.strptime(fields["timestamp"], "%Y-%m-%d %H:%M:%S.%f").timestamp()
      51           1 :     record.thread = int(fields["thread"])
      52           1 :     if level >= logger.getEffectiveLevel():
      53           1 :         logger.handle(record)
      54             : 
      55             : 
      56           1 : def _flush_pending_logs() -> None:
      57           1 :     for entry in get_pending_logs():
      58           1 :         _log_cpp_backend_record(entry.level, entry.message.decode("utf-8", errors="replace"))
      59             : 
      60             : 
      61           1 : def _flush_logs_after(func: Callable[P, R]) -> Callable[P, R]:
      62           1 :     @wraps(func)
      63           1 :     def wrapper(*args: P.args, **kwargs: P.kwargs) -> R:
      64           1 :         result = func(*args, **kwargs)
      65           1 :         _flush_pending_logs()
      66           1 :         return result
      67             : 
      68           1 :     return wrapper
      69             : 
      70             : 
      71           1 : def decorate_module_with_flush_logs(module: object) -> None:
      72           1 :     for name, obj in vars(module).items():
      73           1 :         if inspect.isclass(obj):
      74           1 :             for attr_name, attr in vars(obj).items():
      75           1 :                 if callable(attr) and not attr_name.startswith("__"):
      76           1 :                     setattr(obj, attr_name, _flush_logs_after(attr))
      77           1 :         elif callable(obj) and not name.startswith("__"):
      78           1 :             setattr(module, name, _flush_logs_after(obj))
      79             : 
      80             : 
      81           1 : def configure_logging(
      82             :     level_str: str = "WARNING",
      83             :     fmt: str = ("[%(asctime)s.%(msecs)03d] [%(levelname)s] [%(filename)s:%(lineno)d] %(message)s"),
      84             : ) -> None:
      85             :     """Configure colorfully formatted logging."""
      86             : 
      87           0 :     class ColoredFormatter(logging.Formatter):
      88           0 :         COLORS: ClassVar = {
      89             :             "DEBUG": Fore.BLUE,
      90             :             "INFO": Fore.GREEN,
      91             :             "WARNING": Fore.YELLOW,
      92             :             "ERROR": Fore.RED,
      93             :             "CRITICAL": Fore.RED + Style.BRIGHT,
      94             :         }
      95             : 
      96           0 :         def format(self, record: logging.LogRecord) -> str:
      97           0 :             original_levelname = record.levelname
      98           0 :             record.levelname = f"{self.COLORS[record.levelname]}{record.levelname}{Style.RESET_ALL}"
      99           0 :             formatted = super().format(record)
     100           0 :             record.levelname = original_levelname
     101           0 :             return formatted
     102             : 
     103           0 :     just_fix_windows_console()
     104             : 
     105           0 :     handler = logging.StreamHandler()
     106           0 :     handler.setFormatter(ColoredFormatter(fmt, datefmt="%H:%M:%S"))
     107             : 
     108           0 :     root_logger = logging.getLogger()
     109           0 :     if root_logger.hasHandlers():
     110           0 :         root_logger.handlers.clear()
     111           0 :     root_logger.setLevel(getattr(logging, level_str.upper()))
     112           0 :     root_logger.addHandler(handler)