Coverage for cuda/pathfinder/_utils/windows_arch.py: 88.31%

1# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. 

2# SPDX-License-Identifier: Apache-2.0 

3 

4from __future__ import annotations (empty)

5 

6import platform (empty)

7import sysconfig (empty)

8 

9WINDOWS_PE_MACHINE_BY_ARCH = { (empty)

10 "x64": 0x8664, 

11 "arm64": 0xAA64, 

12} 

13 

14_WINDOWS_ARCH_BY_PE_MACHINE = {machine: arch for arch, machine in WINDOWS_PE_MACHINE_BY_ARCH.items()} (empty)

15 

16 

17class UnsupportedArchError(RuntimeError): (empty)

18 """Raised when Python reports an unsupported Windows architecture.""" 

19 

20 def __init__(self, platform_tag: str) -> None: (empty)

21 self.platform_tag = platform_tag 1 ctx1s

22 super().__init__( 1 ctx1s

23 f"Unsupported Windows Python platform tag: {platform_tag!r}; expected 'win-amd64' or 'win-arm64'" 

24 ) 

25 

26 

27def windows_python_arch() -> str: (empty)

28 """Return the current Windows Python interpreter architecture.""" 

29 raw_platform_tag = sysconfig.get_platform() 11 ctx1abcxdyezBAs

30 platform_tag = raw_platform_tag.lower().replace("_", "-") 11 ctx1abcxdyezBAs

31 

32 if platform_tag == "win-arm64": 11 ctx1abcxdyezBAs

33 return "arm64" 1 ctx1B

34 

35 if platform_tag == "win-amd64": 10 ctx1abcxdyezAs

36 return "x64" 9 ctx1abcxdyezA

37 

38 raise UnsupportedArchError(raw_platform_tag) 1 ctx1s

39 

40 

41def _windows_machine_arch_from_platform() -> str: (empty)

42 """Return the Windows architecture reported by Python's platform module.""" 

43 raw_machine = platform.machine() 4 ctx1tqru

44 machine = raw_machine.lower().replace("_", "-") 4 ctx1tqru

45 

46 if machine in ("amd64", "x86-64"): 4 ctx1tqru

47 return "x64" 2 ctx1tu

48 

49 if machine in ("arm64", "aarch64"): 2 ctx1qr

50 return "arm64" 2 ctx1qr

51 

52 raise RuntimeError(f"Unsupported Windows machine architecture: {raw_machine!r}") 

53 

54 

55def _windows_native_machine() -> int | None: (empty)

56 """Return the native Windows PE machine type, or None on older Windows.""" 

57 import ctypes 3 ctx1gfp

58 from ctypes import wintypes 3 ctx1gfp

59 

60 try: 3 ctx1gfp

61 # These ctypes attributes are absent from the type stubs on non-Windows hosts. 

62 kernel32 = ctypes.WinDLL("kernel32", use_last_error=True) # type: ignore[attr-defined, unused-ignore] 3 ctx1gfp

63 except OSError as exc: 

64 raise RuntimeError("Failed to load kernel32 while detecting the native Windows architecture") from exc 

65 

66 get_current_process = kernel32.GetCurrentProcess 3 ctx1gfp

67 try: 3 ctx1gfp

68 is_wow64_process2 = kernel32.IsWow64Process2 3 ctx1gfp

69 except AttributeError: 1 ctx1p

70 return None 1 ctx1p

71 

72 get_current_process.argtypes = () 2 ctx1gf

73 get_current_process.restype = wintypes.HANDLE 2 ctx1gf

74 is_wow64_process2.argtypes = ( 2 ctx1gf

75 wintypes.HANDLE, 

76 ctypes.POINTER(wintypes.USHORT), 

77 ctypes.POINTER(wintypes.USHORT), 

78 ) 

79 is_wow64_process2.restype = wintypes.BOOL 2 ctx1gf

80 

81 process_machine = wintypes.USHORT() 2 ctx1gf

82 native_machine = wintypes.USHORT() 2 ctx1gf

83 if not is_wow64_process2( 2 ctx1gf

84 get_current_process(), 

85 ctypes.byref(process_machine), 

86 ctypes.byref(native_machine), 

87 ): 

88 error_code = ctypes.get_last_error() # type: ignore[attr-defined, unused-ignore] 1 ctx1f

89 error = ctypes.WinError(error_code) # type: ignore[attr-defined, unused-ignore] 1 ctx1f

90 raise RuntimeError( 1 ctx1f

91 f"IsWow64Process2 failed while detecting the native Windows architecture " 

92 f"(Windows error {error_code}): {error}" 

93 ) from error 

94 return native_machine.value 1 ctx1g

95 

96 

97def windows_machine_arch() -> str: (empty)

98 """Return the native Windows machine architecture, ignoring process emulation.""" 

99 native_machine = _windows_native_machine() 7 ctx1tqruvCD

100 if native_machine is None: 7 ctx1tqruvCD

101 # IsWow64Process2 predates x64-on-Arm emulation, so this fallback is only 

102 # needed on older Windows versions where platform.machine() is sufficient. 

103 return _windows_machine_arch_from_platform() 4 ctx1tqru

104 

105 try: 3 ctx1vCD

106 return _WINDOWS_ARCH_BY_PE_MACHINE[native_machine] 3 ctx1vCD

107 except KeyError: 1 ctx1v

108 raise RuntimeError(f"Unsupported native Windows PE machine type: 0x{native_machine:04x}") from None 1 ctx1v

109 

110 

111def windows_pe_matches_arch(path: str, target_arch: str) -> bool: (empty)

112 """Return whether a Windows Portable Executable (PE) targets the requested architecture. 

113 

114 PE is the file format used for Windows executables and DLLs. This reads the 

115 PE/COFF header's machine field to distinguish x64 images from Arm64 images. 

116 """ 

117 expected_machine = WINDOWS_PE_MACHINE_BY_ARCH.get(target_arch) 14 ctx1abcdehijklmnow

118 if expected_machine is None: 14 ctx1abcdehijklmnow

119 raise ValueError(f"Unsupported Windows target architecture: {target_arch!r}") 

120 

121 try: 14 ctx1abcdehijklmnow

122 with open(path, "rb") as stream: 14 ctx1abcdehijklmnow

123 if stream.read(2) != b"MZ": 14 ctx1abcdehijklmnow

124 return False 1 ctx1w

125 stream.seek(0x3C) 13 ctx1abcdehijklmno

126 pe_offset_bytes = stream.read(4) 13 ctx1abcdehijklmno

127 if len(pe_offset_bytes) != 4: 13 ctx1abcdehijklmno

128 return False 

129 stream.seek(int.from_bytes(pe_offset_bytes, "little")) 13 ctx1abcdehijklmno

130 if stream.read(4) != b"PE\0\0": 13 ctx1abcdehijklmno

131 return False 

132 machine_bytes = stream.read(2) 13 ctx1abcdehijklmno

133 if len(machine_bytes) != 2: 13 ctx1abcdehijklmno

134 return False 

135 except OSError: 

136 return False 

137 

138 return int.from_bytes(machine_bytes, "little") == expected_machine 13 ctx1abcdehijklmno