1# Copyright 2010 Google Inc. All Rights Reserved.
2"""Logging helper module."""
3
4from __future__ import print_function
5
6# System modules
7import os.path
8import sys
9import traceback
10
11
12#TODO(yunlian@google.com): Use GetRoot from misc
13def GetRoot(scr_name):
14  """Break up pathname into (dir+name)."""
15  abs_path = os.path.abspath(scr_name)
16  return (os.path.dirname(abs_path), os.path.basename(abs_path))
17
18
19class Logger(object):
20  """Logging helper class."""
21
22  MAX_LOG_FILES = 10
23
24  def __init__(self, rootdir, basefilename, print_console, subdir='logs'):
25    logdir = os.path.join(rootdir, subdir)
26    basename = os.path.join(logdir, basefilename)
27
28    try:
29      os.makedirs(logdir)
30    except OSError:
31      pass
32      # print("Warning: Logs directory '%s' already exists." % logdir)
33
34    self.print_console = print_console
35
36    self._CreateLogFileHandles(basename)
37
38    self._WriteTo(self.cmdfd, ' '.join(sys.argv), True)
39
40  def _AddSuffix(self, basename, suffix):
41    return '%s%s' % (basename, suffix)
42
43  def _FindSuffix(self, basename):
44    timestamps = []
45    found_suffix = None
46    for i in range(self.MAX_LOG_FILES):
47      suffix = str(i)
48      suffixed_basename = self._AddSuffix(basename, suffix)
49      cmd_file = '%s.cmd' % suffixed_basename
50      if not os.path.exists(cmd_file):
51        found_suffix = suffix
52        break
53      timestamps.append(os.stat(cmd_file).st_mtime)
54
55    if found_suffix:
56      return found_suffix
57
58    # Try to pick the oldest file with the suffix and return that one.
59    suffix = str(timestamps.index(min(timestamps)))
60    # print ("Warning: Overwriting log file: %s" %
61    #       self._AddSuffix(basename, suffix))
62    return suffix
63
64  def _CreateLogFileHandle(self, name):
65    fd = None
66    try:
67      fd = open(name, 'w')
68    except IOError:
69      print('Warning: could not open %s for writing.' % name)
70    return fd
71
72  def _CreateLogFileHandles(self, basename):
73    suffix = self._FindSuffix(basename)
74    suffixed_basename = self._AddSuffix(basename, suffix)
75
76    self.cmdfd = self._CreateLogFileHandle('%s.cmd' % suffixed_basename)
77    self.stdout = self._CreateLogFileHandle('%s.out' % suffixed_basename)
78    self.stderr = self._CreateLogFileHandle('%s.err' % suffixed_basename)
79
80    self._CreateLogFileSymlinks(basename, suffixed_basename)
81
82  # Symlink unsuffixed basename to currently suffixed one.
83  def _CreateLogFileSymlinks(self, basename, suffixed_basename):
84    try:
85      for extension in ['cmd', 'out', 'err']:
86        src_file = '%s.%s' % (os.path.basename(suffixed_basename), extension)
87        dest_file = '%s.%s' % (basename, extension)
88        if os.path.exists(dest_file):
89          os.remove(dest_file)
90        os.symlink(src_file, dest_file)
91    except Exception as ex:
92      print('Exception while creating symlinks: %s' % str(ex))
93
94  def _WriteTo(self, fd, msg, flush):
95    if fd:
96      fd.write(msg)
97      if flush:
98        fd.flush()
99
100  def LogStartDots(self, print_to_console=True):
101    term_fd = self._GetStdout(print_to_console)
102    if term_fd:
103      term_fd.flush()
104      term_fd.write('. ')
105      term_fd.flush()
106
107  def LogAppendDot(self, print_to_console=True):
108    term_fd = self._GetStdout(print_to_console)
109    if term_fd:
110      term_fd.write('. ')
111      term_fd.flush()
112
113  def LogEndDots(self, print_to_console=True):
114    term_fd = self._GetStdout(print_to_console)
115    if term_fd:
116      term_fd.write('\n')
117      term_fd.flush()
118
119  def LogMsg(self, file_fd, term_fd, msg, flush=True):
120    if file_fd:
121      self._WriteTo(file_fd, msg, flush)
122    if self.print_console:
123      self._WriteTo(term_fd, msg, flush)
124
125  def _GetStdout(self, print_to_console):
126    if print_to_console:
127      return sys.stdout
128    return None
129
130  def _GetStderr(self, print_to_console):
131    if print_to_console:
132      return sys.stderr
133    return None
134
135  def LogCmdToFileOnly(self, cmd, machine='', user=None):
136    if not self.cmdfd:
137      return
138
139    host = ('%s@%s' % (user, machine)) if user else machine
140    flush = True
141    cmd_string = 'CMD (%s): %s\n' % (host, cmd)
142    self._WriteTo(self.cmdfd, cmd_string, flush)
143
144  def LogCmd(self, cmd, machine='', user=None, print_to_console=True):
145    if user:
146      host = '%s@%s' % (user, machine)
147    else:
148      host = machine
149
150    self.LogMsg(self.cmdfd, self._GetStdout(print_to_console),
151                'CMD (%s): %s\n' % (host, cmd))
152
153  def LogFatal(self, msg, print_to_console=True):
154    self.LogMsg(self.stderr, self._GetStderr(print_to_console),
155                'FATAL: %s\n' % msg)
156    self.LogMsg(self.stderr, self._GetStderr(print_to_console),
157                '\n'.join(traceback.format_stack()))
158    sys.exit(1)
159
160  def LogError(self, msg, print_to_console=True):
161    self.LogMsg(self.stderr, self._GetStderr(print_to_console),
162                'ERROR: %s\n' % msg)
163
164  def LogWarning(self, msg, print_to_console=True):
165    self.LogMsg(self.stderr, self._GetStderr(print_to_console),
166                'WARNING: %s\n' % msg)
167
168  def LogOutput(self, msg, print_to_console=True):
169    self.LogMsg(self.stdout, self._GetStdout(print_to_console),
170                'OUTPUT: %s\n' % msg)
171
172  def LogFatalIf(self, condition, msg):
173    if condition:
174      self.LogFatal(msg)
175
176  def LogErrorIf(self, condition, msg):
177    if condition:
178      self.LogError(msg)
179
180  def LogWarningIf(self, condition, msg):
181    if condition:
182      self.LogWarning(msg)
183
184  def LogCommandOutput(self, msg, print_to_console=True):
185    self.LogMsg(self.stdout,
186                self._GetStdout(print_to_console),
187                msg,
188                flush=False)
189
190  def LogCommandError(self, msg, print_to_console=True):
191    self.LogMsg(self.stderr,
192                self._GetStderr(print_to_console),
193                msg,
194                flush=False)
195
196  def Flush(self):
197    self.cmdfd.flush()
198    self.stdout.flush()
199    self.stderr.flush()
200
201
202class MockLogger(object):
203  """Logging helper class."""
204
205  MAX_LOG_FILES = 10
206
207  def __init__(self, *_args, **_kwargs):
208    self.stdout = sys.stdout
209    self.stderr = sys.stderr
210    return None
211
212  def _AddSuffix(self, basename, suffix):
213    return '%s%s' % (basename, suffix)
214
215  def _FindSuffix(self, basename):
216    timestamps = []
217    found_suffix = None
218    for i in range(self.MAX_LOG_FILES):
219      suffix = str(i)
220      suffixed_basename = self._AddSuffix(basename, suffix)
221      cmd_file = '%s.cmd' % suffixed_basename
222      if not os.path.exists(cmd_file):
223        found_suffix = suffix
224        break
225      timestamps.append(os.stat(cmd_file).st_mtime)
226
227    if found_suffix:
228      return found_suffix
229
230    # Try to pick the oldest file with the suffix and return that one.
231    suffix = str(timestamps.index(min(timestamps)))
232    # print ("Warning: Overwriting log file: %s" %
233    #       self._AddSuffix(basename, suffix))
234    return suffix
235
236  def _CreateLogFileHandle(self, name):
237    print('MockLogger: creating open file handle for %s (writing)' % name)
238
239  def _CreateLogFileHandles(self, basename):
240    suffix = self._FindSuffix(basename)
241    suffixed_basename = self._AddSuffix(basename, suffix)
242
243    print('MockLogger: opening file %s.cmd' % suffixed_basename)
244    print('MockLogger: opening file %s.out' % suffixed_basename)
245    print('MockLogger: opening file %s.err' % suffixed_basename)
246
247    self._CreateLogFileSymlinks(basename, suffixed_basename)
248
249  # Symlink unsuffixed basename to currently suffixed one.
250  def _CreateLogFileSymlinks(self, basename, suffixed_basename):
251    for extension in ['cmd', 'out', 'err']:
252      src_file = '%s.%s' % (os.path.basename(suffixed_basename), extension)
253      dest_file = '%s.%s' % (basename, extension)
254      print('MockLogger: Calling os.symlink(%s, %s)' % (src_file, dest_file))
255
256  def _WriteTo(self, _fd, msg, _flush):
257    print('MockLogger: %s' % msg)
258
259  def LogStartDots(self, _print_to_console=True):
260    print('. ')
261
262  def LogAppendDot(self, _print_to_console=True):
263    print('. ')
264
265  def LogEndDots(self, _print_to_console=True):
266    print('\n')
267
268  def LogMsg(self, _file_fd, _term_fd, msg, **_kwargs):
269    print('MockLogger: %s' % msg)
270
271  def _GetStdout(self, _print_to_console):
272    return None
273
274  def _GetStderr(self, _print_to_console):
275    return None
276
277  def LogCmdToFileOnly(self, *_args, **_kwargs):
278    return
279
280  # def LogCmdToFileOnly(self, cmd, machine='', user=None):
281  #   host = ('%s@%s' % (user, machine)) if user else machine
282  #   cmd_string = 'CMD (%s): %s\n' % (host, cmd)
283  #   print('MockLogger: Writing to file ONLY: %s' % cmd_string)
284
285  def LogCmd(self, cmd, machine='', user=None, print_to_console=True):
286    if user:
287      host = '%s@%s' % (user, machine)
288    else:
289      host = machine
290
291    self.LogMsg(0, self._GetStdout(print_to_console),
292                'CMD (%s): %s\n' % (host, cmd))
293
294  def LogFatal(self, msg, print_to_console=True):
295    self.LogMsg(0, self._GetStderr(print_to_console), 'FATAL: %s\n' % msg)
296    self.LogMsg(0, self._GetStderr(print_to_console),
297                '\n'.join(traceback.format_stack()))
298    print('MockLogger: Calling sysexit(1)')
299
300  def LogError(self, msg, print_to_console=True):
301    self.LogMsg(0, self._GetStderr(print_to_console), 'ERROR: %s\n' % msg)
302
303  def LogWarning(self, msg, print_to_console=True):
304    self.LogMsg(0, self._GetStderr(print_to_console), 'WARNING: %s\n' % msg)
305
306  def LogOutput(self, msg, print_to_console=True):
307    self.LogMsg(0, self._GetStdout(print_to_console), 'OUTPUT: %s\n' % msg)
308
309  def LogFatalIf(self, condition, msg):
310    if condition:
311      self.LogFatal(msg)
312
313  def LogErrorIf(self, condition, msg):
314    if condition:
315      self.LogError(msg)
316
317  def LogWarningIf(self, condition, msg):
318    if condition:
319      self.LogWarning(msg)
320
321  def LogCommandOutput(self, msg, print_to_console=True):
322    self.LogMsg(self.stdout,
323                self._GetStdout(print_to_console),
324                msg,
325                flush=False)
326
327  def LogCommandError(self, msg, print_to_console=True):
328    self.LogMsg(self.stderr,
329                self._GetStderr(print_to_console),
330                msg,
331                flush=False)
332
333  def Flush(self):
334    print('MockLogger: Flushing cmdfd, stdout, stderr')
335
336
337main_logger = None
338
339
340def InitLogger(script_name, log_dir, print_console=True, mock=False):
341  """Initialize a global logger. To be called only once."""
342  # pylint: disable=global-statement
343  global main_logger
344  assert not main_logger, 'The logger has already been initialized'
345  rootdir, basefilename = GetRoot(script_name)
346  if not log_dir:
347    log_dir = rootdir
348  if not mock:
349    main_logger = Logger(log_dir, basefilename, print_console)
350  else:
351    main_logger = MockLogger(log_dir, basefilename, print_console)
352
353
354def GetLogger(log_dir='', mock=False):
355  if not main_logger:
356    InitLogger(sys.argv[0], log_dir, mock=mock)
357  return main_logger
358
359
360def HandleUncaughtExceptions(fun):
361  """Catches all exceptions that would go outside decorated fun scope."""
362
363  def _Interceptor(*args, **kwargs):
364    try:
365      return fun(*args, **kwargs)
366    except StandardError:
367      GetLogger().LogFatal('Uncaught exception:\n%s' % traceback.format_exc())
368
369  return _Interceptor
370