diff --git a/papermill/clientwrap.py b/papermill/clientwrap.py index 9f00b5ba..e1670d82 100644 --- a/papermill/clientwrap.py +++ b/papermill/clientwrap.py @@ -16,7 +16,7 @@ class PapermillNotebookClient(NotebookClient): stdout_file = Instance(object, default_value=None).tag(config=True) stderr_file = Instance(object, default_value=None).tag(config=True) - def __init__(self, nb_man, km=None, raise_on_iopub_timeout=True, **kw): + def __init__(self, nb_man, km=None, log_cell=None, raise_on_iopub_timeout=True, **kw): """Initializes the execution manager. Parameters @@ -26,9 +26,13 @@ def __init__(self, nb_man, km=None, raise_on_iopub_timeout=True, **kw): km : KernerlManager (optional) Optional kernel manager. If none is provided, a kernel manager will be created. + log_cell : logging.Logger (optional) + Optional logger to use for logging cell output. If not provided, the `log` argument + will be used for cell output. """ super().__init__(nb_man.nb, km=km, raise_on_iopub_timeout=raise_on_iopub_timeout, **kw) self.nb_man = nb_man + self.log_cell = log_cell if log_cell is not None else self.log def execute(self, **kwargs): """ @@ -88,19 +92,20 @@ def log_output_message(self, output): content = "".join(output.text) if output.name == "stdout": if self.log_output: - self.log.info(content) + self.log_cell.info(content) if self.stdout_file: self.stdout_file.write(content) self.stdout_file.flush() elif output.name == "stderr": if self.log_output: # In case users want to redirect stderr differently, pipe to warning - self.log.warning(content) + self.log_cell.warning(content) if self.stderr_file: self.stderr_file.write(content) self.stderr_file.flush() elif self.log_output and ("data" in output and "text/plain" in output.data): - self.log.info("".join(output.data['text/plain'])) + self.log_cell.info("".join(output.data['text/plain'])) + # self.log.info("".join(output.data['text/plain'])) def process_message(self, *arg, **kwargs): output = super().process_message(*arg, **kwargs) diff --git a/papermill/engines.py b/papermill/engines.py index 4d096ba5..a8c0e9e4 100644 --- a/papermill/engines.py +++ b/papermill/engines.py @@ -9,7 +9,7 @@ from .clientwrap import PapermillNotebookClient from .exceptions import PapermillException from .iorw import write_ipynb -from .log import logger +from .log import logger, notebook_logger from .utils import merge_kwargs, nb_kernel_name, nb_language, remove_args @@ -435,6 +435,7 @@ def execute_managed_notebook( startup_timeout=start_timeout, kernel_name=kernel_name, log=logger, + log_cell=notebook_logger, log_output=log_output, stdout_file=stdout_file, stderr_file=stderr_file, diff --git a/papermill/log.py b/papermill/log.py index b90225d2..2508735e 100644 --- a/papermill/log.py +++ b/papermill/log.py @@ -2,3 +2,30 @@ import logging logger = logging.getLogger('papermill') +notebook_logger = logging.getLogger('papermill.notebook') + + +class NbOutputStreamHandler(logging.StreamHandler): + def emit(self, record): + try: + msg = self.format(record) + stream = self.stream + # issue 35046: merged two stream.writes into one. + stream.write(msg) + self.flush() + except RecursionError: # See issue 36272 + raise + except Exception: + self.handleError(record) + + +def _reconfigure_notebook_logger(log: logging.Logger): + log.handlers = [] + custom_handler = NbOutputStreamHandler() + log.addHandler(custom_handler) + formatter = logging.Formatter('%(message)s') + custom_handler.setFormatter(formatter) + log.propagate = False + + +_reconfigure_notebook_logger(notebook_logger) diff --git a/papermill/tests/test_clientwrap.py b/papermill/tests/test_clientwrap.py index cfa2a81a..7202c948 100644 --- a/papermill/tests/test_clientwrap.py +++ b/papermill/tests/test_clientwrap.py @@ -5,7 +5,7 @@ from ..clientwrap import PapermillNotebookClient from ..engines import NotebookExecutionManager -from ..log import logger +from ..log import logger, notebook_logger from . import get_notebook_path @@ -13,22 +13,22 @@ class TestPapermillClientWrapper(unittest.TestCase): def setUp(self): self.nb = nbformat.read(get_notebook_path('test_logging.ipynb'), as_version=4) self.nb_man = NotebookExecutionManager(self.nb) - self.client = PapermillNotebookClient(self.nb_man, log=logger, log_output=True) + self.client = PapermillNotebookClient(self.nb_man, log=logger, log_cell=notebook_logger, log_output=True) def test_logging_stderr_msg(self): - with patch.object(logger, 'warning') as warning_mock: + with patch.object(notebook_logger, 'warning') as warning_mock: for output in self.nb.cells[0].get("outputs", []): self.client.log_output_message(output) warning_mock.assert_called_once_with("INFO:test:test text\n") def test_logging_stdout_msg(self): - with patch.object(logger, 'info') as info_mock: + with patch.object(notebook_logger, 'info') as info_mock: for output in self.nb.cells[1].get("outputs", []): self.client.log_output_message(output) info_mock.assert_called_once_with("hello world\n") def test_logging_data_msg(self): - with patch.object(logger, 'info') as info_mock: + with patch.object(notebook_logger, 'info') as info_mock: for output in self.nb.cells[2].get("outputs", []): self.client.log_output_message(output) info_mock.assert_has_calls( diff --git a/papermill/tests/test_engines.py b/papermill/tests/test_engines.py index db5ee17c..4c3b04d8 100644 --- a/papermill/tests/test_engines.py +++ b/papermill/tests/test_engines.py @@ -9,7 +9,7 @@ from .. import engines, exceptions from ..engines import Engine, NBClientEngine, NotebookExecutionManager from ..iorw import load_notebook_node -from ..log import logger +from ..log import logger, notebook_logger from . import get_notebook_path @@ -434,39 +434,50 @@ def test_nb_convert_engine_execute(self): def test_nb_convert_log_outputs(self): with patch.object(logger, 'info') as info_mock: with patch.object(logger, 'warning') as warning_mock: - with patch.object(NotebookExecutionManager, 'save'): - NBClientEngine.execute_notebook( - self.nb, - 'python', - output_path='foo.ipynb', - progress_bar=False, - log_output=True, - ) - info_mock.assert_has_calls( - [ - call('Executing notebook with kernel: python'), - call('Executing Cell 1---------------------------------------'), - call('Ending Cell 1------------------------------------------'), - call('Executing Cell 2---------------------------------------'), - call('None\n'), - call('Ending Cell 2------------------------------------------'), - ] - ) - warning_mock.is_not_called() + with patch.object(notebook_logger, 'info') as notebook_info_mock: + with patch.object(notebook_logger, 'warning') as notebook_warning_mock: + with patch.object(NotebookExecutionManager, 'save'): + NBClientEngine.execute_notebook( + self.nb, + 'python', + output_path='foo.ipynb', + progress_bar=False, + log_output=True, + ) + info_mock.assert_has_calls( + [ + call('Executing notebook with kernel: python'), + call('Executing Cell 1---------------------------------------'), + call('Ending Cell 1------------------------------------------'), + call('Executing Cell 2---------------------------------------'), + call('Ending Cell 2------------------------------------------'), + ] + ) + notebook_info_mock.assert_has_calls( + [ + call('None\n'), + ] + ) + warning_mock.is_not_called() + notebook_warning_mock.is_not_called() def test_nb_convert_no_log_outputs(self): with patch.object(logger, 'info') as info_mock: with patch.object(logger, 'warning') as warning_mock: - with patch.object(NotebookExecutionManager, 'save'): - NBClientEngine.execute_notebook( - self.nb, - 'python', - output_path='foo.ipynb', - progress_bar=False, - log_output=False, - ) - info_mock.is_not_called() - warning_mock.is_not_called() + with patch.object(notebook_logger, 'info') as notebook_info_mock: + with patch.object(notebook_logger, 'warning') as notebook_warning_mock: + with patch.object(NotebookExecutionManager, 'save'): + NBClientEngine.execute_notebook( + self.nb, + 'python', + output_path='foo.ipynb', + progress_bar=False, + log_output=False, + ) + info_mock.is_not_called() + warning_mock.is_not_called() + notebook_info_mock.is_not_called() + notebook_warning_mock.is_not_called() class TestEngineRegistration(unittest.TestCase):