Enable ruff annotation rules (ANN)

.github/bump_version.py

@@ -47,7 +47,7 @@ def bump_version(version: str, bump: str) -> str:
     return f"{major}.{minor}.{patch + 1}"
 
 
-def update_file(path: Path, old_version: str, new_version: str):
+def update_file(path: Path, old_version: str, new_version: str) -> None:
     text = path.read_text()
     updated = text.replace(
         f'version = "{old_version}"',
@@ -58,7 +58,7 @@ def update_file(path: Path, old_version: str, new_version: str):
         print(f"  Updated {path}")
 
 
-def main():
+def main() -> None:
     root = Path(__file__).resolve().parent.parent
     pyproject = root / "pyproject.toml"
     changelog_dir = root / "changelog.d"

.github/find-api-model-versions.py

@@ -88,7 +88,7 @@ def find_api_model_versions() -> dict[str, str]:
     }
 
 
-def find_api_model_versions_and_output_to_github():
+def find_api_model_versions_and_output_to_github() -> None:
     """
     Find the API model versions and output them to a file for GitHub.
     """

.github/scripts/check_updates.py

@@ -13,6 +13,7 @@
 import argparse
 import re
 import sys
+from typing import TypedDict
 
 import requests
 
@@ -23,7 +24,12 @@
 }
 
 
-def fetch_changelog(package):
+class ChangelogEntry(TypedDict):
+    version: str
+    changes: dict[str, list[str]]
+
+
+def fetch_changelog(package: str) -> str | None:
     """Fetch CHANGELOG.md from the package's GitHub repo."""
     repo = REPO_MAP.get(package)
     if not repo:
@@ -36,18 +42,18 @@ def fetch_changelog(package):
     return None
 
 
-def parse_version(version_str):
+def parse_version(version_str: str) -> tuple[int, ...]:
     """Parse a version string into a comparable tuple."""
     return tuple(map(int, version_str.split(".")))
 
 
-def parse_changelog(text):
+def parse_changelog(text: str | None) -> list[ChangelogEntry]:
     """Parse Keep-a-Changelog markdown into structured entries."""
     if not text:
         return []
 
-    entries = []
-    current_entry = None
+    entries: list[ChangelogEntry] = []
+    current_entry: ChangelogEntry | None = None
     current_category = None
 
     for line in text.splitlines():
@@ -77,7 +83,9 @@ def parse_changelog(text):
     return entries
 
 
-def get_changes_between(changelog, old_version, new_version):
+def get_changes_between(
+    changelog: list[ChangelogEntry], old_version: str, new_version: str
+) -> list[ChangelogEntry]:
     """Extract changelog entries between two versions."""
     old_v = parse_version(old_version)
     new_v = parse_version(new_version)
@@ -88,9 +96,14 @@ def get_changes_between(changelog, old_version, new_version):
     ]
 
 
-def format_changes(entries):
+def format_changes(entries: list[ChangelogEntry]) -> str:
     """Format changelog entries as markdown sections."""
-    buckets = {"added": [], "changed": [], "fixed": [], "removed": []}
+    buckets: dict[str, list[str]] = {
+        "added": [],
+        "changed": [],
+        "fixed": [],
+        "removed": [],
+    }
     for entry in entries:
         for cat, items in entry.get("changes", {}).items():
             if cat in buckets:
@@ -105,7 +118,7 @@ def format_changes(entries):
     return "\n\n".join(sections) if sections else ""
 
 
-def main():
+def main() -> int:
     parser = argparse.ArgumentParser(
         description="Fetch changelog entries between two package versions."
     )

AGENTS.md

@@ -122,8 +122,9 @@ To add or change a rule, read `docs/engineering/skills/rules-registry.md`.
 - `api.quality.environment-variables-load`: Required environment variables and credentials load in a deployed-like environment. (blocking: CI job `test_env_vars`)
 - `api.quality.mypy-v2`: Shared query models, FastAPI dependencies, and v2 routes and services pass mypy with untyped definitions disallowed. (blocking: mypy)
 - `api.quality.ruff-format`: All Python files are formatted with ruff format. (blocking: ruff format)
-- `api.quality.ruff-lint`: All Python code passes ruff check with every rule family selected in pyproject.toml, with no ratchet; families are added one themed group at a time until the selection is ALL. (blocking: ruff check (E4, E7, E9, F, FLY, FURB, I, PIE, PYI, RET, RSE, TC, UP, W))
+- `api.quality.ruff-lint`: All Python code passes ruff check with every rule family selected in pyproject.toml, with no ratchet; families are added one themed group at a time until the selection is ALL. (blocking: ruff check (ANN, E4, E7, E9, F, FLY, FURB, I, PIE, PYI, RET, RSE, TC, UP, W))
 - `api.quality.tests-pass`: The unit, contract and to-refactor test suites pass. (blocking: CI job `test`)
+- `api.quality.tests-ruff-exemptions`: Test files may use pytest's bare assert and need no docstrings, annotations or named constants for magic numbers (S101, D1, ANN, PLR2004); every other ruff rule applies to tests. (blocking: ruff per-file ignore of ANN, D1, PLR2004, S101 for `**/tests/**`)
 
 ### Observability
 

changelog.d/ruff-annotations.changed.md

@@ -0,0 +1 @@
+Enabled ruff's annotation rules (ANN) and annotated every function outside tests; test files are exempt from ANN, missing docstrings, magic numbers and bare-assert rules.

dashboard/experiments/gpt4_api_user.py

@@ -13,7 +13,7 @@
 
 
 @st.cache_data
-def get_embeddings_df():
+def get_embeddings_df() -> pd.DataFrame:
     embeddings_df = pd.read_csv("parameter_embeddings.csv.gz", compression="gzip")
     embeddings_df.parameter_embedding = embeddings_df.parameter_embedding.apply(
         lambda x: eval(x)
@@ -22,7 +22,7 @@ def get_embeddings_df():
 
 
 @st.cache_data
-def embed(prompt, engine="text-embedding-ada-002"):
+def embed(prompt: str, engine: str = "text-embedding-ada-002") -> list[float]:
     return get_embedding(prompt, engine=engine)
 
 
@@ -58,7 +58,7 @@ def embed(prompt, engine="text-embedding-ada-002"):
 # display in streamlit
 
 
-def write_prompt(question, relevant_parameters):
+def write_prompt(question: str, relevant_parameters: object) -> str:
     return f"""Question: {question}
     Here's some metadata about the parameters that are relevant to your question:
     Relevant parameters: {relevant_parameters}
@@ -78,7 +78,7 @@ def write_prompt(question, relevant_parameters):
 
 
 @st.cache_data
-def get_gpt4_response(prompt):
+def get_gpt4_response(prompt: str) -> str:
     return openai.ChatCompletion.create(
         model="gpt-4",
         messages=[{"role": "user", "content": prompt}],
@@ -93,7 +93,7 @@ def get_gpt4_response(prompt):
 
 
 @st.cache_data
-def get_policyengine_impact(reform):
+def get_policyengine_impact(reform: object) -> object:
     payload = {"data": reform}
     url = "https://api.policyengine.org/uk/policy"
 

docs/generated/rules.md

@@ -8,11 +8,11 @@ Every rule that governs `policyengine-api` is listed here with how it is enforce
 
 | Tier | enforced | ratcheted | advisory | doc-only | Total |
 | --- | ---: | ---: | ---: | ---: | ---: |
-| mechanical | 41 | 0 | 1 | 12 | 54 |
+| mechanical | 42 | 0 | 1 | 12 | 55 |
 | metric | 0 | 0 | 0 | 0 | 0 |
 | judgment | 0 | 0 | 0 | 0 | 0 |
 
-**Blocking:** 41 of 54 rules.
+**Blocking:** 42 of 55 rules.
 
 ## Gaps
 
@@ -268,8 +268,9 @@ flowchart LR
 | `api.quality.environment-variables-load` | Required environment variables and credentials load in a deployed-like environment. | mechanical | enforced | yes | CI job `test_env_vars` | `pr.yml:test_env_vars` | — |
 | `api.quality.mypy-v2` | Shared query models, FastAPI dependencies, and v2 routes and services pass mypy with untyped definitions disallowed. | mechanical | enforced | yes | mypy | `pr.yml:typecheck-v2` | — |
 | `api.quality.ruff-format` | All Python files are formatted with ruff format. | mechanical | enforced | yes | ruff format | `pr.yml:lint` | — |
-| `api.quality.ruff-lint` | All Python code passes ruff check with every rule family selected in pyproject.toml, with no ratchet; families are added one themed group at a time until the selection is ALL. | mechanical | enforced | yes | ruff check (E4, E7, E9, F, FLY, FURB, I, PIE, PYI, RET, RSE, TC, UP, W) | `pr.yml:lint` | — |
+| `api.quality.ruff-lint` | All Python code passes ruff check with every rule family selected in pyproject.toml, with no ratchet; families are added one themed group at a time until the selection is ALL. | mechanical | enforced | yes | ruff check (ANN, E4, E7, E9, F, FLY, FURB, I, PIE, PYI, RET, RSE, TC, UP, W) | `pr.yml:lint` | — |
 | `api.quality.tests-pass` | The unit, contract and to-refactor test suites pass. | mechanical | enforced | yes | CI job `test` | `pr.yml:test` | — |
+| `api.quality.tests-ruff-exemptions` | Test files may use pytest's bare assert and need no docstrings, annotations or named constants for magic numbers (S101, D1, ANN, PLR2004); every other ruff rule applies to tests. | mechanical | enforced | yes | ruff per-file ignore of ANN, D1, PLR2004, S101 for `**/tests/**` | `pr.yml:lint` | — |
 
 ```mermaid
 flowchart LR
@@ -286,10 +287,12 @@ flowchart LR
   e_ruff_format["ruff format"]
   j_CI_pr_yml_lint["CI pr.yml:lint"]
   r_api_quality_ruff_lint("api.quality.ruff-lint (enforced)")
-  e_ruff_check__E4__E7__E9__F__FLY__FURB__I__PIE__PYI__RET__RSE__TC__UP__W_["ruff check (E4, E7, E9, F, FLY, FURB, I, PIE, PYI, RET, RSE, TC, UP, W)"]
+  e_ruff_check__ANN__E4__E7__E9__F__FLY__FURB__I__PIE__PYI__RET__RSE__TC__UP__W_["ruff check (ANN, E4, E7, E9, F, FLY, FURB, I, PIE, PYI, RET, RSE, TC, UP, W)"]
   r_api_quality_tests_pass("api.quality.tests-pass (enforced)")
   e_CI_job_test["CI job test"]
   j_CI_pr_yml_test["CI pr.yml:test"]
+  r_api_quality_tests_ruff_exemptions("api.quality.tests-ruff-exemptions (enforced)")
+  e_ruff_per_file_ignore_of_ANN__D1__PLR2004__S101_for____tests___["ruff per-file ignore of ANN, D1, PLR2004, S101 for **/tests/**"]
   r_api_quality_alembic_revisions_unused_imports --> e_ruff_per_file_ignore_of_ALL_for_migrations___versions___py
   e_ruff_per_file_ignore_of_ALL_for_migrations___versions___py -.-> local_only
   r_api_quality_environment_variables_load --> e_CI_job_test_env_vars
@@ -298,10 +301,12 @@ flowchart LR
   e_mypy --> j_CI_pr_yml_typecheck_v2
   r_api_quality_ruff_format --> e_ruff_format
   e_ruff_format --> j_CI_pr_yml_lint
-  r_api_quality_ruff_lint --> e_ruff_check__E4__E7__E9__F__FLY__FURB__I__PIE__PYI__RET__RSE__TC__UP__W_
-  e_ruff_check__E4__E7__E9__F__FLY__FURB__I__PIE__PYI__RET__RSE__TC__UP__W_ --> j_CI_pr_yml_lint
+  r_api_quality_ruff_lint --> e_ruff_check__ANN__E4__E7__E9__F__FLY__FURB__I__PIE__PYI__RET__RSE__TC__UP__W_
+  e_ruff_check__ANN__E4__E7__E9__F__FLY__FURB__I__PIE__PYI__RET__RSE__TC__UP__W_ --> j_CI_pr_yml_lint
   r_api_quality_tests_pass --> e_CI_job_test
   e_CI_job_test --> j_CI_pr_yml_test
+  r_api_quality_tests_ruff_exemptions --> e_ruff_per_file_ignore_of_ANN__D1__PLR2004__S101_for____tests___
+  e_ruff_per_file_ignore_of_ANN__D1__PLR2004__S101_for____tests___ --> j_CI_pr_yml_lint
 ```
 
 ## Observability

migrations/v1/env.py

@@ -1,14 +1,20 @@
 """Alembic environment for the API v1 schema."""
 
+from __future__ import annotations
+
 import os
 from logging.config import fileConfig
+from typing import TYPE_CHECKING
 
 from alembic import context
 from sqlalchemy import engine_from_config, pool
 from sqlalchemy.engine import make_url
 
 from policyengine_api.data.v1_models import V1Base
 
+if TYPE_CHECKING:
+    from sqlalchemy import Connection
+
 config = context.config
 if config.config_file_name is not None:
     fileConfig(config.config_file_name)
@@ -59,7 +65,7 @@ def run_migrations_online() -> None:
         _run_migrations_with_connection(connection)
 
 
-def _run_migrations_with_connection(connection) -> None:
+def _run_migrations_with_connection(connection: Connection) -> None:
     context.configure(
         connection=connection,
         target_metadata=target_metadata,

migrations/v2/env.py

@@ -1,6 +1,9 @@
 """Alembic environment for the isolated API v2-alpha Postgres schema."""
 
+from __future__ import annotations
+
 from logging.config import fileConfig
+from typing import TYPE_CHECKING
 
 from alembic import context
 from alembic.script import ScriptDirectory
@@ -13,6 +16,11 @@
 )
 from policyengine_api.data.v2.models import V2_METADATA
 
+if TYPE_CHECKING:
+    from alembic.runtime.environment import NameFilterType
+    from sqlalchemy import Connection
+    from sqlalchemy.sql.schema import SchemaItem
+
 config = context.config
 if config.config_file_name is not None:
     fileConfig(config.config_file_name)
@@ -22,13 +30,19 @@
 script = ScriptDirectory.from_config(config)
 
 
-def _include_application_object(_object, name, type_, _reflected, _compare_to) -> bool:
+def _include_application_object(
+    _object: SchemaItem,
+    name: str | None,
+    type_: NameFilterType,
+    _reflected: bool,
+    _compare_to: SchemaItem | None,
+) -> bool:
     """Keep Alembic's own version table outside application drift."""
 
     return not (type_ == "table" and name == "alembic_version")
 
 
-def _configure(connection) -> None:
+def _configure(connection: Connection) -> None:
     qualify_v2_connection(connection, settings)
     context.configure(
         connection=connection,

policyengine_api/api.py

@@ -9,7 +9,7 @@
 start_time = time.time()
 
 
-def log_timing(message):
+def log_timing(message: str) -> None:
     elapsed = time.time() - start_time
     print(f"[{elapsed:.2f}s] {message}", file=sys.stderr, flush=True)
 

policyengine_api/asgi_factory.py

@@ -4,7 +4,7 @@
 
 import time
 from contextlib import asynccontextmanager
-from typing import TYPE_CHECKING
+from typing import TYPE_CHECKING, Any
 
 from a2wsgi import WSGIMiddleware
 from fastapi import FastAPI, Request
@@ -51,10 +51,11 @@
 )
 
 if TYPE_CHECKING:
-    from collections.abc import Callable
+    from collections.abc import AsyncIterator, Awaitable, Callable, Iterable
 
     from policyengine_observability import ObservabilityRuntime
-    from starlette.types import ASGIApp
+    from starlette.routing import BaseRoute
+    from starlette.types import ASGIApp, Scope
 
 
 def _apply_request_id_header(
@@ -71,7 +72,7 @@ def _apply_observability_id_header(
     response.headers[OBSERVABILITY_ID_HEADER] = observability_id
 
 
-def _matched_route_template(routes, scope: dict) -> str | None:
+def _matched_route_template(routes: Iterable[BaseRoute], scope: Scope) -> str | None:
     for route in routes:
         match, _ = route.matches(scope)
         if match is Match.FULL:
@@ -93,14 +94,14 @@ def _matched_route_template(routes, scope: dict) -> str | None:
     return None
 
 
-def _native_route_template(app: FastAPI, scope: dict) -> str | None:
+def _native_route_template(app: FastAPI, scope: Scope) -> str | None:
     """Return the matched FastAPI route template, excluding the Flask mount."""
 
     return _matched_route_template(app.router.routes, scope)
 
 
 def create_asgi_app(
-    wsgi_app,
+    wsgi_app: Any,  # noqa: ANN401  # Flask's WSGI signature and a2wsgi's WSGIApp stubs disagree
     *,
     route_settings: RouteImplementationSettings | None = None,
     dependencies: NativeRouteDependencies | None = None,
@@ -116,7 +117,7 @@ def create_asgi_app(
     request_runtime = observability_runtime or get_runtime()
 
     @asynccontextmanager
-    async def lifespan(_app: FastAPI):
+    async def lifespan(_app: FastAPI) -> AsyncIterator[None]:
         try:
             yield
         finally:
@@ -172,7 +173,10 @@ async def oversized_v2_request(
         return v2_error_response(413, str(error))
 
     @app.middleware("http")
-    async def add_request_context_and_migration_logging(request, call_next):
+    async def add_request_context_and_migration_logging(
+        request: Request,
+        call_next: Callable[[Request], Awaitable[Response]],
+    ) -> Response:
         started_at = time.time()
         request_id = request.headers.get(REQUEST_ID_HEADER) or generate_request_id()
         incoming_observability_id = normalize_observability_id(

policyengine_api/constants.py

@@ -3,6 +3,7 @@
 from datetime import datetime
 from importlib.metadata import distribution, distributions
 from pathlib import Path
+from typing import Any
 
 from policyengine_api.dataset_display import (
     DEFAULT_DATASET_DISPLAY_LABEL,
@@ -47,7 +48,7 @@ def _resolve_distribution_version(
     return "0.0.0"
 
 
-def get_py_manifest():
+def get_py_manifest() -> Any:  # noqa: ANN401  # SimplePath, which open() rejects; a Path at runtime
     return distribution("policyengine").locate_file(
         "policyengine/data/bundle/manifest.json"
     )