Migrate to uv, prex, ruff, ty, etc

.github/workflows/macos-test.yml

@@ -1,20 +1,25 @@
-name: macos-test
+name: CI # TODO rename the file in a standalone PR
 
 on:
   # Trigger the workflow manually
   workflow_dispatch: ~
 
-  # Call from other workflows
-  workflow_call: ~
+  push:
+    branches:
+      - 'main'
+      - 'develop'
+
+  pull_request: ~
 
 jobs:
-  macos-pytest:
+  ci:
     strategy:
       fail-fast: true
       matrix:
-        arch_type: [ARM64] # [ARM64, X64]
-        python_version: ["3.11"] # ["3.10", "3.11", "3.12", "3.13"]
-    runs-on: [self-hosted, macOS, "${{ matrix.arch_type }}"]
+        arch_type: ["macos-ARM64", "linux-x86"]
+        python_version: ["3.10", "3.11", "3.12", "3.13"]
+    runs-on: "${{ fromJSON('{\"linux-x86\": [\"self-hosted\", \"Linux\", \"platform-builder-Rocky-8.6\"], \"macos-ARM64\": [\"self-hosted\", \"macOS\", \"ARM64\"]}')[matrix.arch_type] }}"
+    timeout-minutes: 20
     steps:
       - uses: actions/checkout@v4
       - uses: actions/setup-python@v5
@@ -23,7 +28,8 @@ jobs:
       - uses: astral-sh/setup-uv@v6
         with:
           version: 0.7.19
+      - uses: extractions/setup-just@v3
       - run: |
-          uv venv venv && source venv/bin/activate
-          uv pip install '.[tests]'
-          pytest
+          uv sync --python "${{ matrix.python_version }}"
+          just fmt
+          just val

.pre-commit-config.yaml

@@ -1,74 +1,29 @@
 repos:
 - repo: https://github.com/pre-commit/pre-commit-hooks
-  rev: v4.6.0
+  rev: v6.0.0
   hooks:
   - id: check-yaml # Check YAML files for syntax errors only
     args: [--unsafe, --allow-multiple-documents]
-  - id: debug-statements # Check for debugger imports and py37+ breakpoint()
-  - id: end-of-file-fixer # Ensure files end in a newline
-  - id: trailing-whitespace # Trailing whitespace checker
-  - id: no-commit-to-branch # Prevent committing to main / master
-  - id: check-added-large-files # Check for large files added to git
+#  - id: no-commit-to-branch # NOTE prevents commit to main/master, but since we run prek on that branch it makes no sense
+  - id: check-added-large-files
+    exclude: |
+      (?x)(
+        .*uv.lock|
+        .*pylock.toml
+      )
   - id: check-merge-conflict # Check for files that contain merge conflict
--   repo: https://github.com/pre-commit/pygrep-hooks
-    rev: v1.10.0  # Use the ref you want to point at
-    hooks:
-    -   id: python-use-type-annotations # Check for missing type annotations
-    -   id: python-check-blanket-noqa # Check for # noqa: all
-    -   id: python-no-log-warn # Check for log.warn
-- repo: https://github.com/psf/black-pre-commit-mirror
-  rev: 25.1.0
-  hooks:
-  - id: black
-    # args: [--line-length=120]
-- repo: https://github.com/pycqa/isort
-  rev: 5.13.2
-  hooks:
-  - id: isort
-    args:
-    # - -l 120
-    - --profile black
 - repo: https://github.com/astral-sh/ruff-pre-commit
-  rev: v0.6.5
+  rev: v0.12.2
   hooks:
-  - id: ruff
-    # Next line if for documenation cod snippets
-    exclude: '^[^_].*_\.py$'
+  - id: ruff # better black/flake/isort
+    files: ^src/
     args:
-    # - --line-length=120
+    - --select
+    - I # isorting
     - --fix
     - --exit-non-zero-on-fix
-    - --preview
-- repo: https://github.com/sphinx-contrib/sphinx-lint
-  rev: v1.0.0
-  hooks:
-  - id: sphinx-lint
-# For now, we use it. But it does not support a lot of sphinx features
-- repo: https://github.com/dzhu/rstfmt
-  rev: v0.0.14
-  hooks:
-  - id: rstfmt
-    exclude: 'cli/.*' # Because we use argparse
-- repo: https://github.com/b8raoult/pre-commit-docconvert
-  rev: "0.1.5"
-  hooks:
-  - id: docconvert
-    args: ["numpy"]
-# - repo: https://github.com/tox-dev/pyproject-fmt
-#   rev: "2.2.3"
-#   hooks:
-#   - id: pyproject-fmt
-# -   repo: https://github.com/jshwi/docsig # Check docstrings against function sig
-#     rev: v0.60.1
-#     hooks:
-#     -   id: docsig
-#         args:
-#         - --ignore-no-params # Allow docstrings without parameters
-#         - --check-dunders    # Check dunder methods
-#         - --check-overridden # Check overridden methods
-#         - --check-protected  # Check protected methods
-#         - --check-class      # Check class docstrings
-#         - --disable=E113     # Disable empty docstrings
-#         - --summary          # Print a summary
+  - id: ruff-format
+    files: ^backend/
 ci:
   autoupdate_schedule: monthly
+  autoupdate_commit_msg: "chore(deps): pre-commit.ci autoupdate"

.python_version

@@ -0,0 +1 @@
+3.11.9

justfile

@@ -10,12 +10,9 @@
 set dotenv-path := ".env"
 
 val:
-	mypy src/cascade --ignore-missing-imports
-	mypy tests --ignore-missing-imports
-	pytest -n 8 tests
-
+    uv run ty check src/cascade
+    uv run ty check tests/cascade
+    # TODO eventually broaden type coverage to ekw as well
+    uv run pytest -n8 tests
 fmt:
-    # TODO replace with pre-commit
-    isort --profile black .
-    black .
-    flake8 .
+    uv run prek --all-files

pyproject.toml

@@ -7,7 +7,7 @@
 # nor does it submit to any jurisdiction.
 
 [build-system]
-requires = ["setuptools>=65", "setuptools_scm[toml]>=6.2"]
+requires = ["setuptools>=80", "setuptools_scm[toml]>=6.2", "packaging>=25"]
 build-backend = "setuptools.build_meta"
 
 [project]
@@ -39,21 +39,13 @@ dependencies = [
 dynamic = ["version"]
 readme = "README.md"
 
-[project.optional-dependencies]
-tests = ["pytest", "pytest-xdist", "earthkit-data"]
-lint = ["black", "isort", "flake8"]
-gpu = ["jax[cpu]", "jax[cuda11_pip]", "cupy-cuda11x", "numba"]
-examples = ["cftime", "bokeh"]
-earthkit = ["earthkit-data"]
-all = ["cascade[tests,kubernetes,lint,gpu,examples,earthkit]"]
+[dependency-groups]
+dev = ["pytest", "pytest-xdist>=3.8", "prek", "ty==0.0.2", "build", "bokeh"]
 
 [tool.setuptools]
 include-package-data = true
 zip-safe = false
 
-[tool.setuptools.dynamic]
-readme = {file = ["readme.md"], content-type = "text/markdown"}
-
 [tool.setuptools_scm]
 write_to = "src/earthkit/workflows/_version.py"
 write_to_template = '''# Do not change! Do not track in version control!
@@ -65,17 +57,12 @@ local_scheme = "no-local-version"
 include = ["earthkit.workflows*", "cascade*"]
 where = ["src"]
 
-[tool.isort]
-profile = "black"
-
-[tool.mypy]
-plugins = "pydantic.mypy"
-
 [tool.pytest.ini_options]
 log_cli = true
 log_cli_level = "DEBUG"
 testpaths = ["tests/"]
 addopts = "-n8"
 
-[tool.ruff.lint]
-ignore = ["E722", "E731", "E741"]
+[tool.ruff]
+line-length = 140
+lint.ignore = [ "E722", "E731", "E741" ]

src/cascade/benchmarks/anemoi.py

@@ -10,7 +10,7 @@
 
 
 def get_graph(lead_time, ensemble_members, CKPT=None, date="2024-12-02T00:00"):
-    import anemoicascade as ac
+    import anemoicascade as ac  # ty: ignore[unresolved-import]
 
     CKPT = (
         CKPT

src/cascade/benchmarks/dask.py

@@ -5,9 +5,9 @@
 from cascade.low.dask import graph2job
 
 
-def get_job(job: str) -> JobInstance:
+def get_job(job_name: str) -> JobInstance:
 
-    if job == "add":
+    if job_name == "add":
 
         def add(x, y):
             result = x + y
@@ -21,7 +21,7 @@ def add(x, y):
             dataset for task in job.tasks for dataset in job.outputs_of(task)
         ]
         return job
-    elif job == "groupby":
+    elif job_name == "groupby":
         df = dd.DataFrame.from_dict({"x": [0, 0, 1, 1], "y": [1, 2, 3, 4]})
         df = df.groupby("x").sum()
         job = graph2job(df.__dask_graph__())
@@ -30,4 +30,4 @@ def add(x, y):
         ]
         return job
     else:
-        raise NotImplementedError(job)
+        raise NotImplementedError(job_name)

src/cascade/benchmarks/dist.py

@@ -26,7 +26,7 @@ def dist_func_torch(a: int) -> int:
     import datetime as dt
 
     import numpy as np
-    import torch.distributed as dist
+    import torch.distributed as dist  # ty: ignore[unresolved-import]
 
     world_size = int(os.environ["CASCADE_GANG_WORLD_SIZE"])
     rank = int(os.environ["CASCADE_GANG_RANK"])
@@ -61,8 +61,8 @@ def dist_func_jax(a: int) -> int:
     os.environ["JAX_NUM_CPU_DEVICES"] = "1"
     os.environ["JAX_PLATFORM_NAME"] = "cpu"
     os.environ["JAX_PLATFORMS"] = "cpu"
-    import jax
-    import jax.numpy as jp
+    import jax  # ty: ignore[unresolved-import]
+    import jax.numpy as jp  # ty: ignore[unresolved-import]
 
     jax.config.update("jax_platforms", "cpu")
     jax.config.update("jax_platform_name", "cpu")