Support Jupytext (#745)

* wip * wip * more wip * fixups * coverage * add tests * fix myst test * ignore md when metadata doesnt explicitly list python as language * fix test * error 123 if failed notebook, 0 if non-python notebook * document a bit more * catch when jupytext cant parse * add test for blacken-docs on .md file * skip wrong extensoins * fix tet * fix error message if file doesnt exist * clarify instructions * clarify instructions * coverage * fix windows * add test for octave notebook * preserve substitutions * read jupytext config * try-except for reading jupytext config * try-except for reading jupytext config Co-authored-by: MarcoGorelli <>
nbQA-dev · Sep 18, 2022 · b95808d · b95808d
1 parent 8e986fd
commit b95808d
Show file tree

Hide file tree

Showing 19 changed files with 773 additions and 87 deletions.
diff --git a/README.md b/README.md
@@ -6,7 +6,7 @@
 </h1>
 
 <h3 align="center">
-	Run isort, pyupgrade, mypy, pylint, flake8, mdformat, black, blacken-docs, and more on Jupyter Notebooks
+	Run isort, pyupgrade, mypy, pylint, flake8, black, blacken-docs, and more on Jupyter Notebooks
 </h3>
 
 <p align="center">
@@ -69,6 +69,11 @@ In your [virtual environment](https://realpython.com/python-virtual-environments
 $ python -m pip install -U nbqa
 ```
 
+To also install all supported linters/formatters:
+```console
+$ python -m pip install -U "nbqa[toolchain]"
+```
+
 Or, if you are using conda:
 ```console
 $ conda install -c conda-forge nbqa
@@ -91,21 +96,21 @@ All done! ✨ 🍰 ✨
 Sort your imports with [isort](https://timothycrosley.github.io/isort/):
 
 ```console
-$ nbqa isort my_notebook.ipynb
+$ nbqa isort my_notebook.ipynb --float-to-top
 Fixing my_notebook.ipynb
 ```
 
 Upgrade your syntax with [pyupgrade](https://github.com/asottile/pyupgrade):
 
 ```console
-$ nbqa pyupgrade my_notebook.ipynb --py36-plus
+$ nbqa pyupgrade my_notebook.ipynb --py37-plus
 Rewriting my_notebook.ipynb
 ```
 
-Format your markdown cells with [mdformat](https://mdformat.readthedocs.io/en/stable/index.html):
+Format your markdown cells with [blacken-docs](https://github.com/asottile/blacken-docs):
 
 ```console
-$ nbqa mdformat my_notebook.ipynb --nbqa-md --nbqa-diff
+$ nbqa blacken-docs my_notebook.ipynb --nbqa-md --nbqa-diff
 Cell 2
 ------
 --- my_notebook.ipynb
@@ -118,6 +123,15 @@ Cell 2
 To apply these changes, remove the `--nbqa-diff` flag
 ```
 
+Format ``.md`` files saved via [Jupytext](https://github.com/mwouts/jupytext) (requires ``jupytext`` to be installed):
+
+```console
+$ nbqa black my_notebook.md
+reformatted my_notebook.md
+All done! ✨ 🍰 ✨
+1 files reformatted.
+```
+
 See [command-line examples](https://nbqa.readthedocs.io/en/latest/examples.html) for examples involving [doctest](https://docs.python.org/3/library/doctest.html), [flake8](https://flake8.pycqa.org/en/latest/), [mypy](http://mypy-lang.org/), [pylint](https://github.com/PyCQA/pylint), [autopep8](https://github.com/hhatto/autopep8), [pydocstyle](http://www.pydocstyle.org/en/stable/), and [yapf](https://github.com/google/yapf).
 
 ### Pre-commit
@@ -129,9 +143,11 @@ Here's an example of how to set up some pre-commit hooks: put this in your `.pre
   rev: 1.4.0
   hooks:
     - id: nbqa-black
+      additional_dependencies: [jupytext]  # optional, only if you're using Jupytext
     - id: nbqa-pyupgrade
-      args: [--py36-plus]
+      args: ["--py37-plus"]
     - id: nbqa-isort
+      args: ["--float-to-top"]
 ```
 
 If you need to select specific versions of any of these linters/formatters,

diff --git a/docs/configuration.rst b/docs/configuration.rst
@@ -140,20 +140,20 @@ Process markdown cells
 
 You can process markdown cells (instead of code cells) by using the :code:`--nbqa-md` CLI argument.
 
-This is useful when running tools which run on markdown files, such as ``mdformat``.
+This is useful when running tools which run on markdown files, such as ``blacken-docs``.
 
 For example, you could add the following to your :code:`pyproject.toml` file:
 
     .. code-block:: toml
 
         [tool.nbqa.md]
-        mdformat = true
+        blacken-docs = true
 
 or, from the command-line:
 
 .. code-block:: bash
 
-    nbqa mdformat notebook.ipynb --nbqa-md
+    nbqa blacken-docs notebook.ipynb --nbqa-md
 
 Shell commands
 ~~~~~~~~~~~~~~

diff --git a/docs/examples.rst b/docs/examples.rst
@@ -85,20 +85,27 @@ Check docstring style with `pydocstyle`_:
 
    $ nbqa pydocstyle my_notebook.ipynb
 
-Format markdown cells with `mdformat`_:
+Format markdown cells with `blacken-docs`_:
 
 .. code:: console
 
-   $ nbqa mdformat my_notebook.ipynb --nbqa-md
+   $ nbqa blacken-docs my_notebook.ipynb --nbqa-md
+
+Format ``.md`` file saved via `Jupytext`_ (note: requires ``jupytext`` to be installed):
+
+.. code:: console
+
+   $ nbqa black my_notebook.md
 
 .. _black: https://black.readthedocs.io/en/stable/
 .. _doctest: https://docs.python.org/3/library/doctest.html
 .. _flake8: https://flake8.pycqa.org/en/latest/
 .. _isort: https://timothycrosley.github.io/isort/
+.. _Jupytext: https://github.com/mwouts/jupytext
 .. _mypy: http://mypy-lang.org/
 .. _pylint: https://github.com/PyCQA/pylint
 .. _pyupgrade: https://github.com/asottile/pyupgrade
 .. _yapf: https://github.com/google/yapf
 .. _autopep8: https://github.com/hhatto/autopep8
 .. _pydocstyle: http://www.pydocstyle.org/en/stable/
-.. _mdformat: https://mdformat.readthedocs.io/en/stable/index.html
+.. _blacken-docs: https://github.com/asottile/blacken-docs
diff --git a/docs/history.rst b/docs/history.rst
@@ -5,6 +5,12 @@ Changelog
 1.5.0 (???)
 ~~~~~~~~~~~
 ``nbqa`` now removes empty cells which were not empty to begin with.
+Markdown files saved via ``Jupytext`` will now be processed as well
+(thanks @basnijholt and @rgommers for the suggestion!)
+If ``nbqa`` is passed an invalid notebook, it will exit 123. If it's
+passed a non-Python notebook, it'll exit 0.
+Don't try to process files with the wrong extensions, even if passed
+explicitly.
 Added support for subcommands (thanks @dnoliver for the issue, @s-weigand for the fix)
 
 1.4.0 (2022-07-17)

diff --git a/nbqa/__main__.py b/nbqa/__main__.py
@@ -1,7 +1,8 @@
 """Run third-party tool (e.g. :code:`mypy`) against notebook or directory."""
-import json
+import itertools
 import os
 import re
+import string
 import subprocess
 import sys
 import tempfile
@@ -33,7 +34,11 @@
 from nbqa.notebook_info import NotebookInfo
 from nbqa.optional import metadata
 from nbqa.output_parser import Output, map_python_line_to_nb_lines
-from nbqa.path_utils import get_relative_and_absolute_paths, remove_suffix
+from nbqa.path_utils import (
+    get_relative_and_absolute_paths,
+    read_notebook,
+    remove_suffix,
+)
 from nbqa.save_code_source import CODE_SEPARATOR
 from nbqa.text import BOLD, RESET
 
@@ -109,13 +114,33 @@ def _get_notebooks(root_dir: str) -> Iterator[Path]:
     notebooks
         All Jupyter Notebooks found in directory.
     """
-    if not os.path.isdir(root_dir):
+    try:
+        import jupytext  # noqa  # pylint: disable=unused-import,import-outside-toplevel
+    except ImportError:  # pragma: nocover
+        jupytext_installed = False
+    else:
+        jupytext_installed = True
+
+    if os.path.isfile(root_dir):
+        _, ext = os.path.splitext(root_dir)
+        if (jupytext_installed and ext in (".ipynb", ".md")) or (
+            not jupytext_installed and ext == ".ipynb"
+        ):
+            return iter((Path(root_dir),))
+        return iter([])
+
+    if not os.path.exists(root_dir):
+        # Process later, raise appropriate error message after clean up.
         return iter((Path(root_dir),))
-    return (
-        i
-        for i in Path(root_dir).rglob("*.ipynb")
-        if not re.search(EXCLUDES, str(i.resolve().as_posix()))
-    )
+
+    if jupytext_installed:
+        iterable = itertools.chain(
+            Path(root_dir).rglob("*.ipynb"), Path(root_dir).rglob("*.md")
+        )
+    else:  # pragma: nocover
+        iterable = itertools.chain(Path(root_dir).rglob("*.ipynb"))
+
+    return (i for i in iterable if not re.search(EXCLUDES, str(i.resolve().as_posix())))
 
 
 def _filter_by_include_exclude(
@@ -201,16 +226,17 @@ def _replace_temp_python_file_references_in_out_err(
     Output
         Stdout, stderr with temporary directory replaced by current working directory.
     """
+    _, ext = os.path.splitext(notebook)
     py_basename = os.path.basename(temp_python_file)
     nb_basename = os.path.basename(notebook)
     out = out.replace(py_basename, nb_basename)
     err = err.replace(py_basename, nb_basename)
 
     out = out.replace(
-        remove_suffix(py_basename, SUFFIX[md]), remove_suffix(nb_basename, ".ipynb")
+        remove_suffix(py_basename, SUFFIX[md]), remove_suffix(nb_basename, ext)
     )
     err = err.replace(
-        remove_suffix(py_basename, SUFFIX[md]), remove_suffix(nb_basename, ".ipynb")
+        remove_suffix(py_basename, SUFFIX[md]), remove_suffix(nb_basename, ext)
     )
 
     return Output(out, err)
@@ -434,7 +460,9 @@ def _get_nb_to_tmp_mapping(
         nb_to_tmp_mapping[notebook] = TemporaryFile(
             *tempfile.mkstemp(
                 dir=os.path.dirname(notebook),
-                prefix=remove_suffix(os.path.basename(notebook), ".ipynb"),
+                prefix=remove_suffix(
+                    os.path.basename(notebook), os.path.splitext(notebook)[-1]
+                ),
                 suffix=SUFFIX[md],
             )
         )
@@ -469,7 +497,7 @@ def _is_non_python_notebook(notebook: MutableMapping[str, Any]) -> bool:
     if a notebook has empty metadata, we will try to parse it anyway.
     """
     language = notebook.get("metadata", {}).get("language_info", {}).get("name", None)
-    return language is not None and language != "python"
+    return language is not None and language.rstrip(string.digits) != "python"
 
 
 def _save_code_sources(
@@ -489,11 +517,9 @@ def _save_code_sources(
     nb_info_mapping: MutableMapping[str, NotebookInfo] = {}
 
     for notebook, (file_descriptor, _) in nb_to_py_mapping.items():
-        with open(str(notebook), encoding="utf-8") as handle:
-            content = handle.read()
         try:
-            notebook_json = json.loads(content)
-            if _is_non_python_notebook(notebook_json):
+            notebook_json, _ = read_notebook(notebook)
+            if notebook_json is None or _is_non_python_notebook(notebook_json):
                 non_python_notebooks.add(notebook)
                 continue
             nb_info_mapping[notebook] = save_code_source.main(
@@ -522,21 +548,23 @@ def _save_markdown_sources(
     Record which notebooks fail to process.
     """
     failed_notebooks = {}
+    non_python_notebooks = set()
     nb_info_mapping: MutableMapping[str, NotebookInfo] = {}
 
     for notebook, (file_descriptor, _) in nb_to_md_mapping.items():
-        with open(str(notebook), encoding="utf-8") as handle:
-            content = handle.read()
         try:
-            notebook_json = json.loads(content)
+            notebook_json, _ = read_notebook(notebook)
+            if notebook_json is None or _is_non_python_notebook(notebook_json):
+                non_python_notebooks.add(notebook)
+                continue
             nb_info_mapping[notebook] = save_markdown_source.main(
                 notebook_json,
                 file_descriptor,
                 skip_celltags,
             )
         except Exception as exp_repr:  # pylint: disable=W0703
             failed_notebooks[notebook] = repr(exp_repr)
-    return SavedSources(nb_info_mapping, failed_notebooks, set())
+    return SavedSources(nb_info_mapping, failed_notebooks, non_python_notebooks)
 
 
 SAVE_SOURCES = {False: _save_code_sources, True: _save_markdown_sources}
@@ -610,14 +638,10 @@ def _main(cli_args: CLIArgs, configs: Configs) -> int:
     except FileNotFoundError as exc:
         sys.stderr.write(str(exc))
         return 1
-
     try:  # pylint disable=R0912
 
         if not nb_to_tmp_mapping:
-            sys.stderr.write(
-                "No .ipynb notebooks found in given directories: "
-                f"{' '.join(i for i in cli_args.root_dirs if os.path.isdir(i))}\n"
-            )
+            sys.stderr.write("No notebooks found in given path(s)\n")
             return 0
         saved_sources = SAVE_SOURCES[configs["md"]](
             nb_to_tmp_mapping,
@@ -626,11 +650,9 @@ def _main(cli_args: CLIArgs, configs: Configs) -> int:
             configs["dont_skip_bad_cells"],
             cli_args.command,
         )
-
-        if len(saved_sources.failed_notebooks) == len(nb_to_tmp_mapping):
-            sys.stderr.write("No valid .ipynb notebooks found\n")
-            _print_failed_notebook_errors(saved_sources.failed_notebooks)
-            return 123
+        if len(saved_sources.non_python_notebooks) == len(nb_to_tmp_mapping):
+            sys.stderr.write("No valid Python notebooks found in given path(s)\n")
+            return 0
 
         output, output_code, mutated = _run_command(
             cli_args.command,