import logging import sys import threading from time import time from ..binding.frameworks import _patched_call # noqa from ..config import running_remotely, config class StdStreamPatch(object): _stdout_proxy = None _stderr_proxy = None _stdout_original_write = None _stderr_original_write = None @staticmethod def patch_std_streams(a_logger, connect_stdout=True, connect_stderr=True): if (connect_stdout or connect_stderr) and not PrintPatchLogger.patched and not running_remotely(): StdStreamPatch._stdout_proxy = PrintPatchLogger(sys.stdout, a_logger, level=logging.INFO) \ if connect_stdout else None StdStreamPatch._stderr_proxy = PrintPatchLogger(sys.stderr, a_logger, level=logging.ERROR) \ if connect_stderr else None if StdStreamPatch._stdout_proxy: # noinspection PyBroadException try: if StdStreamPatch._stdout_original_write is None: StdStreamPatch._stdout_original_write = sys.stdout.write # this will only work in python 3, guard it with try/catch if not hasattr(sys.stdout, '_original_write'): sys.stdout._original_write = sys.stdout.write sys.stdout.write = StdStreamPatch._stdout__patched__write__ except Exception: pass sys.stdout = StdStreamPatch._stdout_proxy # noinspection PyBroadException try: sys.__stdout__ = sys.stdout except Exception: pass if StdStreamPatch._stderr_proxy: # noinspection PyBroadException try: if StdStreamPatch._stderr_original_write is None: StdStreamPatch._stderr_original_write = sys.stderr.write if not hasattr(sys.stderr, '_original_write'): sys.stderr._original_write = sys.stderr.write sys.stderr.write = StdStreamPatch._stderr__patched__write__ except Exception: pass sys.stderr = StdStreamPatch._stderr_proxy # patch the base streams of sys (this way colorama will keep its ANSI colors) # noinspection PyBroadException try: sys.__stderr__ = sys.stderr except Exception: pass # now check if we have loguru and make it re-register the handlers # because it stores internally the stream.write function, which we cant patch # noinspection PyBroadException try: from loguru import logger # noqa register_stderr = None register_stdout = None for k, v in logger._handlers.items(): # noqa if connect_stderr and v._name == '': # noqa register_stderr = k elif connect_stdout and v._name == '': # noqa register_stderr = k if register_stderr is not None: logger.remove(register_stderr) logger.add(sys.stderr) if register_stdout is not None: logger.remove(register_stdout) logger.add(sys.stdout) except Exception: pass elif (connect_stdout or connect_stderr) and not running_remotely(): if StdStreamPatch._stdout_proxy and connect_stdout: StdStreamPatch._stdout_proxy.connect(a_logger) if StdStreamPatch._stderr_proxy and connect_stderr: StdStreamPatch._stderr_proxy.connect(a_logger) @staticmethod def patch_logging_formatter(a_logger, logging_handler=None): if not logging_handler: import logging logging_handler = logging.Handler logging_handler.format = _patched_call(logging_handler.format, HandlerFormat(a_logger)) @staticmethod def remove_patch_logging_formatter(logging_handler=None): if not logging_handler: import logging logging_handler = logging.Handler # remove the function, Hack calling patched logging.Handler.format() returns the original function # noinspection PyBroadException try: logging_handler.format = logging_handler.format() # noqa except Exception: pass @staticmethod def remove_std_logger(logger=None): if isinstance(sys.stdout, PrintPatchLogger): # noinspection PyBroadException try: sys.stdout.disconnect(logger) except Exception: pass if isinstance(sys.stderr, PrintPatchLogger): # noinspection PyBroadException try: sys.stderr.disconnect(logger) except Exception: pass @staticmethod def stdout_original_write(*args, **kwargs): if StdStreamPatch._stdout_original_write: StdStreamPatch._stdout_original_write(*args, **kwargs) else: sys.stdout.write(*args, **kwargs) @staticmethod def stderr_original_write(*args, **kwargs): if StdStreamPatch._stderr_original_write: StdStreamPatch._stderr_original_write(*args, **kwargs) else: sys.stderr.write(*args, **kwargs) @staticmethod def _stdout__patched__write__(*args, **kwargs): if StdStreamPatch._stdout_proxy: return StdStreamPatch._stdout_proxy.write(*args, **kwargs) return sys.stdout._original_write(*args, **kwargs) # noqa @staticmethod def _stderr__patched__write__(*args, **kwargs): if StdStreamPatch._stderr_proxy: return StdStreamPatch._stderr_proxy.write(*args, **kwargs) return sys.stderr._original_write(*args, **kwargs) # noqa class HandlerFormat(object): def __init__(self, logger): self._logger = logger def __call__(self, original_format_func, *args): # hack get back original function, so we can remove it if all(a is None for a in args): return original_format_func if len(args) == 1: record = args[0] msg = original_format_func(record) else: handler = args[0] record = args[1] msg = original_format_func(handler, record) self._logger.report_text(msg=msg, level=record.levelno, print_console=False) return msg class PrintPatchLogger(object): """ Allowed patching a stream into the logger. Used for capturing and logging stdin and stderr when running in development mode pseudo worker. """ patched = False lock = threading.Lock() recursion_protect_lock = threading.RLock() cr_flush_period = config.get("development.worker.console_cr_flush_period", 0) def __init__(self, stream, logger=None, level=logging.INFO): PrintPatchLogger.patched = True self._terminal = stream self._log = logger self._log_level = level self._cur_line = '' self._force_lf_flush = False def write(self, message): # make sure that we do not end up in infinite loop (i.e. log.console ends up calling us) if self._log and not PrintPatchLogger.recursion_protect_lock._is_owned(): # noqa try: self.lock.acquire() with PrintPatchLogger.recursion_protect_lock: if hasattr(self._terminal, '_original_write'): self._terminal._original_write(message) # noqa else: self._terminal.write(message) do_flush = '\n' in message do_cr = '\r' in message self._cur_line += message if not do_flush and do_cr and PrintPatchLogger.cr_flush_period and self._force_lf_flush: self._cur_line += '\n' do_flush = True self._force_lf_flush = False if (not do_flush and (PrintPatchLogger.cr_flush_period or not do_cr)) or not message: return if PrintPatchLogger.cr_flush_period and self._cur_line: self._cur_line = '\n'.join(line.split('\r')[-1] for line in self._cur_line.split('\n')) last_lf = self._cur_line.rindex('\n' if do_flush else '\r') next_line = self._cur_line[last_lf + 1:] cur_line = self._cur_line[:last_lf + 1].rstrip() self._cur_line = next_line finally: self.lock.release() if cur_line: with PrintPatchLogger.recursion_protect_lock: # noinspection PyBroadException try: if self._log: # noinspection PyProtectedMember self._log._console(cur_line, level=self._log_level, omit_console=True) except Exception: # what can we do, nothing pass else: if hasattr(self._terminal, '_original_write'): self._terminal._original_write(message) # noqa else: self._terminal.write(message) def connect(self, logger): self._cur_line = '' self._log = logger def disconnect(self, logger=None): # disconnect the logger only if it was registered if not logger or self._log == logger: self.connect(None) def force_lf_flush(self): self._force_lf_flush = True def __getattr__(self, attr): if attr in ['_log', '_terminal', '_log_level', '_cur_line', '_cr_overwrite', '_force_lf_flush']: return self.__dict__.get(attr) return getattr(self._terminal, attr) def __setattr__(self, key, value): if key in ['_log', '_terminal', '_log_level', '_cur_line', '_cr_overwrite', '_force_lf_flush']: self.__dict__[key] = value else: return setattr(self._terminal, key, value) class LogFlusher(threading.Thread): def __init__(self, logger, period, **kwargs): super(LogFlusher, self).__init__(**kwargs) self.daemon = True self._period = period self._logger = logger self._exit_event = threading.Event() self._cr_last_flush = 0 try: self._cr_flush_period = float(PrintPatchLogger.cr_flush_period) except (ValueError, TypeError): self._cr_flush_period = 0 @property def period(self): return self._period def run(self): self._logger.flush() # store original wait period while True: period = self._period while not self._exit_event.wait(period or 1.0): if self._cr_flush_period and time() - self._cr_last_flush > self._cr_flush_period: if isinstance(sys.stdout, PrintPatchLogger): sys.stdout.force_lf_flush() if isinstance(sys.stderr, PrintPatchLogger): sys.stderr.force_lf_flush() self._cr_last_flush = time() # now signal the real flush self._logger.flush() # check if period is negative or None we should exit if self._period is None or self._period < 0: break # check if period was changed, we should restart self._exit_event.clear() def exit(self): self._period = None self._exit_event.set() def set_period(self, period): self._period = period # make sure we exit the previous wait self._exit_event.set()