Share encoded CUDA version decoding logic.

Centralize encoded CUDA version parsing and validation so toolkit and driver version helpers stay aligned and cuda.h parsing gets consistent string conversion and error reporting.

Made-with: Cursor

Author: rwgk

cuda_pathfinder/cuda/pathfinder/_utils/driver_info.py

@@ -7,19 +7,21 @@
 import functools
 from collections.abc import Callable
 from dataclasses import dataclass
+from typing import cast
 
 from cuda.pathfinder._dynamic_libs.load_nvidia_dynamic_lib import (
     load_nvidia_dynamic_lib as _load_nvidia_dynamic_lib,
 )
 from cuda.pathfinder._utils.platform_aware import IS_WINDOWS
+from cuda.pathfinder._utils.toolkit_info import EncodedCudaVersion
 
 
 class QueryDriverCudaVersionError(RuntimeError):
     """Raised when ``query_driver_cuda_version()`` cannot determine the CUDA driver version."""
 
 
 @dataclass(frozen=True, slots=True)
-class DriverCudaVersion:
+class DriverCudaVersion(EncodedCudaVersion):
     """
     CUDA-facing driver version reported by ``cuDriverGetVersion()``.
 
@@ -41,21 +43,13 @@ class DriverCudaVersion:
     to ``Driver Version: 595.58.03``.
     """
 
-    encoded: int
-    major: int
-    minor: int
-
 
 @functools.cache
 def query_driver_cuda_version() -> DriverCudaVersion:
     """Return the CUDA driver version parsed into its major/minor components."""
     try:
         encoded = _query_driver_cuda_version_int()
-        return DriverCudaVersion(
-            encoded=encoded,
-            major=encoded // 1000,
-            minor=(encoded % 1000) // 10,
-        )
+        return cast(DriverCudaVersion, DriverCudaVersion.from_encoded(encoded))
     except Exception as exc:
         raise QueryDriverCudaVersionError("Failed to query the CUDA driver version.") from exc
 

cuda_pathfinder/cuda/pathfinder/_utils/toolkit_info.py

@@ -7,34 +7,63 @@
 import re
 from dataclasses import dataclass
 from pathlib import Path
+from typing import TypeVar
 
 _CUDA_VERSION_RE = re.compile(r"^\s*#\s*define\s+CUDA_VERSION\s+(?P<encoded>\d+)\b", re.MULTILINE)
+EncodedCudaVersionT = TypeVar("EncodedCudaVersionT", bound="EncodedCudaVersion")
+
+
+@dataclass(frozen=True, slots=True)
+class EncodedCudaVersion:
+    """CUDA major/minor version represented in CUDA's integer ``encoded`` form."""
+
+    encoded: int
+    major: int
+    minor: int
+
+    @classmethod
+    def from_encoded(cls: type[EncodedCudaVersionT], encoded: int | str) -> EncodedCudaVersionT:
+        if isinstance(encoded, str):
+            try:
+                encoded_int = int(encoded)
+            except ValueError as exc:
+                raise ValueError(
+                    f"{cls.__name__}.from_encoded() expected an integer or decimal string, got {encoded!r}."
+                ) from exc
+        elif isinstance(encoded, int):
+            encoded_int = encoded
+        else:
+            raise TypeError(
+                f"{cls.__name__}.from_encoded() expected an integer or decimal string, got {type(encoded).__name__}."
+            )
+        if encoded_int < 0:
+            raise ValueError(
+                f"{cls.__name__}.from_encoded() expected a non-negative encoded CUDA version, got {encoded_int}."
+            )
+        # CUDA encodes versions as major * 1000 + minor * 10. The least-significant
+        # decimal is ignored here: it is 0 in all CUDA releases and is not a patch version.
+        return cls(
+            encoded=encoded_int,
+            major=encoded_int // 1000,
+            minor=(encoded_int % 1000) // 10,
+        )
 
 
 class ReadCudaHeaderVersionError(RuntimeError):
     """Raised when ``read_cuda_header_version()`` cannot determine the CTK version from ``cuda.h``."""
 
 
 @dataclass(frozen=True, slots=True)
-class CudaToolkitVersion:
+class CudaToolkitVersion(EncodedCudaVersion):
     """CUDA Toolkit version encoded by the ``CUDA_VERSION`` macro in ``cuda.h``."""
 
-    encoded: int
-    major: int
-    minor: int
-
 
 def parse_cuda_header_version(header_text: str) -> CudaToolkitVersion | None:
     """Parse the CUDA Toolkit major/minor version from ``cuda.h`` text."""
     match = _CUDA_VERSION_RE.search(header_text)
     if match is None:
         return None
-    encoded = int(match.group("encoded"))
-    return CudaToolkitVersion(
-        encoded=encoded,
-        major=encoded // 1000,
-        minor=(encoded % 1000) // 10,
-    )
+    return CudaToolkitVersion.from_encoded(match.group("encoded"))
 
 
 @functools.cache

cuda_pathfinder/tests/test_utils_driver_info.py

@@ -73,6 +73,17 @@ def test_query_driver_cuda_version_returns_parsed_dataclass(monkeypatch):
     )
 
 
+def test_driver_cuda_version_from_encoded_returns_subclass_instance():
+    version = driver_info.DriverCudaVersion.from_encoded(12080)
+
+    assert version == driver_info.DriverCudaVersion(
+        encoded=12080,
+        major=12,
+        minor=8,
+    )
+    assert type(version) is driver_info.DriverCudaVersion
+
+
 def test_query_driver_cuda_version_wraps_internal_failures(monkeypatch):
     root_cause = RuntimeError("low-level query failed")
 

cuda_pathfinder/tests/test_utils_toolkit_info.py

@@ -13,6 +13,39 @@ def _clear_cuda_header_version_cache():
     toolkit_info.read_cuda_header_version.cache_clear()
 
 
+def test_encoded_cuda_version_from_encoded_decodes_major_minor():
+    assert toolkit_info.EncodedCudaVersion.from_encoded(13020) == toolkit_info.EncodedCudaVersion(
+        encoded=13020,
+        major=13,
+        minor=2,
+    )
+
+
+def test_encoded_cuda_version_from_encoded_accepts_decimal_string():
+    assert toolkit_info.EncodedCudaVersion.from_encoded("13020") == toolkit_info.EncodedCudaVersion(
+        encoded=13020,
+        major=13,
+        minor=2,
+    )
+
+
+def test_encoded_cuda_version_from_encoded_raises_helpful_error_for_invalid_string():
+    with pytest.raises(
+        ValueError,
+        match=r"EncodedCudaVersion\.from_encoded\(\) expected an integer or decimal string, got '13\.2'",
+    ):
+        toolkit_info.EncodedCudaVersion.from_encoded("13.2")
+
+
+@pytest.mark.parametrize("encoded", [-1, "-1"])
+def test_encoded_cuda_version_from_encoded_rejects_negative_values(encoded):
+    with pytest.raises(
+        ValueError,
+        match=r"EncodedCudaVersion\.from_encoded\(\) expected a non-negative encoded CUDA version, got -1",
+    ):
+        toolkit_info.EncodedCudaVersion.from_encoded(encoded)
+
+
 def test_parse_cuda_header_version_returns_parsed_dataclass():
     header_text = """
     #ifndef CUDA_H
@@ -28,6 +61,17 @@ def test_parse_cuda_header_version_returns_parsed_dataclass():
     )
 
 
+def test_cuda_toolkit_version_from_encoded_returns_subclass_instance():
+    version = toolkit_info.CudaToolkitVersion.from_encoded(12090)
+
+    assert version == toolkit_info.CudaToolkitVersion(
+        encoded=12090,
+        major=12,
+        minor=9,
+    )
+    assert type(version) is toolkit_info.CudaToolkitVersion
+
+
 def test_parse_cuda_header_version_returns_none_when_macro_is_missing():
     header_text = """
     #ifndef CUDA_H