import logging import sys import threading from ..backend_interface.task.development.worker import DevWorker from ..backend_interface.task.log import TaskHandler from ..config import running_remotely class StdStreamPatch(object): _stdout_proxy = None _stderr_proxy = None _stdout_original_write = None _stderr_original_write = None @staticmethod def patch_std_streams(logger): if DevWorker.report_stdout and not PrintPatchLogger.patched and not running_remotely(): StdStreamPatch._stdout_proxy = PrintPatchLogger(sys.stdout, logger, level=logging.INFO) StdStreamPatch._stderr_proxy = PrintPatchLogger(sys.stderr, logger, level=logging.ERROR) logger._task_handler = TaskHandler(task=logger._task, capacity=100) # noinspection PyBroadException try: if StdStreamPatch._stdout_original_write is None: StdStreamPatch._stdout_original_write = sys.stdout.write if StdStreamPatch._stderr_original_write is None: StdStreamPatch._stderr_original_write = sys.stderr.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__ 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.stdout = StdStreamPatch._stdout_proxy 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 # noinspection PyBroadException try: sys.__stdout__ = sys.stdout 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 register_stderr = None register_stdout = None for k, v in logger._handlers.items(): if v._name == '': register_stderr = k elif v._name == '': 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 DevWorker.report_stdout and not running_remotely(): logger._task_handler = TaskHandler(task=logger._task, capacity=100) if StdStreamPatch._stdout_proxy: StdStreamPatch._stdout_proxy.connect(logger) if StdStreamPatch._stderr_proxy: StdStreamPatch._stderr_proxy.connect(logger) @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) @staticmethod def stderr_original_write(*args, **kwargs): if StdStreamPatch._stderr_original_write: StdStreamPatch._stderr_original_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) @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) 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() 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 = '' 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(): try: self.lock.acquire() with PrintPatchLogger.recursion_protect_lock: if hasattr(self._terminal, '_original_write'): self._terminal._original_write(message) else: self._terminal.write(message) do_flush = '\n' in message do_cr = '\r' in message self._cur_line += message if (not do_flush and not do_cr) or not message: return 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: 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) 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 __getattr__(self, attr): if attr in ['_log', '_terminal', '_log_level', '_cur_line']: return self.__dict__.get(attr) return getattr(self._terminal, attr) def __setattr__(self, key, value): if key in ['_log', '_terminal', '_log_level', '_cur_line']: 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() @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): 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()