New argument --set-kernel in Jupytext command line (#230)

HISTORY.rst

@@ -7,6 +7,7 @@ Release History
 ++++++++++++++++++++++
 
 **Improvements**
+- New argument ``--set-kernel`` in Jupytext command line (#230)
 
 **BugFixes**
 - Invalid notebooks may cause a warning, but not a fatal error (#234)

jupytext/cli.py

@@ -8,12 +8,14 @@
 import argparse
 import json
 from copy import copy
+from jupyter_client.kernelspec import find_kernel_specs, get_kernel_spec
 from .jupytext import readf, reads, writef, writes
 from .formats import _VALID_FORMAT_OPTIONS, _BINARY_FORMAT_OPTIONS, check_file_version
 from .formats import long_form_one_format, long_form_multiple_formats, short_form_one_format
 from .paired_paths import paired_paths, base_path, full_path, InconsistentPath
 from .combine import combine_inputs_with_outputs
 from .compare import test_round_trip_conversion, NotebookDifference
+from .kernels import kernelspec_from_language
 from .version import __version__
 
 
@@ -69,6 +71,12 @@ def parse_jupytext_args(args=None):
                         type=str,
                         help='Set jupytext.formats metadata to the given value. Use this to activate pairing on a '
                              'notebook, with e.g. --set-formats ipynb,py:light')
+    parser.add_argument('--kernel', '-k',
+                        type=str,
+                        help="Set the kernel with the given name on the notebook. Use '--kernel -' to set "
+                             "a kernel matching the current environment on Python notebooks, and matching the "
+                             "notebook language otherwise "
+                             "(get the list of available kernels with 'jupyter kernelspec list')")
     parser.add_argument('--update-metadata',
                         default={},
                         type=json.loads,
@@ -86,7 +94,7 @@ def parse_jupytext_args(args=None):
 
     # Action: convert(default)/version/list paired paths/sync/apply/test
     action = parser.add_mutually_exclusive_group()
-    action.add_argument('--version',
+    action.add_argument('--version', '-v',
                         action='store_true',
                         help="Show jupytext's version number and exit")
     action.add_argument('--paired-paths', '-p',
@@ -258,6 +266,28 @@ def writef_git_add(notebook_, nb_file_, fmt_):
             set_prefix_and_suffix(fmt, notebook, nb_file)
             notebook, inputs_nb_file, outputs_nb_file = load_paired_notebook(notebook, fmt, nb_file, log)
 
+        # Set the kernel
+        if args.kernel == '-':
+            language = notebook.metadata.get('jupytext', {})['main_language'] \
+                       or notebook.metadata['kernelspec']['language']
+            if not language:
+                raise ValueError('Cannot infer a kernel as notebook language is not defined')
+            kernelspec = kernelspec_from_language(language)
+            if not kernelspec:
+                raise ValueError('Found no kernel for {}'.format(language))
+            notebook.metadata['kernelspec'] = kernelspec
+            if 'main_language' in notebook.metadata.get('jupytext', {}):
+                notebook.metadata['jupytext'].pop('main_language')
+        elif args.kernel:
+            try:
+                kernelspec = get_kernel_spec(args.kernel)
+            except KeyError:
+                raise KeyError('Please choose a kernel name among {}'
+                               .format([name for name in find_kernel_specs()]))
+            notebook.metadata['kernelspec'] = {'name': args.kernel,
+                                               'language': kernelspec.language,
+                                               'display_name': kernelspec.display_name}
+
         # II. ### Apply commands onto the notebook ###
         # Pipe the notebook into the desired commands
         for cmd in args.pipe or []:
@@ -268,7 +298,7 @@ def writef_git_add(notebook_, nb_file_, fmt_):
             pipe_notebook(notebook, cmd, args.pipe_fmt, update=False)
 
         # III. ### Possible actions ###
-        modified = args.update_metadata or args.pipe
+        modified = args.update_metadata or args.pipe or args.kernel
         # a. Test round trip conversion
         if args.test or args.test_strict:
             try:

jupytext/contentsmanager.py

@@ -20,26 +20,13 @@
     # Older versions of notebook do not have the LargeFileManager #217
     from notebook.services.contents.filemanager import FileContentsManager as LargeFileManager
 
-from jupyter_client.kernelspec import find_kernel_specs, get_kernel_spec
-
 from .jupytext import reads, writes, create_prefix_dir
 from .combine import combine_inputs_with_outputs
 from .formats import rearrange_jupytext_metadata, check_file_version
 from .formats import NOTEBOOK_EXTENSIONS, long_form_one_format, long_form_multiple_formats
 from .formats import short_form_one_format, short_form_multiple_formats
 from .paired_paths import paired_paths, find_base_path_and_format, base_path, full_path, InconsistentPath
-
-
-def kernelspec_from_language(language):
-    """Return the kernel specification for the first kernel with a matching language"""
-    try:
-        for name in find_kernel_specs():
-            kernel_specs = get_kernel_spec(name)
-            if kernel_specs.language == language or (language == 'c++' and kernel_specs.language.startswith('C++')):
-                return {'name': name, 'language': language, 'display_name': kernel_specs.display_name}
-    except (KeyError, ValueError):
-        pass
-    return None
+from .kernels import set_kernelspec_from_language
 
 
 def preferred_format(incomplete_format, preferred_formats):
@@ -409,12 +396,7 @@ def get(self, path, content=True, type=None, format=None, load_alternative_forma
         if model_outputs:
             combine_inputs_with_outputs(model['content'], model_outputs['content'], fmt_inputs)
         elif not path.endswith('.ipynb'):
-            nbk = model['content']
-            language = nbk.metadata.get('jupytext', {}).get('main_language', 'python')
-            if 'kernelspec' not in nbk.metadata and language != 'python':
-                kernelspec = kernelspec_from_language(language)
-                if kernelspec:
-                    nbk.metadata['kernelspec'] = kernelspec
+            set_kernelspec_from_language(model['content'])
 
         # Trust code cells when they have no output
         for cell in model['content'].cells:

jupytext/kernels.py

@@ -0,0 +1,33 @@
+"""Find kernel specifications for a given language"""
+
+import sys
+from jupyter_client.kernelspec import find_kernel_specs, get_kernel_spec
+
+
+def set_kernelspec_from_language(notebook):
+    """Set the kernel specification based on the 'main_language' metadata"""
+    language = notebook.metadata.get('jupytext', {}).get('main_language')
+    if 'kernelspec' not in notebook.metadata and language:
+        kernelspec = kernelspec_from_language(language)
+        if kernelspec:
+            notebook.metadata['kernelspec'] = kernelspec
+            notebook.metadata.get('jupytext', {}).pop('main_language')
+
+
+def kernelspec_from_language(language):
+    """Return the python kernel that matches the current env, or the first kernel that matches the given language"""
+    try:
+        if language == 'python':
+            # Return the kernel that matches the current Python executable
+            for name in find_kernel_specs():
+                kernel_specs = get_kernel_spec(name)
+                if kernel_specs.argv[0] == sys.executable:
+                    return {'name': name, 'language': language, 'display_name': kernel_specs.display_name}
+
+        for name in find_kernel_specs():
+            kernel_specs = get_kernel_spec(name)
+            if kernel_specs.language == language or (language == 'c++' and kernel_specs.language.startswith('C++')):
+                return {'name': name, 'language': language, 'display_name': kernel_specs.display_name}
+    except (KeyError, ValueError):
+        pass
+    return None

tests/test_cli.py

@@ -1,4 +1,5 @@
 import os
+import sys
 import time
 import stat
 import mock
@@ -9,6 +10,7 @@
 from testfixtures import compare
 from argparse import ArgumentTypeError
 from nbformat.v4.nbbase import new_notebook, new_markdown_cell, new_code_cell
+from jupyter_client.kernelspec import find_kernel_specs, get_kernel_spec
 from jupytext import __version__
 from jupytext import readf, writef, writes
 from jupytext.cli import parse_jupytext_args, jupytext, jupytext_cli, system, str2bool
@@ -556,6 +558,37 @@ def test_update_metadata(py_file, tmpdir, capsys):
     assert 'invalid' in err
 
 
+@pytest.mark.parametrize('py_file', list_notebooks('python'))
+def test_set_kernel_auto(py_file, tmpdir):
+    tmp_py = str(tmpdir.join('notebook.py'))
+    tmp_ipynb = str(tmpdir.join('notebook.ipynb'))
+
+    copyfile(py_file, tmp_py)
+
+    jupytext(['--to', 'ipynb', tmp_py, '--kernel', '-'])
+
+    nb = readf(tmp_ipynb)
+    kernel_name = nb.metadata['kernelspec']['name']
+    assert get_kernel_spec(kernel_name).argv[0] == sys.executable
+
+
+@pytest.mark.parametrize('py_file', list_notebooks('python'))
+def test_set_kernel_with_name(py_file, tmpdir):
+    tmp_py = str(tmpdir.join('notebook.py'))
+    tmp_ipynb = str(tmpdir.join('notebook.ipynb'))
+
+    copyfile(py_file, tmp_py)
+
+    for kernel in find_kernel_specs():
+        jupytext(['--to', 'ipynb', tmp_py, '--kernel', kernel])
+
+        nb = readf(tmp_ipynb)
+        assert nb.metadata['kernelspec']['name'] == kernel
+
+    with pytest.raises(KeyError) as info:
+        jupytext(['--to', 'ipynb', tmp_py, '--kernel', 'non_existing_env'])
+
+
 @pytest.mark.parametrize('nb_file', list_notebooks('ipynb_py'))
 def test_paired_paths(nb_file, tmpdir, capsys):
     tmp_ipynb = str(tmpdir.join('notebook.ipynb'))

tests/test_contentsmanager.py

@@ -14,7 +14,7 @@
 from jupytext.compare import compare_notebooks
 from jupytext.header import header_to_metadata_and_cell
 from jupytext.formats import read_format_from_metadata, auto_ext_from_metadata
-from jupytext.contentsmanager import kernelspec_from_language
+from jupytext.kernels import kernelspec_from_language
 from .utils import list_notebooks, requires_sphinx_gallery, requires_pandoc, skip_if_dict_is_not_ordered
 
 
@@ -528,6 +528,33 @@ def test_open_using_preferred_and_default_format_174(nb_file, tmpdir):
     assert not os.path.isfile(str(tmpdir.join('other/notebook.ipynb')))
 
 
+@skip_if_dict_is_not_ordered
+@pytest.mark.parametrize('nb_file', list_notebooks('ipynb_py', skip='many hash'))
+def test_kernelspec_are_preserved(nb_file, tmpdir):
+    tmp_ipynb = str(tmpdir.join('notebook.ipynb'))
+    tmp_py = str(tmpdir.join('notebook.py'))
+    shutil.copyfile(nb_file, tmp_ipynb)
+
+    cm = jupytext.TextFileContentsManager()
+    cm.root_dir = str(tmpdir)
+    cm.default_jupytext_formats = "ipynb,py"
+    cm.default_notebook_metadata_filter = "-all"
+
+    # load notebook
+    model = cm.get('notebook.ipynb')
+    model['content'].metadata['kernelspec'] = {'display_name': 'Kernel name',
+                                               'language': 'python',
+                                               'name': 'custom'}
+
+    # save to ipynb and py
+    cm.save(model=model, path='notebook.ipynb')
+    assert os.path.isfile(tmp_py)
+
+    # read ipynb
+    model2 = cm.get('notebook.ipynb')
+    compare_notebooks(model['content'], model2['content'])
+
+
 @skip_if_dict_is_not_ordered
 @pytest.mark.parametrize('nb_file', list_notebooks('ipynb_py'))
 def test_save_to_light_percent_sphinx_format(nb_file, tmpdir):
@@ -816,7 +843,7 @@ def test_metadata_filter_is_effective(nb_file, tmpdir):
     nb2 = jupytext.readf(str(tmpdir.join(tmp_script)))
 
     # test no metadata
-    assert set(nb2.metadata.keys()) <= {'jupytext'}
+    assert set(nb2.metadata.keys()) <= {'jupytext', 'kernelspec'}
     for cell in nb2.cells:
         assert not cell.metadata