ci: Switch from flake8, isort and black to ruff (#554)

Structural changes:
* Add ruff dependency, configuration and pre-commit hook
* Add ruff in the typing check and rename it code-quality
* Remove check-todos (now done by ruff in CI)

Code changes to adapt to ruff:
* Use strict=True when zipping
* Use | None instead of Optional
* Use | in isinstance instead of tuple
* Rename unused for-loop counters to _
* Stop using .format in strings
* Stop importing Iterable and Callable from typing
* Group "as" imports
* Stop opening files explicitly with "r" mode
---------
Co-authored-by: Valérian Rey <valerian.rey@gmail.com>

Author: PierreQuinton

.github/workflows/checks.yml

@@ -109,8 +109,8 @@ jobs:
           # This reduces false positives due to rate limits
           GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }}
 
-  typing:
-    name: Typing correctness
+  code-quality:
+    name: Code quality (ty and ruff)
     runs-on: ubuntu-latest
     steps:
       - name: Checkout repository
@@ -128,32 +128,5 @@ jobs:
       - name: Run ty
         run: uv run ty check --output-format=github
 
-  check-todos:
-    name: Absence of TODOs
-    runs-on: ubuntu-latest
-    steps:
-      - name: Checkout code
-        uses: actions/checkout@v6
-
-      - name: Scan for TODO strings
-        run: |
-          echo "Scanning codebase for TODOs..."
-
-          git grep -nE "TODO" -- . ':(exclude).github/workflows/*' > todos_found.txt || true
-
-          if [ -s todos_found.txt ]; then
-            echo "❌ ERROR: Found TODOs in the following files:"
-            echo "-------------------------------------------"
-
-            while IFS=: read -r file line content; do
-              echo "::error file=$file,line=$line::TODO found at $file:$line - must be resolved before merge:%0A$content"
-            done < todos_found.txt
-
-            echo "-------------------------------------------"
-            echo "Please resolve these TODOs or track them in an issue before merging."
-
-            exit 1
-          else
-            echo "✅ No TODOs found. Codebase is clean!"
-            exit 0
-          fi
+      - name: Run ruff
+        run: uv run ruff check --output-format=github

.pre-commit-config.yaml

@@ -9,32 +9,12 @@ repos:
     -   id: check-docstring-first  # Check a common error of defining a docstring after code.
     -   id: check-merge-conflict  # Check for files that contain merge conflict strings.
 
--   repo: https://github.com/PyCQA/flake8
-    rev: 7.3.0
+-   repo: https://github.com/astral-sh/ruff-pre-commit
+    rev: v0.9.0
     hooks:
-    -   id: flake8  # Check style and syntax. Does not modify code, issues have to be solved manually.
-        args: [
-            '--ignore=E501,E203,W503,E402',  # Ignore line length problems, space after colon problems, line break occurring before a binary operator problems, module level import not at top of file problems.
-        ]
-
--   repo: https://github.com/pycqa/isort
-    rev: 7.0.0
-    hooks:
-    -   id: isort  # Sort imports.
-        args: [
-            --multi-line=3,
-            --line-length=100,
-            --trailing-comma,
-            --force-grid-wrap=0,
-            --use-parentheses,
-            --ensure-newline-before-comments,
-        ]
-
--   repo: https://github.com/psf/black-pre-commit-mirror
-    rev: 25.12.0
-    hooks:
-    -   id: black  # Format code.
-        args: [--line-length=100]
+      - id: ruff
+        args: [ --fix, --ignore, FIX ]
+      - id: ruff-format
 
 ci:
     autoupdate_commit_msg: 'chore: Update pre-commit hooks'

pyproject.toml

@@ -66,6 +66,7 @@ Changelog = "https://github.com/TorchJD/torchjd/blob/main/CHANGELOG.md"
 
 [dependency-groups]
 check = [
+    "ruff>=0.14.14",
     "ty>=0.0.14",
     "pre-commit>=2.9.2",  # isort doesn't work before 2.9.2
 ]
@@ -122,6 +123,32 @@ exclude_lines = [
     "@overload",
 ]
 
+[tool.ruff]
+line-length = 100
+target-version = "py310"
+
+[tool.ruff.lint]
+select = [
+    "E",   # pycodestyle Error
+    "F",   # Pyflakes
+    "W",   # pycodestyle Warning
+    "I",   # isort
+    "UP",  # pyupgrade
+    "B",   # flake8-bugbear
+    "FIX",  # flake8-fixme
+]
+
+ignore = [
+    "E501", # line-too-long (handled by the formatter)
+    "E402", # module-import-not-at-top-of-file
+]
+
+[tool.ruff.lint.isort]
+combine-as-imports = true
+
+[tool.ruff.format]
+quote-style = "double"
+
 [tool.ty.src]
 include = ["src", "tests"]
 exclude = ["src/torchjd/aggregation/_nash_mtl.py"]

src/torchjd/__init__.py

@@ -1,8 +1,7 @@
 from collections.abc import Callable
 from warnings import warn as _warn
 
-from .autojac import backward as _backward
-from .autojac import mtl_backward as _mtl_backward
+from .autojac import backward as _backward, mtl_backward as _mtl_backward
 
 _deprecated_items: dict[str, tuple[str, Callable]] = {
     "backward": ("autojac", _backward),

src/torchjd/aggregation/_graddrop.py

@@ -74,5 +74,5 @@ def __str__(self) -> str:
         if self.leak is None:
             leak_str = ""
         else:
-            leak_str = f"([{', '.join(['{:.2f}'.format(l_).rstrip('0') for l_ in self.leak])}])"
+            leak_str = f"([{', '.join([f'{l_:.2f}'.rstrip('0') for l_ in self.leak])}])"
         return f"GradDrop{leak_str}"

src/torchjd/aggregation/_mgda.py

@@ -53,7 +53,7 @@ def forward(self, gramian: PSDMatrix, /) -> Tensor:
         dtype = gramian.dtype
 
         alpha = torch.ones(gramian.shape[0], device=device, dtype=dtype) / gramian.shape[0]
-        for i in range(self.max_iters):
+        for _ in range(self.max_iters):
             t = torch.argmin(gramian @ alpha)
             e_t = torch.zeros(gramian.shape[0], device=device, dtype=dtype)
             e_t[t] = 1.0

src/torchjd/aggregation/_utils/str.py

@@ -7,5 +7,5 @@ def vector_to_str(vector: Tensor) -> str:
     `1.23, 1., ...`.
     """
 
-    weights_str = ", ".join(["{:.2f}".format(value).rstrip("0") for value in vector])
+    weights_str = ", ".join([f"{value:.2f}".rstrip("0") for value in vector])
     return weights_str

src/torchjd/autogram/_gramian_accumulator.py

@@ -1,5 +1,3 @@
-from typing import Optional
-
 from torchjd._linalg import PSDMatrix
 
 
@@ -13,7 +11,7 @@ class GramianAccumulator:
     """
 
     def __init__(self) -> None:
-        self._gramian: Optional[PSDMatrix] = None
+        self._gramian: PSDMatrix | None = None
 
     def reset(self) -> None:
         self._gramian = None
@@ -25,7 +23,7 @@ def accumulate_gramian(self, gramian: PSDMatrix) -> None:
             self._gramian = gramian
 
     @property
-    def gramian(self) -> Optional[PSDMatrix]:
+    def gramian(self) -> PSDMatrix | None:
         """
         Get the Gramian matrix accumulated so far.
 

src/torchjd/autogram/_gramian_computer.py

@@ -1,5 +1,5 @@
 from abc import ABC, abstractmethod
-from typing import Optional, cast
+from typing import cast
 
 from torch import Tensor
 from torch.utils._pytree import PyTree
@@ -16,12 +16,14 @@ def __call__(
         grad_outputs: tuple[Tensor, ...],
         args: tuple[PyTree, ...],
         kwargs: dict[str, PyTree],
-    ) -> Optional[PSDMatrix]:
+    ) -> PSDMatrix | None:
         """Compute what we can for a module and optionally return the gramian if it's ready."""
 
+    @abstractmethod
     def track_forward_call(self) -> None:
         """Track that the module's forward was called. Necessary in some implementations."""
 
+    @abstractmethod
     def reset(self) -> None:
         """Reset state if any. Necessary in some implementations."""
 
@@ -40,7 +42,7 @@ class JacobianBasedGramianComputerWithCrossTerms(JacobianBasedGramianComputer):
     def __init__(self, jacobian_computer: JacobianComputer):
         super().__init__(jacobian_computer)
         self.remaining_counter = 0
-        self.summed_jacobian: Optional[Matrix] = None
+        self.summed_jacobian: Matrix | None = None
 
     def reset(self) -> None:
         self.remaining_counter = 0
@@ -55,7 +57,7 @@ def __call__(
         grad_outputs: tuple[Tensor, ...],
         args: tuple[PyTree, ...],
         kwargs: dict[str, PyTree],
-    ) -> Optional[PSDMatrix]:
+    ) -> PSDMatrix | None:
         """Compute what we can for a module and optionally return the gramian if it's ready."""
 
         jacobian_matrix = self.jacobian_computer(rg_outputs, grad_outputs, args, kwargs)

src/torchjd/autogram/_module_hook_manager.py

@@ -141,7 +141,7 @@ def __call__(
             *rg_outputs,
         )
 
-        for idx, output in zip(rg_output_indices, autograd_fn_rg_outputs):
+        for idx, output in zip(rg_output_indices, autograd_fn_rg_outputs, strict=True):
             flat_outputs[idx] = output
 
         return tree_unflatten(flat_outputs, output_spec)