Coverage for src/km3flux/logger.py: 59%

1#!/usr/bin/env python3 

2from hashlib import sha256 

3import os 

4import re 

5import sys 

6import logging 

7import logging.handlers 

8 

9loggers = {} # this holds all the registered loggers 

10 

11 

12ATTRIBUTES = dict( 

13 list( 

14 zip( 

15 ["bold", "dark", "", "underline", "blink", "", "reverse", "concealed"], 

16 list(range(1, 9)), 

17 ) 

18 ) 

19) 

20del ATTRIBUTES[""] 

21 

22ATTRIBUTES_RE = r"\033\[(?:%s)m" % "|".join(["%d" % v for v in ATTRIBUTES.values()]) 

23 

24HIGHLIGHTS = dict( 

25 list( 

26 zip( 

27 [ 

28 "on_grey", 

29 "on_red", 

30 "on_green", 

31 "on_yellow", 

32 "on_blue", 

33 "on_magenta", 

34 "on_cyan", 

35 "on_white", 

36 ], 

37 list(range(40, 48)), 

38 ) 

39 ) 

40) 

41 

42HIGHLIGHTS_RE = r"\033\[(?:%s)m" % "|".join(["%d" % v for v in HIGHLIGHTS.values()]) 

43 

44COLORS = dict( 

45 list( 

46 zip( 

47 [ 

48 "grey", 

49 "red", 

50 "green", 

51 "yellow", 

52 "blue", 

53 "magenta", 

54 "cyan", 

55 "white", 

56 ], 

57 list(range(30, 38)), 

58 ) 

59 ) 

60) 

61 

62COLORS_RE = r"\033\[(?:%s)m" % "|".join(["%d" % v for v in COLORS.values()]) 

63 

64RESET = r"\033[0m" 

65RESET_RE = r"\033\[0m" 

66 

67 

68def supports_color(): 

69 """Checks if the terminal supports color.""" 

70 if isnotebook(): 

71 return True 

72 supported_platform = sys.platform != "win32" or "ANSICON" in os.environ 

73 is_a_tty = hasattr(sys.stdout, "isatty") and sys.stdout.isatty() 

74 

75 if not supported_platform or not is_a_tty: 

76 return False 

77 

78 return True 

79 

80 

81def colored(text, color=None, on_color=None, attrs=None, ansi_code=None): 

82 """Colorize text, while stripping nested ANSI color sequences. 

83 

84 Author: Konstantin Lepa <konstantin.lepa@gmail.com> / termcolor 

85 

86 Available text colors: 

87 red, green, yellow, blue, magenta, cyan, white. 

88 Available text highlights: 

89 on_red, on_green, on_yellow, on_blue, on_magenta, on_cyan, on_white. 

90 Available attributes: 

91 bold, dark, underline, blink, reverse, concealed. 

92 Example: 

93 colored('Hello, World!', 'red', 'on_grey', ['blue', 'blink']) 

94 colored('Hello, World!', 'green') 

95 """ 

96 if os.getenv("ANSI_COLORS_DISABLED") is None: 

97 if ansi_code is not None: 

98 return "\033[38;5;{}m{}\033[0m".format(ansi_code, text) 

99 fmt_str = "\033[%dm%s" 

100 if color is not None: 

101 text = re.sub(COLORS_RE + "(.*?)" + RESET_RE, r"\1", text) 

102 text = fmt_str % (COLORS[color], text) 

103 if on_color is not None: 

104 text = re.sub(HIGHLIGHTS_RE + "(.*?)" + RESET_RE, r"\1", text) 

105 text = fmt_str % (HIGHLIGHTS[on_color], text) 

106 if attrs is not None: 

107 text = re.sub(ATTRIBUTES_RE + "(.*?)" + RESET_RE, r"\1", text) 

108 for attr in attrs: 

109 text = fmt_str % (ATTRIBUTES[attr], text) 

110 return text + RESET 

111 else: 

112 return text 

113 

114 

115def hash_coloured_escapes(text): 

116 """Return the ANSI hash colour prefix and suffix for a given text""" 

117 ansi_code = int(sha256(text.encode("utf-8")).hexdigest(), 16) % 230 

118 prefix, suffix = colored("SPLIT", ansi_code=ansi_code).split("SPLIT") 

119 return prefix, suffix 

120 

121 

122def isnotebook(): 

123 """Check if running within a Jupyter notebook""" 

124 try: 

125 shell = get_ipython().__class__.__name__ 

126 if shell == "ZMQInteractiveShell": 

127 return True # Jupyter notebook or qtconsole 

128 elif shell == "TerminalInteractiveShell": 

129 return False # Terminal running IPython 

130 else: 

131 return False # Other type (?) 

132 except NameError: 

133 return False 

134 

135 

136def get_logger(name, filename=None, stream_loglevel="INFO", file_loglevel="DEBUG"): 

137 """Helper function to get a logger""" 

138 if name in loggers: 

139 return loggers[name] 

140 logger = logging.getLogger(name) 

141 logger.propagate = False 

142 

143 with_color = supports_color() 

144 

145 pre1, suf1 = hash_coloured_escapes(name) if with_color else ("", "") 

146 pre2, suf2 = hash_coloured_escapes(name + "salt") if with_color else ("", "") 

147 formatter = logging.Formatter( 

148 "%(asctime)s %(levelname)s {}+{}+{} " 

149 "%(name)s: %(message)s".format(pre1, pre2, suf1), 

150 datefmt="%Y-%m-%d %H:%M:%S", 

151 ) 

152 if filename is not None: 

153 ch_file = logging.handlers.RotatingFileHandler( 

154 filename, maxBytes=5 * 1024 * 1024, backupCount=10 

155 ) 

156 ch_file.setLevel(file_loglevel) 

157 ch_file.setFormatter(formatter) 

158 logger.addHandler(ch_file) 

159 ch = logging.StreamHandler() 

160 ch.setLevel(stream_loglevel) 

161 ch.setFormatter(formatter) 

162 logger.addHandler(ch) 

163 

164 loggers[name] = logger 

165 

166 logger.once_dict = {} 

167 

168 return logger 

169 

170 

171def set_level(log_or_name, level): 

172 """Set the log level for given logger""" 

173 if isinstance(log_or_name, str): 

174 log = get_logger(log_or_name) 

175 else: 

176 log = log_or_name 

177 log.setLevel(level) 

178 for handler in log.handlers: 

179 handler.setLevel(level) 

180 

181 

182if supports_color(): 

183 logging.addLevelName( 

184 logging.INFO, "\033[1;32m%s\033[1;0m" % logging.getLevelName(logging.INFO) 

185 ) 

186 logging.addLevelName( 

187 logging.DEBUG, "\033[1;34m%s\033[1;0m" % logging.getLevelName(logging.DEBUG) 

188 ) 

189 logging.addLevelName( 

190 logging.WARNING, "\033[1;33m%s\033[1;0m" % logging.getLevelName(logging.WARNING) 

191 ) 

192 logging.addLevelName( 

193 logging.ERROR, "\033[1;31m%s\033[1;0m" % logging.getLevelName(logging.ERROR) 

194 ) 

195 logging.addLevelName( 

196 logging.CRITICAL, 

197 "\033[1;101m%s\033[1;0m" % logging.getLevelName(logging.CRITICAL), 

198 ) 

199 

200 

201log = get_logger("km3flux") 

202set_level(log, "WARNING")