From 4d13e2ec4d69a2813b57e714dd3db47c30f1eed9 Mon Sep 17 00:00:00 2001
From: Remco de Boer <29308176+redeboer@users.noreply.github.com>
Date: Mon, 23 Sep 2024 17:36:26 +0200
Subject: [PATCH] FEAT: move `import`s to top in notebooks (#392)

* BEHAVIOR: do not remove `nbqa.isort` table if `--imports-on-top`
* ENH: make `isort` compatible with Ruff
* FEAT: allow sorting `imports` to top of notebook
* FIX: update description of `--outsource-pixi-to-tox` flag
---
 src/compwa_policy/check_dev_files/__init__.py | 10 +++-
 src/compwa_policy/check_dev_files/ruff.py     | 56 +++++++++++++++++--
 2 files changed, 59 insertions(+), 7 deletions(-)

diff --git a/src/compwa_policy/check_dev_files/__init__.py b/src/compwa_policy/check_dev_files/__init__.py
index 3bec606d..0f1d44c2 100644
--- a/src/compwa_policy/check_dev_files/__init__.py
+++ b/src/compwa_policy/check_dev_files/__init__.py
@@ -111,7 +111,7 @@ def main(argv: Sequence[str] | None = None) -> int:
             do(pytest.main)
             do(pyupgrade.main, precommit_config, args.no_ruff)
             if not args.no_ruff:
-                do(ruff.main, precommit_config, has_notebooks)
+                do(ruff.main, precommit_config, has_notebooks, args.imports_on_top)
         if args.pin_requirements != "no":
             do(
                 update_pip_constraints.main,
@@ -197,11 +197,17 @@ def _create_argparse() -> ArgumentParser:
         action="store_true",
         default=False,
     )
+    parser.add_argument(
+        "--imports-on-top",
+        action="store_true",
+        default=False,
+        help="Sort notebook imports on the top",
+    )
     parser.add_argument(
         "--outsource-pixi-to-tox",
         action="store_true",
         default=False,
-        help="Run ",
+        help="Run tox command through pixi",
     )
     parser.add_argument(
         "--no-cspell-update",
diff --git a/src/compwa_policy/check_dev_files/ruff.py b/src/compwa_policy/check_dev_files/ruff.py
index 574b4ad6..84d2d4c6 100644
--- a/src/compwa_policy/check_dev_files/ruff.py
+++ b/src/compwa_policy/check_dev_files/ruff.py
@@ -27,7 +27,11 @@
     from compwa_policy.utilities.precommit import ModifiablePrecommit
 
 
-def main(precommit: ModifiablePrecommit, has_notebooks: bool) -> None:
+def main(
+    precommit: ModifiablePrecommit,
+    has_notebooks: bool,
+    imports_on_top: bool,
+) -> None:
     with Executor() as do, ModifiablePyproject.load() as pyproject:
         do(
             add_badge,
@@ -36,10 +40,12 @@ def main(precommit: ModifiablePrecommit, has_notebooks: bool) -> None:
         do(pyproject.remove_dependency, "radon")
         do(_remove_black, precommit, pyproject)
         do(_remove_flake8, precommit, pyproject)
-        do(_remove_isort, precommit, pyproject)
+        do(_remove_isort, precommit, pyproject, imports_on_top)
         do(_remove_pydocstyle, precommit, pyproject)
         do(_remove_pylint, precommit, pyproject)
         do(_move_ruff_lint_config, pyproject)
+        if has_notebooks and imports_on_top:
+            do(_sort_imports_on_top, precommit, pyproject)
         do(_update_ruff_config, precommit, pyproject, has_notebooks)
         do(_update_precommit_hook, precommit, has_notebooks)
         do(_update_lint_dependencies, pyproject)
@@ -87,14 +93,16 @@ def _remove_flake8(
 def _remove_isort(
     precommit: ModifiablePrecommit,
     pyproject: ModifiablePyproject,
+    imports_on_top: bool,
 ) -> None:
     with Executor() as do:
         do(__remove_nbqa_option, pyproject, "black")
-        do(__remove_nbqa_option, pyproject, "isort")
-        do(__remove_tool_table, pyproject, "isort")
         do(vscode.remove_extension_recommendation, "ms-python.isort", unwanted=True)
         do(precommit.remove_hook, "isort")
-        do(precommit.remove_hook, "nbqa-isort")
+        if not imports_on_top:
+            do(__remove_tool_table, pyproject, "isort")
+            do(__remove_nbqa_option, pyproject, "isort")
+            do(precommit.remove_hook, "nbqa-isort")
         do(vscode.remove_settings, ["isort.check", "isort.importStrategy"])
         do(remove_badge, r".*https://img\.shields\.io/badge/%20imports\-isort")
 
@@ -574,6 +582,44 @@ def _update_precommit_hook(precommit: ModifiablePrecommit, has_notebooks: bool)
     precommit.update_single_hook_repo(expected_repo)
 
 
+def _sort_imports_on_top(
+    precommit: ModifiablePrecommit, pyproject: ModifiablePyproject
+) -> None:
+    __add_isort_configuration(pyproject)
+    __add_nbqa_isort_pre_commit(precommit)
+
+
+def __add_isort_configuration(pyproject: ModifiablePyproject) -> None:
+    isort_settings = pyproject.get_table("tool.isort", create=True)
+    minimal_settings = dict(
+        profile="black",
+    )
+    if not complies_with_subset(isort_settings, minimal_settings):
+        isort_settings.update(minimal_settings)
+        msg = "Made isort configuration compatible with Ruff"
+        pyproject.changelog.append(msg)
+
+
+def __add_nbqa_isort_pre_commit(precommit: ModifiablePrecommit) -> None:
+    existing_repo = precommit.find_repo("https://github.com/nbQA-dev/nbQA")
+    excludes = None
+    if existing_repo is not None and existing_repo.get("hooks"):
+        nbqa_hook_candidates = [
+            h for h in existing_repo["hooks"] if h["id"] == "nbqa-isort"
+        ]
+        if nbqa_hook_candidates:
+            nbqa_hook = nbqa_hook_candidates[0]
+            excludes = nbqa_hook.get("exclude")
+    expected_repo = Repo(
+        repo="https://github.com/nbQA-dev/nbQA",
+        rev="1.9.0",
+        hooks=[Hook(id="nbqa-isort", args=YAML(typ="rt").load("[--float-to-top]"))],
+    )
+    if excludes is not None:
+        expected_repo["hooks"][0]["exclude"] = excludes
+    precommit.update_single_hook_repo(expected_repo)
+
+
 def _update_lint_dependencies(pyproject: ModifiablePyproject) -> None:
     if get_build_system() is None:
         return