ci(branch-checks): align Python checks with pre-commit (NVIDIA#1908)

Signed-off-by: Taylor Mutch <taylormutch@gmail.com>

Author: TaylorMutch

.github/workflows/branch-checks.yml

@@ -160,9 +160,15 @@ jobs:
       - name: Install dependencies
         run: uv sync --frozen
 
+      - name: Format
+        run: mise run python:format:check
+
       - name: Lint
         run: mise run python:lint
 
+      - name: Typecheck
+        run: mise run python:typecheck
+
       - name: Test
         run: mise run test:python
 

pyproject.toml

@@ -54,6 +54,8 @@ dev = [
     "setuptools-scm>=8",
     "grpcio-tools>=1.60",
     "pyelftools>=0.30",
+    "pyyaml==6.0.2",
+    "types-pyyaml>=6.0.12.20260518",
 ]
 
 [tool.uv]

tasks/scripts/sync_docs_website.py

@@ -16,13 +16,14 @@
 import shutil
 import sys
 import tempfile
-from collections.abc import MutableMapping, MutableSequence
 from dataclasses import dataclass
 from pathlib import Path
+from typing import cast
 
 import yaml
 
 SLUG_RE = re.compile(r"^[A-Za-z0-9._-]+$")
+YamlMapping = dict[str, object]
 
 
 @dataclass
@@ -39,7 +40,9 @@ def parse_args() -> argparse.Namespace:
     parser.add_argument("--operation", choices=["sync", "remove"], default="sync")
     parser.add_argument("--source-root", type=Path)
     parser.add_argument("--docs-website-root", required=True, type=Path)
-    parser.add_argument("--channel", required=True, choices=["dev", "latest", "version"])
+    parser.add_argument(
+        "--channel", required=True, choices=["dev", "latest", "version"]
+    )
     parser.add_argument("--source-ref", default="")
     parser.add_argument("--version-slug", default="")
     parser.add_argument("--display-name", default="")
@@ -58,11 +61,15 @@ def resolve_slug(channel: str, version_slug: str) -> str:
     if not version_slug:
         raise ValueError("--version-slug is required when --channel=version")
     if not SLUG_RE.fullmatch(version_slug):
-        raise ValueError(f"version slug contains unsupported characters: {version_slug}")
+        raise ValueError(
+            f"version slug contains unsupported characters: {version_slug}"
+        )
     return version_slug
 
 
-def resolve_display_name(channel: str, slug: str, source_ref: str, override: str) -> str:
+def resolve_display_name(
+    channel: str, slug: str, source_ref: str, override: str
+) -> str:
     if override:
         return override
     if channel == "dev":
@@ -116,15 +123,15 @@ def copy_if_exists(src: Path, dst: Path) -> None:
         shutil.copy2(src, dst)
 
 
-def read_yaml(path: Path) -> dict:
+def read_yaml(path: Path) -> YamlMapping:
     ensure_existing(path, "YAML file")
     data = yaml.safe_load(path.read_text(encoding="utf-8"))
     if not isinstance(data, dict):
         raise ValueError(f"expected YAML mapping in {path}")
-    return data
+    return cast("YamlMapping", data)
 
 
-def write_yaml(path: Path, data: dict) -> None:
+def write_yaml(path: Path, data: YamlMapping) -> None:
     path.write_text(
         yaml.safe_dump(data, sort_keys=False, allow_unicode=True),
         encoding="utf-8",
@@ -140,20 +147,23 @@ def prefix_path(value: object, pages_dir: str) -> object:
 
 
 def prefix_navigation_paths(value: object, pages_dir: str) -> object:
-    if isinstance(value, MutableMapping):
+    if isinstance(value, dict):
+        mapping = cast("YamlMapping", value)
         for key in ("path", "folder"):
-            if key in value:
-                value[key] = prefix_path(value[key], pages_dir)
-        for child in value.values():
+            if key in mapping:
+                mapping[key] = prefix_path(mapping[key], pages_dir)
+        for child in mapping.values():
             prefix_navigation_paths(child, pages_dir)
-    elif isinstance(value, MutableSequence):
-        for child in value:
+    elif isinstance(value, list):
+        for child in cast("list[object]", value):
             prefix_navigation_paths(child, pages_dir)
     return value
 
 
-def version_navigation(source_index: Path, pages_dir: str) -> dict:
-    return prefix_navigation_paths(read_yaml(source_index), pages_dir)
+def version_navigation(source_index: Path, pages_dir: str) -> YamlMapping:
+    data = read_yaml(source_index)
+    prefix_navigation_paths(data, pages_dir)
+    return data
 
 
 def parse_versions(raw_versions: object) -> list[VersionEntry]:
@@ -162,20 +172,27 @@ def parse_versions(raw_versions: object) -> list[VersionEntry]:
     if not isinstance(raw_versions, list):
         raise ValueError("docs.yml versions must be a list")
     entries: list[VersionEntry] = []
-    for raw in raw_versions:
+    for raw in cast("list[object]", raw_versions):
         if not isinstance(raw, dict):
             continue
-        slug = raw.get("slug")
-        display_name = raw.get("display-name")
-        path = raw.get("path")
-        if isinstance(slug, str) and isinstance(display_name, str) and isinstance(path, str):
+        entry = cast("YamlMapping", raw)
+        slug = entry.get("slug")
+        display_name = entry.get("display-name")
+        path = entry.get("path")
+        if (
+            isinstance(slug, str)
+            and isinstance(display_name, str)
+            and isinstance(path, str)
+        ):
             entries.append(
                 VersionEntry(slug=slug, display_name=display_name, path=path)
             )
     return entries
 
 
-def ordered_entries(existing: list[VersionEntry], updated: VersionEntry) -> list[VersionEntry]:
+def ordered_entries(
+    existing: list[VersionEntry], updated: VersionEntry
+) -> list[VersionEntry]:
     by_slug = {entry.slug: entry for entry in existing}
     by_slug[updated.slug] = updated
     existing_order = [entry.slug for entry in existing if entry.slug != updated.slug]
@@ -206,7 +223,9 @@ def render_versions(entries: list[VersionEntry]) -> list[dict[str, str]]:
 def component_dirs(fern_dir: Path) -> list[str]:
     dirs: list[str] = []
     preferred = ["pages-latest", "pages-dev"]
-    all_page_dirs = sorted(path.name for path in fern_dir.glob("pages-*") if path.is_dir())
+    all_page_dirs = sorted(
+        path.name for path in fern_dir.glob("pages-*") if path.is_dir()
+    )
     for name in preferred + all_page_dirs:
         path = fern_dir / name / "_components"
         component = f"./{name}/_components"
@@ -229,7 +248,9 @@ def update_docs_yml(docs_yml: Path, updated: VersionEntry, fern_dir: Path) -> No
 
 def remove_docs_yml_entry(docs_yml: Path, slug: str, fern_dir: Path) -> None:
     data = read_yaml(docs_yml)
-    entries = [entry for entry in parse_versions(data.get("versions")) if entry.slug != slug]
+    entries = [
+        entry for entry in parse_versions(data.get("versions")) if entry.slug != slug
+    ]
     data["experimental"] = {
         "mdx-components": component_dirs(fern_dir),
     }
@@ -274,7 +295,9 @@ def sync_docs(args: argparse.Namespace) -> None:
     )
     if refresh_shared:
         copy_if_exists(source_fern / "main.css", target_fern / "main.css")
-        copy_if_exists(source_fern / "fern.config.json", target_fern / "fern.config.json")
+        copy_if_exists(
+            source_fern / "fern.config.json", target_fern / "fern.config.json"
+        )
 
     versions_dir = target_fern / "versions"
     versions_dir.mkdir(parents=True, exist_ok=True)
@@ -294,8 +317,7 @@ def sync_docs(args: argparse.Namespace) -> None:
     )
 
     print(
-        f"Synced {channel} docs from {source_ref} to fern/{pages_dir} "
-        f"({display_name})"
+        f"Synced {channel} docs from {source_ref} to fern/{pages_dir} ({display_name})"
     )
 
 

tasks/scripts/sync_docs_website_test.py

@@ -11,7 +11,7 @@
 from __future__ import annotations
 
 from argparse import Namespace
-from typing import TYPE_CHECKING
+from typing import TYPE_CHECKING, cast
 
 import pytest
 import sync_docs_website as sdw
@@ -66,7 +66,7 @@ def test_ordered_entries_pins_latest_then_dev() -> None:
 
 
 def test_prefix_navigation_paths() -> None:
-    nav = {
+    nav: dict[str, object] = {
         "navigation": [
             {"page": "Intro", "path": "intro.mdx"},
             {
@@ -78,11 +78,14 @@ def test_prefix_navigation_paths() -> None:
         ]
     }
     sdw.prefix_navigation_paths(nav, "pages-dev")
-    assert nav["navigation"][0]["path"] == "../pages-dev/intro.mdx"
-    assert nav["navigation"][1]["folder"] == "../pages-dev/guide"
-    assert nav["navigation"][1]["contents"][0]["path"] == "../pages-dev/guide/a.mdx"
+    navigation = cast("list[dict[str, object]]", nav["navigation"])
+    guide = navigation[1]
+    contents = cast("list[dict[str, object]]", guide["contents"])
+    assert navigation[0]["path"] == "../pages-dev/intro.mdx"
+    assert guide["folder"] == "../pages-dev/guide"
+    assert contents[0]["path"] == "../pages-dev/guide/a.mdx"
     # Absolute URLs are left untouched.
-    assert nav["navigation"][2]["path"] == "https://example.com"
+    assert navigation[2]["path"] == "https://example.com"
 
 
 def _make_source_tree(root: Path) -> None: