chore: enforce linting and type checks (#15)

.github/actions/setup-poetry-environment/action.yml

@@ -0,0 +1,39 @@
+name: "Setup test environment"
+description: "Sets up Python, Poetry and dependencies"
+
+inputs:
+  python:
+    description: "Python version to use"
+    default: "3.11"
+    required: false
+  poetry:
+    description: "Poetry version to use"
+    default: 1.7.1
+    required: false
+
+runs:
+  using: "composite"
+
+  steps:
+    - uses: actions/setup-python@v5
+      with:
+        python-version: ${{ inputs.python }}
+
+    - uses: snok/install-poetry@v1
+      with:
+        version: ${{ inputs.poetry }}
+        virtualenvs-create: true
+        virtualenvs-in-project: true
+
+    - name: Load cached venv
+      id: cached-poetry-dependencies
+      uses: actions/cache@v4
+      with:
+        path: .venv
+        key: venv-${{ runner.os }}-${{ steps.setup-python.outputs.python-version }}-${{ hashFiles('**/poetry.lock') }}
+
+    - if: steps.cached-poetry-dependencies.outputs.cache-hit != 'true'
+      run: |
+        poetry lock --no-update
+        poetry install --no-interaction
+      shell: bash

.github/workflows/pr-report.yml

@@ -0,0 +1,23 @@
+name: PR Report
+
+on:
+  pull_request:
+    branches:
+      - main
+  workflow_dispatch:
+
+env:
+  PYTHON_VERSION: 3.11
+
+jobs:
+  pre-commit:
+    name: Linting and type checking
+    runs-on: ubuntu-latest
+    steps:
+      - uses: actions/checkout@v4
+      - name: Setup poetry environment
+        uses: ./.github/actions/setup-poetry-environment
+      - name: Run linting, type checking and testing
+        uses: pre-commit/[email protected]
+        with:
+          extra_args: "--all-files --hook-stage=push"

.pre-commit-config.yaml

@@ -0,0 +1,34 @@
+default_language_version:
+  python: python3.11
+
+repos:
+  - repo: https://github.com/pre-commit/pre-commit-hooks
+    rev: f71fa2c1f9cf5cb705f73dffe4b21f7c61470ba9 # 4.4.0
+    hooks:
+      - id: check-yaml
+      - id: check-json
+      - id: check-toml
+      - id: end-of-file-fixer
+      - id: trailing-whitespace
+  - repo: https://github.com/pappasam/toml-sort
+    rev: b9b6210da457c38122995e434b314f4c4a4a923e # 0.23.1
+    hooks:
+      - id: toml-sort-fix
+        files: ^.+.toml$
+  - repo: https://github.com/astral-sh/ruff-pre-commit
+    rev: v0.2.1
+    hooks:
+      - id: ruff
+        args:
+          - --fix
+          - --exit-non-zero-on-fix
+#          - --ignore=E501 # line-too-long
+#          - --ignore=F631 # assert-tuple
+#          - --ignore=E741 # ambiguous-variable-name
+      - id: ruff-format
+        files: ^src\/.+\.py$
+#  - repo: https://github.com/pre-commit/mirrors-mypy
+#    rev: v1.8.0
+#    hooks:
+#      - id: mypy
+default_stages: [push]

dev-setup.sh

@@ -0,0 +1,5 @@
+#!/bin/zsh
+poetry env use 3.11 # Create the virtual environment if it does not exist
+source $(poetry env info --path)/bin/activate # Activate and enter the virtual environment
+poetry install --with=dev # Install dev dependencies
+pre-commit install --install-hooks --overwrite -t pre-push # Set up pre-commit hooks

pyproject.toml

@@ -29,6 +29,7 @@ matplotlib = "^3.8.2"
 mypy = "^1.8.0"
 ruff = "^0.2.1"
 toml-sort = "^0.23.1"
+pre-commit = "^3.6.1"
 
 [tool.ruff]
 ignore = []

src/evaluation/knn.py

@@ -41,9 +41,7 @@ def __init__(
         self.num_classes = num_classes or labels.max() + 1
 
         embeddings = self.normalize(embeddings)
-        index, self.index_infos = build_index(
-            embeddings, save_on_disk=False, verbose=logging.ERROR
-        )
+        index, self.index_infos = build_index(embeddings, save_on_disk=False, verbose=logging.ERROR)
 
         if index is None:
             raise ValueError("Failed to build an index for knn search")
@@ -71,9 +69,7 @@ def predict(self, embeddings: Embeddings) -> tuple[ProbabilityArray, ClassArray]
         # We can shape of a factor self.k if we count differently here.
         weighted_count = np.zeros((n, self.num_classes, self.k), dtype=np.float32)
         weighted_count[
-            np.tile(np.arange(n), (self.k,)).reshape(
-                -1
-            ),  # [0, 0, .., 0_k, 1, 1, .., 1_k, ..]
+            np.tile(np.arange(n), (self.k,)).reshape(-1),  # [0, 0, .., 0_k, 1, 1, .., 1_k, ..]
             nearest_classes.reshape(-1),  # [class numbers]
             np.tile(np.arange(self.k), (n,)),  # [0, 1, .., k-1, 0, 1, .., k-1, ..]
         ] = 1 / dists.reshape(-1)

src/evaluation/linear_probe.py

@@ -40,9 +40,7 @@ def __init__(
                 Cs=[1e-3, 1e-2, 1e-1, 1e0, 1e1, 1e2, 1e3], random_state=0, **params
             ).fit(embeddings, labels)  # type: ignore
         else:
-            self.classifier = LogisticRegression(random_state=0, **params).fit(
-                embeddings, labels
-            )
+            self.classifier = LogisticRegression(random_state=0, **params).fit(embeddings, labels)
 
     @property
     def dim(self) -> int: