test(python): Add test coverage for _cpu_check module

.github/workflows/release-python.yml

@@ -135,6 +135,7 @@ jobs:
         env:
           IS_LTS_CPU: ${{ matrix.package == 'polars-lts-cpu' }}
           IS_MACOS: ${{ matrix.os == 'macos-latest' }}
+        # IMPORTANT: All features enabled here should also be included in py-polars/polars/_cpu_check.py
         run: |
           if [[ "$IS_LTS_CPU" = true ]]; then
             FEATURES=+sse3,+ssse3,+sse4.1,+sse4.2,+popcnt

py-polars/polars/_cpu_check.py

@@ -122,10 +122,6 @@ class CPUID_struct(ctypes.Structure):
 
 class CPUID:
     def __init__(self) -> None:
-        if _POLARS_ARCH != "x86-64":
-            msg = "CPUID is only available for x86"
-            raise SystemError(msg)
-
         if _IS_WINDOWS:
             if _IS_64BIT:
                 # VirtualAlloc seems to fail under some weird
@@ -187,11 +183,7 @@ def __del__(self) -> None:
             self.win.VirtualFree(self.addr, 0, _MEM_RELEASE)
 
 
-def read_cpu_flags() -> dict[str, bool]:
-    # Right now we only enable extra feature flags for x86.
-    if _POLARS_ARCH != "x86-64":
-        return {}
-
+def _read_cpu_flags() -> dict[str, bool]:
     # CPU flags from https://en.wikipedia.org/wiki/CPUID
     cpuid = CPUID()
     cpuid1 = cpuid(1, 0)
@@ -223,12 +215,12 @@ def check_cpu_flags() -> None:
         return
 
     expected_cpu_flags = [f.lstrip("+") for f in _POLARS_FEATURE_FLAGS.split(",")]
-    supported_cpu_flags = read_cpu_flags()
+    supported_cpu_flags = _read_cpu_flags()
 
     missing_features = []
     for f in expected_cpu_flags:
         if f not in supported_cpu_flags:
-            msg = f'unknown feature flag "{f}"'
+            msg = f"unknown feature flag: {f!r}"
             raise RuntimeError(msg)
 
         if not supported_cpu_flags[f]:

py-polars/tests/unit/test_cpu_check.py

@@ -0,0 +1,80 @@
+from unittest.mock import Mock
+
+import pytest
+
+from polars import _cpu_check
+from polars._cpu_check import check_cpu_flags
+
+
+@pytest.fixture()
+def _feature_flags(monkeypatch: pytest.MonkeyPatch) -> None:
+    """Use the default set of feature flags."""
+    feature_flags = "+sse3,+ssse3"
+    monkeypatch.setattr(_cpu_check, "_POLARS_FEATURE_FLAGS", feature_flags)
+
+
+@pytest.mark.usefixtures("_feature_flags")
+def test_check_cpu_flags(
+    monkeypatch: pytest.MonkeyPatch, recwarn: pytest.WarningsRecorder
+) -> None:
+    cpu_flags = {"sse3": True, "ssse3": True}
+    mock_read_cpu_flags = Mock(return_value=cpu_flags)
+    monkeypatch.setattr(_cpu_check, "_read_cpu_flags", mock_read_cpu_flags)
+
+    check_cpu_flags()
+
+    assert len(recwarn) == 0
+
+
+@pytest.mark.usefixtures("_feature_flags")
+def test_check_cpu_flags_missing_features(monkeypatch: pytest.MonkeyPatch) -> None:
+    cpu_flags = {"sse3": True, "ssse3": False}
+    mock_read_cpu_flags = Mock(return_value=cpu_flags)
+    monkeypatch.setattr(_cpu_check, "_read_cpu_flags", mock_read_cpu_flags)
+
+    with pytest.warns(RuntimeWarning, match="Missing required CPU features") as w:
+        check_cpu_flags()
+
+    assert "ssse3" in str(w[0].message)
+
+
+def test_check_cpu_flags_unknown_flag(
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    feature_flags = "+sse3,+ssse3,+HelloWorld!"
+    monkeypatch.setattr(_cpu_check, "_POLARS_FEATURE_FLAGS", feature_flags)
+    with pytest.raises(RuntimeError, match="unknown feature flag: 'HelloWorld!'"):
+        check_cpu_flags()
+
+
+def test_check_cpu_flags_skipped_no_flags(monkeypatch: pytest.MonkeyPatch) -> None:
+    mock_read_cpu_flags = Mock()
+    monkeypatch.setattr(_cpu_check, "_read_cpu_flags", mock_read_cpu_flags)
+
+    check_cpu_flags()
+
+    assert mock_read_cpu_flags.call_count == 0
+
+
+@pytest.mark.usefixtures("_feature_flags")
+def test_check_cpu_flags_skipped_lts_cpu(monkeypatch: pytest.MonkeyPatch) -> None:
+    monkeypatch.setattr(_cpu_check, "_POLARS_LTS_CPU", True)
+
+    mock_read_cpu_flags = Mock()
+    monkeypatch.setattr(_cpu_check, "_read_cpu_flags", mock_read_cpu_flags)
+
+    check_cpu_flags()
+
+    assert mock_read_cpu_flags.call_count == 0
+
+
+@pytest.mark.usefixtures("_feature_flags")
+def test_check_cpu_flags_skipped_env_var(monkeypatch: pytest.MonkeyPatch) -> None:
+    monkeypatch.setenv("POLARS_SKIP_CPU_CHECK", "1")
+
+    mock_read_cpu_flags = Mock()
+    monkeypatch.setattr(_cpu_check, "_read_cpu_flags", mock_read_cpu_flags)
+
+    check_cpu_flags()
+
+    assert mock_read_cpu_flags.call_count == 0