Fix log output

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)

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,

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)

papermill/tests/test_clientwrap.py

@@ -5,30 +5,30 @@
 
 from ..clientwrap import PapermillNotebookClient
 from ..engines import NotebookExecutionManager
-from ..log import logger
+from ..log import logger, notebook_logger
 from . import get_notebook_path
 
 
 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(

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):