Merge pull request #3907 from PolicyEngine/ruff-docstrings

Enable ruff docstring rules (D)

Author: anth-volk

.github/bump_version.py

@@ -6,6 +6,7 @@
 
 
 def get_current_version(pyproject_path: Path) -> str:
+    """Return the version in pyproject.toml, exiting with status 1 if it is absent."""
     text = pyproject_path.read_text()
     match = re.search(r'^version\s*=\s*"(\d+\.\d+\.\d+)"', text, re.MULTILINE)
     if not match:
@@ -18,6 +19,11 @@ def get_current_version(pyproject_path: Path) -> str:
 
 
 def infer_bump(changelog_dir: Path) -> str:
+    """Return "major", "minor" or "patch" from the pending changelog fragment types.
+
+    Any breaking fragment makes a major bump and any added or removed fragment a
+    minor one; exit with status 1 when there are no fragments.
+    """
     fragments = [
         f for f in changelog_dir.iterdir() if f.is_file() and f.name != ".gitkeep"
     ]
@@ -39,6 +45,7 @@ def infer_bump(changelog_dir: Path) -> str:
 
 
 def bump_version(version: str, bump: str) -> str:
+    """Return ``version`` raised by one ``bump`` level, resetting the lower parts."""
     major, minor, patch = (int(x) for x in version.split("."))
     if bump == "major":
         return f"{major + 1}.0.0"
@@ -48,6 +55,7 @@ def bump_version(version: str, bump: str) -> str:
 
 
 def update_file(path: Path, old_version: str, new_version: str) -> None:
+    """Rewrite the ``version = "..."`` line in ``path``; leave other files untouched."""
     text = path.read_text()
     updated = text.replace(
         f'version = "{old_version}"',
@@ -59,6 +67,7 @@ def update_file(path: Path, old_version: str, new_version: str) -> None:
 
 
 def main() -> None:
+    """Bump the version in pyproject.toml according to the changelog.d fragments."""
     root = Path(__file__).resolve().parent.parent
     pyproject = root / "pyproject.toml"
     changelog_dir = root / "changelog.d"

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

@@ -1,3 +1,9 @@
+"""CI helper that reads the versions pinned by the installed PolicyEngine bundle.
+
+It checks the package and data versions against the API's version constants,
+then writes them to GITHUB_ENV or, with --shell, prints shell assignments.
+"""
+
 import argparse
 import os
 import shlex
@@ -48,9 +54,7 @@ def _data_release_version(bundle: dict, country_id: str) -> str:
 
 
 def find_api_model_versions() -> dict[str, str]:
-    """
-    Find and validate the package and data versions in the installed bundle.
-    """
+    """Find and validate the package and data versions in the installed bundle."""
     bundle = get_current_bundle()
     if not isinstance(bundle, dict):
         raise RuntimeError("Installed PolicyEngine bundle manifest is not an object.")
@@ -89,9 +93,7 @@ def find_api_model_versions() -> dict[str, str]:
 
 
 def find_api_model_versions_and_output_to_github() -> None:
-    """
-    Find the API model versions and output them to a file for GitHub.
-    """
+    """Find the API model versions and output them to a file for GitHub."""
     versions = find_api_model_versions()
     with open(os.environ["GITHUB_ENV"], "a") as f:
         for key, value in versions.items():

.github/scripts/check_updates.py

@@ -1,6 +1,5 @@
 #!/usr/bin/env python3
-"""
-Fetch and format changelog entries between two versions of a package.
+"""Fetch and format changelog entries between two versions of a package.
 
 Fetches CHANGELOG.md from the package's GitHub repo and extracts
 entries between the old and new version, formatted as markdown.
@@ -25,6 +24,8 @@
 
 
 class ChangelogEntry(TypedDict):
+    """One released version's changelog items, grouped by lowercase category."""
+
     version: str
     changes: dict[str, list[str]]
 
@@ -119,6 +120,10 @@ def format_changes(entries: list[ChangelogEntry]) -> str:
 
 
 def main() -> int:
+    """Print the markdown changes between two versions of a package.
+
+    Always return 0, so an unreachable changelog never fails the update job.
+    """
     parser = argparse.ArgumentParser(
         description="Fetch changelog entries between two package versions."
     )

.github/scripts/verify_observability_collector.py

@@ -64,6 +64,8 @@
 
 
 class VerificationProbe(NamedTuple):
+    """Identifiers of one verification export, for finding its stored telemetry."""
+
     verification_id: str
     trace_id: str
     started_at: float
@@ -315,7 +317,6 @@ def verify(
     delivery_timeout_seconds: float = 600.0,
 ) -> VerificationProbe:
     """Export real telemetry, verify storage, and require log rejection."""
-
     started_at = time.time()
     verification_id = f"collector-verify-{uuid.uuid4().hex}"
     trace_request, metric_request, log_request, trace_id = _probe_requests(
@@ -370,6 +371,11 @@ def verify(
 
 
 def main() -> None:
+    """Verify the collector, reading both tokens from the environment.
+
+    COLLECTOR_ID_TOKEN authenticates to the collector and GOOGLE_OAUTH_ACCESS_TOKEN
+    reads Cloud Trace and Cloud Monitoring.
+    """
     parser = argparse.ArgumentParser()
     parser.add_argument("--endpoint", required=True)
     parser.add_argument("--project-id", required=True)

AGENTS.md

@@ -123,7 +123,7 @@ 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, v2 routes and services, and the user sign-in and account-deletion libraries 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 (ANN, 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, D, E4, E7, E9, F, FLY, FURB, I, PIE, PYI, RET, RSE, TC, UP, W))
 - `api.quality.runtime-evaluated-models-resolve`: Every package model is one ruff treats as runtime-evaluated (a base resolves, through the import in its own module, to a name in runtime-evaluated-base-classes, or is a same-module class that does), and every Pydantic model and SQLAlchemy mapper resolves its annotations once the package is imported. (blocking: pytest `tests/unit/test_runtime_evaluated_annotations.py`)
 - `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/**`)

changelog.d/ruff-docstrings.changed.md

@@ -0,0 +1 @@
+Enabled ruff's docstring rules (D) and documented every module, class and function outside tests; published models and routes now carry these descriptions in their JSON schema and OpenAPI, with no structural schema change.

dashboard/experiments/gpt4_api_user.py

@@ -1,3 +1,10 @@
+"""Streamlit experiment that answers a UK policy question with GPT-4 and the API.
+
+It finds related parameters by embedding similarity, asks GPT-4 to write the
+reform as JSON, runs it on the public PolicyEngine API, and asks GPT-4 to answer
+from the resulting economic impact.
+"""
+
 import json
 import os
 import time
@@ -14,6 +21,7 @@
 
 @st.cache_data
 def get_embeddings_df() -> pd.DataFrame:
+    """Load parameter_embeddings.csv.gz, evaluating each embedding string to a list."""
     embeddings_df = pd.read_csv("parameter_embeddings.csv.gz", compression="gzip")
     embeddings_df.parameter_embedding = embeddings_df.parameter_embedding.apply(
         lambda x: eval(x)
@@ -23,6 +31,7 @@ def get_embeddings_df() -> pd.DataFrame:
 
 @st.cache_data
 def embed(prompt: str, engine: str = "text-embedding-ada-002") -> list[float]:
+    """Return the OpenAI embedding of ``prompt``, cached across Streamlit reruns."""
     return get_embedding(prompt, engine=engine)
 
 
@@ -59,6 +68,7 @@ def embed(prompt: str, engine: str = "text-embedding-ada-002") -> list[float]:
 
 
 def write_prompt(question: str, relevant_parameters: object) -> str:
+    """Build the GPT-4 prompt that rewrites ``question`` as a reform JSON object."""
     return f"""Question: {question}
     Here's some metadata about the parameters that are relevant to your question:
     Relevant parameters: {relevant_parameters}
@@ -79,6 +89,7 @@ def write_prompt(question: str, relevant_parameters: object) -> str:
 
 @st.cache_data
 def get_gpt4_response(prompt: str) -> str:
+    """Return GPT-4's reply to ``prompt`` as one user message, cached across reruns."""
     return openai.ChatCompletion.create(
         model="gpt-4",
         messages=[{"role": "user", "content": prompt}],
@@ -94,6 +105,10 @@ def get_gpt4_response(prompt: str) -> str:
 
 @st.cache_data
 def get_policyengine_impact(reform: object) -> object:
+    """Create ``reform`` on the public UK API and return its 2023 UK economic impact.
+
+    Poll every five seconds, with no time limit, while the impact is computing.
+    """
     payload = {"data": reform}
     url = "https://api.policyengine.org/uk/policy"
 

docs/generated/rules.md

@@ -272,7 +272,7 @@ 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, v2 routes and services, and the user sign-in and account-deletion libraries 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 (ANN, 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, D, E4, E7, E9, F, FLY, FURB, I, PIE, PYI, RET, RSE, TC, UP, W) | `pr.yml:lint` | — |
 | `api.quality.runtime-evaluated-models-resolve` | Every package model is one ruff treats as runtime-evaluated (a base resolves, through the import in its own module, to a name in runtime-evaluated-base-classes, or is a same-module class that does), and every Pydantic model and SQLAlchemy mapper resolves its annotations once the package is imported. | mechanical | enforced | yes | pytest `tests/unit/test_runtime_evaluated_annotations.py` | `pr.yml:test` | **unproven** |
 | `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` | — |
@@ -292,7 +292,7 @@ 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__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)"]
+  e_ruff_check__ANN__D__E4__E7__E9__F__FLY__FURB__I__PIE__PYI__RET__RSE__TC__UP__W_["ruff check (ANN, D, E4, E7, E9, F, FLY, FURB, I, PIE, PYI, RET, RSE, TC, UP, W)"]
   r_api_quality_runtime_evaluated_models_resolve("api.quality.runtime-evaluated-models-resolve (enforced)")
   e_pytest_tests_unit_test_runtime_evaluated_annotations_py["pytest tests/unit/test_runtime_evaluated_annotations.py"]
   j_CI_pr_yml_test["CI pr.yml:test"]
@@ -308,8 +308,8 @@ 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__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_ruff_lint --> e_ruff_check__ANN__D__E4__E7__E9__F__FLY__FURB__I__PIE__PYI__RET__RSE__TC__UP__W_
+  e_ruff_check__ANN__D__E4__E7__E9__F__FLY__FURB__I__PIE__PYI__RET__RSE__TC__UP__W_ --> j_CI_pr_yml_lint
   r_api_quality_runtime_evaluated_models_resolve --> e_pytest_tests_unit_test_runtime_evaluated_annotations_py
   e_pytest_tests_unit_test_runtime_evaluated_annotations_py --> j_CI_pr_yml_test
   r_api_quality_tests_pass --> e_CI_job_test

migrations/v1/env.py

@@ -40,6 +40,7 @@
 
 
 def run_migrations_offline() -> None:
+    """Emit the migration SQL for the configured MySQL URL without connecting."""
     context.configure(
         url=config.get_main_option("sqlalchemy.url"),
         target_metadata=target_metadata,
@@ -52,6 +53,7 @@ def run_migrations_offline() -> None:
 
 
 def run_migrations_online() -> None:
+    """Migrate the connection the caller supplied, or else a new configured engine."""
     if provided_connection is not None:
         _run_migrations_with_connection(provided_connection)
         return

migrations/v2/env.py

@@ -38,7 +38,6 @@ def _include_application_object(
     _compare_to: SchemaItem | None,
 ) -> bool:
     """Keep Alembic's own version table outside application drift."""
-
     return not (type_ == "table" and name == "alembic_version")
 
 
@@ -65,13 +64,19 @@ def _configure(connection: Connection) -> None:
 
 
 def run_migrations_offline() -> None:
+    """Raise, because v2 migrations need a live connection to qualify the target."""
     raise RuntimeError(
         "v2 migrations require an online connection so target identity and "
         "generated application-data before/after states can be verified"
     )
 
 
 def run_migrations_online() -> None:
+    """Migrate the caller's connection, or else a new engine in one transaction.
+
+    The new engine's transaction commits, so the target qualification's reads
+    cannot leave the completed migration to be rolled back.
+    """
     provided_connection = config.attributes.get("connection")
     if provided_connection is not None:
         _configure(provided_connection)