diff --git a/loguru/_logger.py b/loguru/_logger.py index 2ef78c41f..fa04ef5d1 100644 --- a/loguru/_logger.py +++ b/loguru/_logger.py @@ -1752,6 +1752,40 @@ def configure(self, *, handlers=None, levels=None, extra=None, patcher=None, act return [self.add(**params) for params in handlers] + def reinstall(self): + """Reinstall the core of logger. + + When using multiprocessing, you can pass logger as a parameter to the target of + ``multiprocessing.Process``, and run this method once, thus you don't need to pass + logger to every function you called in the same process with spawn multiprocessing. + + Examples + -------- + >>> def subworker(logger): + ... logger.reinstall() + ... logger.info("Child") + ... deeper_subworker() + + >>> def deeper_subworker(): + ... logger.info("Grandchild") + + >>> def test_process_spawn(): + ... spawn_context = multiprocessing.get_context("spawn") + ... logger.add("file.log", context=spawn_context, enqueue=True, catch=False) + ... + ... process = spawn_context.Process(target=subworker, args=(logger,)) + ... process.start() + ... process.join() + + ... assert process.exitcode == 0 + + ... logger.info("Main") + ... logger.remove() + """ + from loguru import logger + + logger._core = self._core + def _change_activation(self, name, status): if not (name is None or isinstance(name, str)): raise TypeError( diff --git a/tests/test_reinstall.py b/tests/test_reinstall.py new file mode 100644 index 000000000..188879e90 --- /dev/null +++ b/tests/test_reinstall.py @@ -0,0 +1,72 @@ +import multiprocessing +import os + +import pytest + +from loguru import logger + + +@pytest.fixture +def fork_context(): + yield multiprocessing.get_context("fork") + + +@pytest.fixture +def spawn_context(): + yield multiprocessing.get_context("spawn") + + +class Writer: + def __init__(self): + self._output = "" + + def write(self, message): + self._output += message + + def read(self): + return self._output + + +def subworker(logger): + logger.reinstall() + logger.info("Child") + deeper_subworker() + + +def deeper_subworker(): + logger.info("Grandchild") + + +@pytest.mark.skipif(os.name == "nt", reason="Windows does not support forking") +def test_process_fork(fork_context): + writer = Writer() + + logger.add(writer, context=fork_context, format="{message}", enqueue=True, catch=False) + + process = fork_context.Process(target=subworker, args=(logger,)) + process.start() + process.join() + + assert process.exitcode == 0 + + logger.info("Main") + logger.remove() + + assert writer.read() == "Child\nGrandchild\nMain\n" + + +def test_process_spawn(spawn_context): + writer = Writer() + + logger.add(writer, context=spawn_context, format="{message}", enqueue=True, catch=False) + + process = spawn_context.Process(target=subworker, args=(logger,)) + process.start() + process.join() + + assert process.exitcode == 0 + + logger.info("Main") + logger.remove() + + assert writer.read() == "Child\nGrandchild\nMain\n"