Support for subcommands (#747)

* ✨ Added support for subcommands

* add test

* document

Co-authored-by: s-weigand <s.weigand.phy@gmail.com>

docs/history.rst

@@ -5,6 +5,7 @@ Changelog
 1.5.0 (???)
 ~~~~~~~~~~~
 ``nbqa`` now removes empty cells which were not empty to begin with.
+Added support for subcommands (thanks @dnoliver for the issue, @s-weigand for the fix)
 
 1.4.0 (2022-07-17)
 ~~~~~~~~~~~~~~~~~~

nbqa/__main__.py

@@ -20,6 +20,7 @@
     Sequence,
     Set,
     Tuple,
+    cast,
 )
 
 import tomli
@@ -267,15 +268,18 @@ def _run_command(
     """
     before = [_get_mtimes(i) for i in args]
 
+    main_command, *sub_commands = command.split()
+
     my_env = os.environ.copy()
-    if command == "mypy" and "MYPY_FORCE_COLOR" not in my_env:
+    if main_command == "mypy" and "MYPY_FORCE_COLOR" not in my_env:
         my_env["MYPY_FORCE_COLOR"] = "1"
 
     if shell:
-        cmd = [command]
+        # We already checked that which does not return None
+        cmd = [cast(str, which(main_command)), *sub_commands]
     else:
-        python_module = COMMAND_TO_PYTHON_MODULE.get(command, command)
-        cmd = [sys.executable, "-m", python_module]
+        python_module = COMMAND_TO_PYTHON_MODULE.get(main_command, main_command)
+        cmd = [sys.executable, "-m", python_module, *sub_commands]
 
     output = subprocess.run(
         [*cmd, *args, *cmd_args],
@@ -695,12 +699,13 @@ def _check_command_is_installed(command: str, *, shell: bool) -> None:
     CommandNotFoundError
         If third-party tool isn't available as a script in $PATH.
     """
+    main_command, *_ = command.split()
     if shell:
-        if which(command):
+        if which(main_command):
             return
-        raise CommandNotFoundError(command)
+        raise CommandNotFoundError(main_command)
 
-    python_module = COMMAND_TO_PYTHON_MODULE.get(command, command)
+    python_module = COMMAND_TO_PYTHON_MODULE.get(main_command, main_command)
     try:
         command_version = metadata.version(python_module)
     except metadata.PackageNotFoundError:
@@ -712,13 +717,15 @@ def _check_command_is_installed(command: str, *, shell: bool) -> None:
             ):  # pragma: nocover(py<37)
                 # I presume the lack of coverage in Python3.6 here is a bug, as all
                 # these branches are actually covered.
-                raise ModuleNotFoundError(_get_command_not_found_msg(command)) from None
+                raise ModuleNotFoundError(
+                    _get_command_not_found_msg(main_command)
+                ) from None
     else:
-        if command in MIN_VERSIONS:
-            min_version = MIN_VERSIONS[command]
+        if main_command in MIN_VERSIONS:
+            min_version = MIN_VERSIONS[main_command]
             if parse_version(command_version) < parse_version(min_version):
                 raise UnsupportedPackageVersionError(
-                    command, command_version, min_version
+                    main_command, command_version, min_version
                 )
 
 

tests/local_script.py

@@ -0,0 +1,18 @@
+"""Local module with subcommand."""
+import argparse
+import sys
+from typing import Optional, Sequence
+
+
+def main(argv: Optional[Sequence[str]] = None) -> int:
+    """Print word (subcommand), ignore paths"""
+    parser = argparse.ArgumentParser()
+    parser.add_argument("word", nargs=1)
+    parser.add_argument("paths", nargs="*")
+    args = parser.parse_args(argv)
+    print(args.word)
+    return 0
+
+
+if __name__ == "__main__":
+    sys.exit(main())

tests/test_local_script.py

@@ -1,10 +1,14 @@
 """Tets running local script."""
 import os
+from typing import TYPE_CHECKING
 
 import pytest
 
 from nbqa.__main__ import main
 
+if TYPE_CHECKING:
+    from _pytest.capture import CaptureFixture
+
 
 def test_local_script() -> None:
     """Test local script is picked up."""
@@ -45,3 +49,10 @@ def test_local_nonfound() -> None:
             main(["fdsfda", "."])
     finally:
         os.chdir(cwd)
+
+
+def test_with_subcommand(capsys: "CaptureFixture") -> None:
+    """Check subcommand is picked up by module."""
+    main(["tests.local_script foo", "."])
+    out, _ = capsys.readouterr()
+    assert out == "['foo']\n"

tests/test_nbqa_shell.py

@@ -1,6 +1,7 @@
 """Ensure the --nbqa-shell flag correctly calls the underlying command."""
 import os
 import sys
+from shutil import which
 from subprocess import CompletedProcess
 from typing import List
 
@@ -27,11 +28,11 @@ def test_nbqa_shell(monkeypatch: MonkeyPatch, capsys: CaptureFixture) -> None:
     monkeypatch.setattr("subprocess.run", subprocess_run)
 
     args = ["black", "--nbqa-shell", path]
-    expected_run = ["black", path]
+    expected_run = [which("black"), path]
     main(args)
     out, err = capsys.readouterr()
     received = err.strip()
-    expected = _message(args=expected_run)
+    expected = _message(args=expected_run)  # type:ignore[arg-type]
     assert (
         received == expected
     ), f"nbqa called unexpected `{received}` instead of `{expected}` for args `{args}`"