Merge pull request #9 from analysiscenter/minor_fix

Minor fixes

Author: SergeyTsimfer

nbtools/__init__.py

@@ -4,4 +4,4 @@
 from .run_notebook import run_notebook
 from .pylint_notebook import pylint_notebook
 
-__version__ = '0.9.11'
+__version__ = '0.9.12'

nbtools/pylint_notebook.py

@@ -1,6 +1,8 @@
 """ Functions for code quality control of Jupyter Notebooks. """
 #pylint: disable=import-outside-toplevel
 import os
+from io import StringIO
+from contextlib import redirect_stderr, redirect_stdout
 
 from .core import StringWithDisabledRepr, get_notebook_path, notebook_to_script
 
@@ -101,7 +103,13 @@ def pylint_notebook(path=None, options=(), disable=(), enable=(), printer=print,
     pylint_params : dict
         Additional parameter of linting. Each is converted to a separate valid entry in the `pylintrc` file.
     """
-    from pylint import epylint as lint
+    try:
+        from pylint.lint import Run
+        from pylint.reporters.text import TextReporter
+    except ImportError as exception:
+        raise ImportError('Install pylint') from exception
+
+
     path = path or get_notebook_path()
     if path is None:
         raise ValueError('Provide path to Jupyter Notebook or run `pylint_notebook` inside of it!')
@@ -117,25 +125,32 @@ def pylint_notebook(path=None, options=(), disable=(), enable=(), printer=print,
     pylintrc = generate_pylintrc(path_pylintrc, disable=disable, enable=enable, **pylint_params)
 
     # Run pylint on script with pylintrc configuration
-    pylint_cmdline = ' '.join([path_pylintrc, f'--rcfile {path_pylintrc}', *options])
-    pylint_stdout, pylint_stderr = lint.py_run(pylint_cmdline, return_std=True)
+    pylint_cmdline = [path_pylintrc, f'--rcfile={path_pylintrc}', *options]
+
+    # Run pylint and catch messagies
+    with redirect_stdout(StringIO()) as pylint_stdout, redirect_stderr(StringIO()) as pylint_stderr:
+        reporter = TextReporter(pylint_stdout)
+        Run(pylint_cmdline, reporter=reporter, exit=False)
 
-    errors = pylint_stderr.getvalue()
-    report = pylint_stdout.getvalue()
+        report = pylint_stdout.getvalue()
+        errors = pylint_stderr.getvalue()
 
     # Prepare custom report
     output = []
 
     for line in report.split('\n'):
+        # error line is a str in the format:
+        # "filename.py:line_num:position_num: error_code: error_text (error_name)"
         if 'rated' in line:
             output.insert(0, line.strip(' '))
             output.insert(1, '–' * (len(line) - 1))
 
         elif path_script in line or script_name in line:
             line = line.replace(path_script, '__').replace(script_name, '__')
+            line_semicolon_split = line.split(':')
 
             # Locate the cell and line inside the cell
-            code_line_number = int(line.split(':')[1])
+            code_line_number = int(line_semicolon_split[1])
             for cell_number, cell_ranges in cell_line_numbers.items():
                 if code_line_number in cell_ranges:
                     cell_line_number = code_line_number - cell_ranges[0]
@@ -144,12 +159,13 @@ def pylint_notebook(path=None, options=(), disable=(), enable=(), printer=print,
                 cell_number, cell_line_number = -1, code_line_number
 
             # Find error_code and error_name: for example, `C0123` and `invalid-name`
-            position_left  = line.find('(') + 1
-            position_right = line.find(')') - 1
-            error_message = line[position_left : position_right]
-            error_human_message = line[position_right + 2:]
-            error_code, error_name, *_ = error_message.split(',')
-            error_name = error_name.strip()
+            error_code = line_semicolon_split[3].strip()
+
+            error_message = line_semicolon_split[4].strip()
+            error_name_start  = error_message.rfind('(')
+
+            error_human_message = error_message[:error_name_start].strip()
+            error_name = error_message[error_name_start+1:-1]
 
             # Make new message
             message = f'Cell {cell_number}:{cell_line_number}, code={error_code}, name={error_name}'

nbtools/run_notebook.py

@@ -19,24 +19,29 @@ def run_in_process(func):
     """ Decorator to run the `func` in a separated process for terminating all relevant processes properly. """
     @wraps(func)
     def _wrapper(*args, **kwargs):
-        # pylint: disable=bare-except
-        returned_value = Queue()
-        kwargs = {**kwargs, 'returned_value': returned_value}
+        # pylint: disable=broad-exception-caught, broad-exception-raised
+        _output_queue = Queue()
+        kwargs = {**kwargs, '_output_queue': _output_queue}
 
         json_path = None
 
+        output = {'failed': True, 'traceback': ''}
+
         try:
             process = Process(target=func, args=args, kwargs=kwargs)
             process.start()
 
             path = args[0] if args else kwargs['path']
             json_path = f'{TMP_DIR}/{process.pid}.json'
+
             with open(json_path, 'w', encoding='utf-8') as file:
                 json.dump({'path': path}, file)
 
-            output = returned_value.get()
+            output = _output_queue.get()
             process.join()
-        except:
+        except Exception as e:
+            output = {'failed': True, 'traceback': e}
+
             # Terminate all relevant processes when something went wrong, e.g. Keyboard Interrupt
             for child in psutil.Process(process.pid).children():
                 if psutil.pid_exists(child.pid):
@@ -48,6 +53,9 @@ def _wrapper(*args, **kwargs):
             if json_path is not None and os.path.exists(json_path):
                 os.remove(json_path)
 
+            if kwargs.get('raise_exception', False) and output['failed']:
+                raise Exception(output['traceback'])
+
         return output
     return _wrapper
 
@@ -130,7 +138,7 @@ def run_notebook(path, inputs=None, outputs=None, inputs_pos=1, replace_inputs_p
                  working_dir = './', execute_kwargs=None,
                  out_path_db=None, out_path_ipynb=None, out_path_html=None, remove_db='always', add_timestamp=True,
                  hide_code_cells=False, mask_extra_code=False, display_links=True,
-                 raise_exception=False, return_notebook=False, returned_value=None):
+                 raise_exception=False, return_notebook=False, _output_queue=None):
     """ Execute a Jupyter Notebook programmatically.
     Heavily inspired by https://github.com/tritemio/nbrun.
 
@@ -190,6 +198,9 @@ def run_notebook(path, inputs=None, outputs=None, inputs_pos=1, replace_inputs_p
         Note, that database exists only if inputs and/or outputs are provided.
     execute_kwargs : dict, optional
         Parameters of `:class:ExecutePreprocessor`.
+        For example, you can provide timeout, kernel_name, resources (such as metadata)
+        and other `nbclient.client.NotebookClient` arguments from :ref:`the NotebookClient doc page
+        <https://nbclient.readthedocs.io/en/latest/reference/nbclient.html#nbclient.client.NotebookClient>`.
     add_timestamp : bool, optional
         Whether to add a cell with execution information at the beginning of the saved notebook.
     hide_code_cells : bool, optional
@@ -203,7 +214,7 @@ def run_notebook(path, inputs=None, outputs=None, inputs_pos=1, replace_inputs_p
         Whether to re-raise exceptions from the notebook.
     return_notebook : bool, optional
         Whether to return the notebook object from this function.
-    returned_value : None
+    _output_queue : None
         Placeholder for the :func:`~.run_in_process` decorator to return this function result.
 
     Returns
@@ -224,7 +235,7 @@ def run_notebook(path, inputs=None, outputs=None, inputs_pos=1, replace_inputs_p
            Executed notebook object.
            Note that this output is provided only if `return_notebook` is True.
     """
-    # pylint: disable=bare-except, lost-exception
+    # pylint: disable=bare-except, lost-exception, return-in-finally
     import nbformat
     from jupyter_client.manager import KernelManager
     from nbconvert.preprocessors import ExecutePreprocessor
@@ -310,9 +321,6 @@ def run_notebook(path, inputs=None, outputs=None, inputs_pos=1, replace_inputs_p
         if outputs is not None:
             executor.kc = kernel_manager.client() # For compatibility with 5.x.x version of `nbconvert`
             executor.preprocess_cell(output_cell, {'metadata': {'path': working_dir}}, -1)
-
-        if raise_exception:
-            raise
     finally:
         # Shutdown kernel
         kernel_manager.cleanup_resources()
@@ -340,6 +348,11 @@ def run_notebook(path, inputs=None, outputs=None, inputs_pos=1, replace_inputs_p
         # Prepare execution results: execution state, notebook outputs and error info (if exists)
         if failed:
             exec_res = {'failed': failed, 'failed cell number': error_cell_num, 'traceback': traceback_message}
+
+            # Re-raise exception if needed
+            if raise_exception:
+                _output_queue.put(exec_res)
+                return None
         else:
             exec_res = {'failed': failed, 'failed cell number': None, 'traceback': ''}
 
@@ -372,8 +385,8 @@ def run_notebook(path, inputs=None, outputs=None, inputs_pos=1, replace_inputs_p
         if return_notebook:
             exec_res['notebook'] = notebook
 
-        returned_value.put(exec_res) # return for parent process
-        return
+        _output_queue.put(exec_res) # return for parent process
+    return None
 
 # Mask functions for database operations cells
 def mask_inputs_reading(notebook, pos):