diff --git a/src/winml/modelkit/commands/sys.py b/src/winml/modelkit/commands/sys.py index dc682bc30..e5f70c503 100644 --- a/src/winml/modelkit/commands/sys.py +++ b/src/winml/modelkit/commands/sys.py @@ -23,6 +23,8 @@ from __future__ import annotations +import ctypes +import ctypes.wintypes import json import logging import platform @@ -68,10 +70,80 @@ def _get_python_info() -> dict[str, Any]: } +# Intentionally limited to the host architectures Windows 11 ships on. +# 32-bit ARM (ARMNT 0xC4, ARM 0x1C4) and IA64 are not Windows 11 host +# targets, so IsWow64Process2 will not report them in practice; an +# unmapped value falls through to None and logs at debug level. +_IMAGE_FILE_MACHINE_TO_NAME = { + 0x8664: "AMD64", + 0xAA64: "ARM64", + 0x14C: "x86", +} + + +if sys.platform == "win32": + try: + _KERNEL32 = ctypes.WinDLL("kernel32", use_last_error=True) + _IS_WOW64_PROCESS_2 = _KERNEL32.IsWow64Process2 + _IS_WOW64_PROCESS_2.argtypes = [ + ctypes.wintypes.HANDLE, + ctypes.POINTER(ctypes.c_ushort), + ctypes.POINTER(ctypes.c_ushort), + ] + _IS_WOW64_PROCESS_2.restype = ctypes.wintypes.BOOL + except (OSError, AttributeError): + _KERNEL32 = None # type: ignore[assignment] + _IS_WOW64_PROCESS_2 = None # type: ignore[assignment] +else: + _KERNEL32 = None # type: ignore[assignment] + _IS_WOW64_PROCESS_2 = None # type: ignore[assignment] + + +def _query_native_machine_via_win32() -> int | None: + """Call IsWow64Process2 and return the native IMAGE_FILE_MACHINE_* code. + + Returns None on non-Windows or when the API call fails. Logs the + GetLastError code at debug level on failure so a regression in + IsWow64Process2 wiring is not silently swallowed. + """ + if sys.platform != "win32" or _IS_WOW64_PROCESS_2 is None or _KERNEL32 is None: + return None + + proc = ctypes.c_ushort(0) + native = ctypes.c_ushort(0) + if not _IS_WOW64_PROCESS_2( + _KERNEL32.GetCurrentProcess(), ctypes.byref(proc), ctypes.byref(native) + ): + logger.debug("IsWow64Process2 failed: GetLastError=%d", ctypes.get_last_error()) + return None + return native.value + + +def _get_windows_native_machine() -> str | None: + """Return the host architecture name, or None when unavailable. + + platform.machine() returns the *process* arch (PROCESSOR_ARCHITECTURE), + so an x64 Python running under ARM64 emulation reports "AMD64". This + consults IsWow64Process2 for the real host machine type, which is what + the user expects to see in `winml sys`. PROCESSOR_ARCHITEW6432 is + unreliable on ARM64 (Prism emulation does not set it on Snapdragon X). + """ + if sys.platform != "win32": + return None + raw = _query_native_machine_via_win32() + if raw is None: + return None + name = _IMAGE_FILE_MACHINE_TO_NAME.get(raw) + if name is None: + logger.debug("IsWow64Process2 returned unmapped native machine: 0x%x", raw) + return name + + def _get_platform_info() -> dict[str, Any]: """Gather OS and platform information.""" system = platform.system() release = platform.release() + machine = platform.machine() # For Windows, use OS class for accurate Windows 11 detection # platform.release() may incorrectly report '10' on some Python versions @@ -86,10 +158,14 @@ def _get_platform_info() -> dict[str, Any]: # Fallback to platform.release() if OS detection fails pass + native_machine = _get_windows_native_machine() + if native_machine: + machine = native_machine + return { "system": system, "release": release, - "machine": platform.machine(), + "machine": machine, "processor": platform.processor() or "Unknown", } diff --git a/tests/unit/sysinfo/test_sysinfo.py b/tests/unit/sysinfo/test_sysinfo.py index fbb5c6b41..f831af027 100644 --- a/tests/unit/sysinfo/test_sysinfo.py +++ b/tests/unit/sysinfo/test_sysinfo.py @@ -12,13 +12,69 @@ import sys from unittest.mock import MagicMock, patch +import pytest + + +class TestGetWindowsNativeMachine: + """Test _get_windows_native_machine and the IMAGE_FILE_MACHINE mapping. + + These tests lock in the IMAGE_FILE_MACHINE_* → display-name mapping so + accidentally renaming "ARM64" → "Arm64" (or similar) ships with a test + failure rather than silently changing user-visible output. The + higher-level _get_platform_info tests mock the helper's return value + and would not catch a mapping rename. + """ + + @pytest.mark.parametrize( + ("raw", "expected"), + [ + (0x8664, "AMD64"), + (0xAA64, "ARM64"), + (0x14C, "x86"), + # ARMNT (0xC4) is not a Windows 11 host arch — intentionally unmapped + (0xC4, None), + # IMAGE_FILE_MACHINE_UNKNOWN — IsWow64Process2 returns this for the + # process slot when the process is native on the host + (0x0, None), + ], + ) + def test_native_machine_mapping( + self, raw: int, expected: str | None, monkeypatch: pytest.MonkeyPatch + ) -> None: + from winml.modelkit.commands import sys as sys_mod + + monkeypatch.setattr(sys, "platform", "win32") + monkeypatch.setattr(sys_mod, "_query_native_machine_via_win32", lambda: raw) + + assert sys_mod._get_windows_native_machine() == expected + + def test_returns_none_when_win32_query_fails( + self, monkeypatch: pytest.MonkeyPatch + ) -> None: + from winml.modelkit.commands import sys as sys_mod + + monkeypatch.setattr(sys, "platform", "win32") + monkeypatch.setattr(sys_mod, "_query_native_machine_via_win32", lambda: None) + + assert sys_mod._get_windows_native_machine() is None + + def test_returns_none_on_non_windows(self, monkeypatch: pytest.MonkeyPatch) -> None: + from winml.modelkit.commands import sys as sys_mod + + monkeypatch.setattr(sys, "platform", "linux") + + assert sys_mod._get_windows_native_machine() is None + class TestGetPlatformInfo: """Test _get_platform_info function.""" + @patch("winml.modelkit.commands.sys._get_windows_native_machine", return_value=None) @patch("winml.modelkit.commands.sys.OS") @patch("winml.modelkit.commands.sys.platform") - def test_windows_11_detection(self, mock_platform: MagicMock, mock_os_class: MagicMock) -> None: + def test_windows_11_detection( + self, mock_platform: MagicMock, mock_os_class: MagicMock, _mock_native: MagicMock + ) -> None: """Test Windows 11 is correctly detected.""" from winml.modelkit.commands.sys import _get_platform_info @@ -39,9 +95,12 @@ def test_windows_11_detection(self, mock_platform: MagicMock, mock_os_class: Mag assert result["machine"] == "AMD64" mock_os_class.get.assert_called_once() + @patch("winml.modelkit.commands.sys._get_windows_native_machine", return_value=None) @patch("winml.modelkit.commands.sys.OS") @patch("winml.modelkit.commands.sys.platform") - def test_windows_10_detection(self, mock_platform: MagicMock, mock_os_class: MagicMock) -> None: + def test_windows_10_detection( + self, mock_platform: MagicMock, mock_os_class: MagicMock, _mock_native: MagicMock + ) -> None: """Test Windows 10 is correctly detected.""" from winml.modelkit.commands.sys import _get_platform_info @@ -61,9 +120,12 @@ def test_windows_10_detection(self, mock_platform: MagicMock, mock_os_class: Mag assert result["release"] == "10" assert result["machine"] == "AMD64" + @patch("winml.modelkit.commands.sys._get_windows_native_machine", return_value=None) @patch("winml.modelkit.commands.sys.OS") @patch("winml.modelkit.commands.sys.platform") - def test_windows_7_preserved(self, mock_platform: MagicMock, mock_os_class: MagicMock) -> None: + def test_windows_7_preserved( + self, mock_platform: MagicMock, mock_os_class: MagicMock, _mock_native: MagicMock + ) -> None: """Test Windows 7 version is preserved (not changed to 10).""" from winml.modelkit.commands.sys import _get_platform_info @@ -83,9 +145,12 @@ def test_windows_7_preserved(self, mock_platform: MagicMock, mock_os_class: Magi assert result["release"] == "7" # Should keep original value assert result["machine"] == "AMD64" + @patch("winml.modelkit.commands.sys._get_windows_native_machine", return_value=None) @patch("winml.modelkit.commands.sys.OS") @patch("winml.modelkit.commands.sys.platform") - def test_windows_81_preserved(self, mock_platform: MagicMock, mock_os_class: MagicMock) -> None: + def test_windows_81_preserved( + self, mock_platform: MagicMock, mock_os_class: MagicMock, _mock_native: MagicMock + ) -> None: """Test Windows 8.1 version is preserved (not changed to 10).""" from winml.modelkit.commands.sys import _get_platform_info @@ -105,10 +170,11 @@ def test_windows_81_preserved(self, mock_platform: MagicMock, mock_os_class: Mag assert result["release"] == "8.1" # Should keep original value assert result["machine"] == "AMD64" + @patch("winml.modelkit.commands.sys._get_windows_native_machine", return_value=None) @patch("winml.modelkit.commands.sys.OS") @patch("winml.modelkit.commands.sys.platform") def test_windows_detection_fallback_on_exception( - self, mock_platform: MagicMock, mock_os_class: MagicMock + self, mock_platform: MagicMock, mock_os_class: MagicMock, _mock_native: MagicMock ) -> None: """Test fallback to platform.release() when OS detection fails.""" from winml.modelkit.commands.sys import _get_platform_info @@ -129,6 +195,64 @@ def test_windows_detection_fallback_on_exception( assert result["release"] == "10" assert result["machine"] == "AMD64" + @patch("winml.modelkit.commands.sys._get_windows_native_machine") + @patch("winml.modelkit.commands.sys.OS") + @patch("winml.modelkit.commands.sys.platform") + def test_windows_arm64_host_with_x64_python( + self, + mock_platform: MagicMock, + mock_os_class: MagicMock, + mock_native: MagicMock, + ) -> None: + """x64 Python on ARM64 host: IsWow64Process2 reveals the real host arch. + + platform.machine() returns the process arch ("AMD64") under emulation; + winml sys should display the host arch ("ARM64") instead. + """ + from winml.modelkit.commands.sys import _get_platform_info + + mock_platform.system.return_value = "Windows" + mock_platform.release.return_value = "10" + mock_platform.machine.return_value = "AMD64" + mock_platform.processor.return_value = "Snapdragon" + + mock_os_instance = MagicMock() + mock_os_instance.is_windows_11.return_value = True + mock_os_class.get.return_value = mock_os_instance + + mock_native.return_value = "ARM64" + + result = _get_platform_info() + + assert result["machine"] == "ARM64" + + @patch("winml.modelkit.commands.sys._get_windows_native_machine") + @patch("winml.modelkit.commands.sys.OS") + @patch("winml.modelkit.commands.sys.platform") + def test_windows_native_lookup_failure_falls_back( + self, + mock_platform: MagicMock, + mock_os_class: MagicMock, + mock_native: MagicMock, + ) -> None: + """When IsWow64Process2 yields None, fall back to platform.machine().""" + from winml.modelkit.commands.sys import _get_platform_info + + mock_platform.system.return_value = "Windows" + mock_platform.release.return_value = "10" + mock_platform.machine.return_value = "AMD64" + mock_platform.processor.return_value = "Intel64 Family 6" + + mock_os_instance = MagicMock() + mock_os_instance.is_windows_11.return_value = True + mock_os_class.get.return_value = mock_os_instance + + mock_native.return_value = None + + result = _get_platform_info() + + assert result["machine"] == "AMD64" + @patch("winml.modelkit.commands.sys.platform") def test_non_windows_platform(self, mock_platform: MagicMock) -> None: """Test non-Windows platforms pass through unchanged."""