Coverage for gpaw/new/logger.py: 81%

1from __future__ import annotations 

2 

3import contextlib 

4import io 

5import os 

6import sys 

7from functools import cache 

8from pathlib import Path 

9from typing import IO, Any, Sequence 

10 

11from gpaw.mpi import MPIComm, world 

12 

13 

14def indent(text: Any, indentation=' ') -> str: 

15 r"""Indent text blob. 

16 

17 >>> indent('line 1\nline 2', '..') 

18 '..line 1\n..line 2' 

19 """ 

20 if not isinstance(text, str): 

21 text = str(text) 

22 return indentation + text.replace('\n', '\n' + indentation) 

23 

24 

25class Logger: 

26 def __init__(self, 

27 filename: str | Path | IO[str] | None = '-', 

28 comm: MPIComm | Sequence[int] | None = None): 

29 if comm is None: 

30 comm = world 

31 elif not hasattr(comm, 'rank'): 

32 comm = world.new_communicator(list(comm)) 

33 

34 self.comm: MPIComm = comm # type: ignore 

35 

36 self.fd: IO[str] 

37 

38 if self.comm.rank > 0 or filename is None: 

39 self.fd = open(os.devnull, 'w', encoding='utf-8') 

40 self.close_fd = True 

41 elif filename == '-': 

42 self.fd = sys.stdout 

43 self.close_fd = False 

44 elif isinstance(filename, (str, Path)): 

45 self.fd = open(filename, 'w', encoding='utf-8') 

46 self.close_fd = True 

47 else: 

48 self.fd = filename 

49 self.close_fd = False 

50 

51 self.indentation = '' 

52 

53 self.use_colors = can_colorize(file=self.fd) 

54 if self.use_colors: 

55 self.green = '\x1b[32m' 

56 self.reset = '\x1b[0m' 

57 else: 

58 self.green = '' 

59 self.reset = '' 

60 

61 def __del__(self): 

62 self.close() 

63 

64 def close(self) -> None: 

65 if self.close_fd: 

66 self.fd.close() 

67 

68 @contextlib.contextmanager 

69 def indent(self, text): 

70 self(text) 

71 self.indentation += ' ' 

72 yield 

73 self.indentation = self.indentation[2:] 

74 

75 def __call__(self, *args, end=None, flush=False) -> None: 

76 if self.fd.closed: 

77 return 

78 i = self.indentation 

79 text = ' '.join(str(arg) for arg in args) 

80 if i: 

81 text = (i + text.replace('\n', '\n' + i)).rstrip(' ') 

82 print(text, file=self.fd, end=end, flush=flush) 

83 

84 

85def can_colorize(*, file: IO[str] | IO[bytes] | None = None) -> bool: 

86 """Code from Python 3.14b1: cpython/Lib/_colorize.py.""" 

87 ok = _can_colorize() 

88 if ok is not None: 

89 return ok 

90 

91 if file is None: 

92 file = sys.stdout 

93 

94 if not hasattr(file, 'fileno'): 

95 return False 

96 

97 try: 

98 return os.isatty(file.fileno()) 

99 except io.UnsupportedOperation: 

100 return hasattr(file, 'isatty') and file.isatty() 

101 

102 

103@cache 

104def _can_colorize() -> bool | None: 

105 """Check standard envvars for colors. 

106 

107 See https://docs.python.org/3/using/cmdline.html#controlling-color 

108 

109 Returns None if undecided. 

110 """ 

111 if not sys.flags.ignore_environment: 

112 if os.environ.get('PYTHON_COLORS') == '0': 

113 return False 

114 if os.environ.get('PYTHON_COLORS') == '1': 

115 return True 

116 if os.environ.get('NO_COLOR'): 

117 return False 

118 if os.environ.get('FORCE_COLOR'): 

119 return True 

120 if os.environ.get('TERM') == 'dumb': 

121 return False 

122 if sys.platform == 'win32': 

123 try: 

124 import nt 

125 

126 if not nt._supports_virtual_terminal(): 

127 return False 

128 except (ImportError, AttributeError): 

129 return False 

130 return None