Coverage for src/flag_gems/logging_utils.py: 79%

1"""Logging helpers for flag_gems. 

2 

3Notes 

4----- 

51) When you enter through the public APIs `enable`, `only_enable`, or the 

6 context manager `use_gems`, the `record` flag controls whether op-level 

7 logging is enabled and where it is written. 

82) If you import `flag_gems` and call operators directly (e.g., `flag_gems.mm`) 

9 without those helpers, call `setup_flaggems_logging()` yourself to initialize 

10 the logging mode and file handler. 

11""" 

12 

13import logging 

14from pathlib import Path 

15 

16 

17class LogOncePerLocationFilter(logging.Filter): 

18 def __init__(self): 

19 super().__init__() 

20 self.logged_locations = set() 

21 

22 def filter(self, record): 

23 key = (record.pathname, record.lineno) 

24 if key in self.logged_locations: 

25 return False 

26 self.logged_locations.add(key) 

27 return True 

28 

29 

30def _remove_file_handlers(logger: logging.Logger): 

31 # Remove and close only the FileHandlers created by setup_flaggems_logging. 

32 # This avoids touching unrelated FileHandlers attached by other modules. 

33 removed = False 

34 for h in list(logger.handlers): 

35 if isinstance(h, logging.FileHandler) and getattr(h, "_flaggems_owned", False): 

36 h.close() 

37 logger.removeHandler(h) 

38 removed = True 

39 return removed 

40 

41 

42def setup_flaggems_logging(path=None, record=True, once=False): 

43 logger = logging.getLogger("flag_gems") 

44 

45 # If caller asks for recording, refresh file handler (new path overwrites old). 

46 if record: 

47 _remove_file_handlers(logger) 

48 else: 

49 return 

50 

51 filename = Path(path or Path.home() / ".flaggems/oplist.log") 

52 handler = logging.FileHandler(filename, mode="w") 

53 handler._flaggems_owned = True 

54 

55 if once: 

56 handler.addFilter(LogOncePerLocationFilter()) 

57 

58 formatter = logging.Formatter("[%(levelname)s] %(name)s.%(funcName)s: %(message)s") 

59 handler.setFormatter(formatter) 

60 

61 logger.setLevel(logging.DEBUG) 

62 logger.addHandler(handler) 

63 logger.propagate = False 

64 

65 

66def teardown_flaggems_logging(logger: logging.Logger | None = None): 

67 """Remove file handlers for the flag_gems logger (used on context exit).""" 

68 

69 logger = logger or logging.getLogger("flag_gems") 

70 _remove_file_handlers(logger)