chore(typecheck): include top-level tools

pyproject.toml

@@ -373,6 +373,4 @@ exclude = [
     ".uv_cache", # Cache Dir in CI
     "./docs/**/*.ipynb",
     "./src/nemo_safe_synthesizer/pii_replacer/",
-    "./tools/*.py",
-    "!./tools/gen_cuda_deps.py",
 ]

tests/tools/test_diff_lockfile.py

@@ -0,0 +1,56 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+from __future__ import annotations
+
+import importlib.util
+import sys
+from pathlib import Path
+from typing import Protocol, cast
+
+import pytest
+
+
+class DiffLockfileTool(Protocol):
+    """Typed interface exposed by the dynamically loaded lockfile tool."""
+
+    def parse_packages(self, content: str) -> dict[str, object]: ...
+
+
+def _load_tool(tool_path: Path) -> DiffLockfileTool:
+    spec = importlib.util.spec_from_file_location("diff_lockfile", tool_path)
+    if spec is None or spec.loader is None:
+        raise RuntimeError(f"Could not load {tool_path}")
+    module = importlib.util.module_from_spec(spec)
+    sys.modules[spec.name] = module
+    spec.loader.exec_module(module)
+    return cast(DiffLockfileTool, module)
+
+
+@pytest.fixture(scope="module")
+def diff_lockfile_tool(pytestconfig: pytest.Config) -> DiffLockfileTool:
+    """Load the lockfile helper from pytest's discovered repository root."""
+    return _load_tool(Path(pytestconfig.rootpath) / "tools" / "diff-lockfile.py")
+
+
+def test_parse_packages_preserves_same_source_versions(diff_lockfile_tool: DiffLockfileTool) -> None:
+    content = """
+version = 1
+
+[[package]]
+name = "cuda-python"
+version = "12.9.4"
+source = { registry = "https://pypi.org/simple" }
+
+[[package]]
+name = "cuda-python"
+version = "13.1.1"
+source = { registry = "https://pypi.org/simple" }
+"""
+
+    packages = diff_lockfile_tool.parse_packages(content)
+
+    assert set(packages) == {
+        "cuda-python@https://pypi.org/simple@12.9.4",
+        "cuda-python@https://pypi.org/simple@13.1.1",
+    }

tools/diff-lockfile.py

@@ -36,7 +36,7 @@
 import subprocess
 import tomllib
 from enum import Enum
-from typing import Annotated, Optional
+from typing import Annotated, Any, Optional
 
 import git
 import typer
@@ -116,12 +116,17 @@ def get_lockfile_content(repo: git.Repo, ref: str, path: str) -> str:
     return blob.data_stream.read().decode()
 
 
-def _extract_source(raw: dict) -> str:
+def _extract_source(raw: dict[str, Any]) -> str:
     """Return a human-readable source string from a ``[[package]]`` entry."""
-    src = raw.get("source", {})
-    if isinstance(src, dict):
-        return src.get("registry", src.get("git", src.get("path", "")))
-    return str(src) if src else ""
+    match raw.get("source", {}):
+        case {"registry": value} | {"git": value} | {"path": value}:
+            return str(value) if value is not None else ""
+        case dict():
+            return ""
+        case src if src:
+            return str(src)
+        case _:
+            return ""
 
 
 def parse_packages(content: str) -> dict[str, Package]:
@@ -131,9 +136,8 @@ def parse_packages(content: str) -> dict[str, Package]:
         content: Raw UTF-8 text of the lockfile.
 
     Returns:
-        A dict keyed by package name (or ``name@source`` when a name
-        appears more than once with different sources).  Values are
-        :class:`Package` instances.
+        A dict keyed by package name (or ``name@source@version`` when a name
+        appears more than once). Values are :class:`Package` instances.
     """
     data = tomllib.loads(content)
 
@@ -151,8 +155,10 @@ def parse_packages(content: str) -> dict[str, Package]:
             continue
         source = _extract_source(pkg)
 
-        # Disambiguate packages that appear under multiple sources
-        key = f"{name}@{source}" if name_counts.get(name, 1) > 1 else name
+        # uv.lock may contain several versions of a package from one source.
+        # Keep every variant instead of letting later entries overwrite earlier
+        # entries in this mapping.
+        key = f"{name}@{source}@{version}" if name_counts.get(name, 1) > 1 else name
         packages[key] = Package(name=name, version=Version(version), source=source)
 
     return packages

tools/hf_network_guard_proxy.py

@@ -8,6 +8,7 @@
 #     "pydantic",
 #     "proxy.py>=2.4.10",
 #     "structlog",
+#     "typing-extensions>=4.15.0",
 # ]
 # ///
 # pyright: reportMissingImports=false
@@ -74,7 +75,7 @@
 import sys
 import threading
 import time
-from contextlib import nullcontext
+from contextlib import AbstractContextManager, nullcontext
 from dataclasses import dataclass, field
 from enum import IntEnum
 from multiprocessing import Manager
@@ -90,6 +91,7 @@
 from proxy.http.parser import HttpParser
 from proxy.http.proxy import HttpProxyBasePlugin
 from pydantic import BaseModel, ConfigDict, Field
+from typing_extensions import override
 
 DEFAULT_HOST = "127.0.0.1"
 DEFAULT_PORT = 8765
@@ -170,7 +172,7 @@ def summary(self) -> ProxySummary:
 class RunningServer:
     """Background proxy.py server instance."""
 
-    server: Proxy
+    server: AbstractContextManager[object]
     host: str
     port: int
     state: ProxyState
@@ -201,6 +203,7 @@ def configure(cls, *, state: ProxyState, allow_passthrough_hf: bool) -> None:
         cls.state = state
         cls.allow_passthrough_hf = allow_passthrough_hf
 
+    @override
     def before_upstream_connection(self, request: HttpParser) -> HttpParser | None:
         """Record the destination and decide whether proxy.py may connect upstream."""
         proxy_request = _proxy_request_from_proxy_py(request)