From 617f8ea4bb53359a0d5c9a1531848766d503a73d Mon Sep 17 00:00:00 2001 From: "Ralf W. Grosse-Kunstleve" Date: Wed, 4 Mar 2026 00:15:04 -0800 Subject: [PATCH] Fix site-packages search order for venv with --system-site-packages (fixes #1716) In a venv created with --system-site-packages, user-site-packages were being searched before the virtual environment's site-packages, which violates PEP 405. The fix ensures that: - In a venv: venv site-packages come first, then user-site-packages, then system site-packages (PEP 405) - Outside a venv: user-site-packages come before system site-packages (PEP 370) Added tests to verify the correct ordering in both scenarios. Made-with: Cursor --- .../cuda/pathfinder/_utils/find_sub_dirs.py | 9 +- .../tests/test_utils_find_sub_dirs.py | 169 ++++++++++++++++++ 2 files changed, 177 insertions(+), 1 deletion(-) diff --git a/cuda_pathfinder/cuda/pathfinder/_utils/find_sub_dirs.py b/cuda_pathfinder/cuda/pathfinder/_utils/find_sub_dirs.py index a37f2dc8244..ebb7b13f488 100644 --- a/cuda_pathfinder/cuda/pathfinder/_utils/find_sub_dirs.py +++ b/cuda_pathfinder/cuda/pathfinder/_utils/find_sub_dirs.py @@ -54,5 +54,12 @@ def find_sub_dirs_all_sitepackages(sub_dirs: Sequence[str]) -> list[str]: if site.ENABLE_USER_SITE: user_site = site.getusersitepackages() if user_site: - parent_dirs.insert(0, user_site) + # Determine insertion index based on whether we're in a virtual environment (fixes #1716): + # - In venv (PEP 405): venv site-packages should come first, then user-site-packages, + # then system site-packages. Insert at index 1 (after venv, before system). + # - Not in venv (PEP 370): user-site-packages should come before system site-packages. + # Insert at index 0. + # Detect venv by checking if sys.prefix differs from sys.base_prefix + insert_idx = 1 if sys.prefix != sys.base_prefix else 0 + parent_dirs.insert(insert_idx, user_site) return find_sub_dirs(parent_dirs, sub_dirs) diff --git a/cuda_pathfinder/tests/test_utils_find_sub_dirs.py b/cuda_pathfinder/tests/test_utils_find_sub_dirs.py index 80bab252910..a647e66099b 100644 --- a/cuda_pathfinder/tests/test_utils_find_sub_dirs.py +++ b/cuda_pathfinder/tests/test_utils_find_sub_dirs.py @@ -89,3 +89,172 @@ def test_find_sub_dirs_sys_path_no_math(): def test_find_sub_dirs_all_sitepackages_no_match(): result = find_sub_dirs_all_sitepackages((NONEXISTENT,)) assert result == [] + + +def test_find_sub_dirs_all_sitepackages_venv_order(mocker, tmp_path): + """Test that in a venv with --system-site-packages, search order is: venv, user, system. + + This test verifies fix for issue #1716: user-site-packages should come after + venv site-packages but before system site-packages. + """ + # Create test directories + venv_site = tmp_path / "venv" / "lib" / "python3.12" / "site-packages" + user_site = tmp_path / "user" / ".local" / "lib" / "python3.12" / "site-packages" + system_site = tmp_path / "system" / "lib" / "python3.12" / "dist-packages" + + venv_site.mkdir(parents=True) + user_site.mkdir(parents=True) + system_site.mkdir(parents=True) + + # Create a test subdirectory in each + test_subdir = ("nvidia", "cuda_runtime", "lib") + (venv_site / "nvidia" / "cuda_runtime" / "lib").mkdir(parents=True) + (user_site / "nvidia" / "cuda_runtime" / "lib").mkdir(parents=True) + (system_site / "nvidia" / "cuda_runtime" / "lib").mkdir(parents=True) + + # Mock site.getsitepackages() to return venv first, then system + mocker.patch( + "cuda.pathfinder._utils.find_sub_dirs.site.getsitepackages", + return_value=[str(venv_site), str(system_site)], + ) + # Mock user site-packages + mocker.patch( + "cuda.pathfinder._utils.find_sub_dirs.site.getusersitepackages", + return_value=str(user_site), + ) + mocker.patch( + "cuda.pathfinder._utils.find_sub_dirs.site.ENABLE_USER_SITE", + True, + ) + # Mock sys.prefix != sys.base_prefix to simulate venv + mocker.patch( + "cuda.pathfinder._utils.find_sub_dirs.sys.prefix", + str(tmp_path / "venv"), + ) + mocker.patch( + "cuda.pathfinder._utils.find_sub_dirs.sys.base_prefix", + str(tmp_path / "system"), + ) + + # Clear cache to ensure mocks take effect + from cuda.pathfinder._utils.find_sub_dirs import find_sub_dirs_cached + + find_sub_dirs_cached.cache_clear() + + result = find_sub_dirs_all_sitepackages(test_subdir) + + # Verify order: venv should come first, then user, then system + assert len(result) == 3 + assert result[0] == str(venv_site / "nvidia" / "cuda_runtime" / "lib") + assert result[1] == str(user_site / "nvidia" / "cuda_runtime" / "lib") + assert result[2] == str(system_site / "nvidia" / "cuda_runtime" / "lib") + + +def test_find_sub_dirs_all_sitepackages_non_venv_order(mocker, tmp_path): + """Test that outside a venv, search order is: user, system. + + This verifies PEP 370 behavior: user-site-packages should come before + system site-packages when not in a venv. + """ + # Create test directories + user_site = tmp_path / "user" / ".local" / "lib" / "python3.12" / "site-packages" + system_site = tmp_path / "system" / "lib" / "python3.12" / "dist-packages" + + user_site.mkdir(parents=True) + system_site.mkdir(parents=True) + + # Create a test subdirectory in each + test_subdir = ("nvidia", "cuda_runtime", "lib") + (user_site / "nvidia" / "cuda_runtime" / "lib").mkdir(parents=True) + (system_site / "nvidia" / "cuda_runtime" / "lib").mkdir(parents=True) + + # Mock site.getsitepackages() to return only system site-packages + mocker.patch( + "cuda.pathfinder._utils.find_sub_dirs.site.getsitepackages", + return_value=[str(system_site)], + ) + # Mock user site-packages + mocker.patch( + "cuda.pathfinder._utils.find_sub_dirs.site.getusersitepackages", + return_value=str(user_site), + ) + mocker.patch( + "cuda.pathfinder._utils.find_sub_dirs.site.ENABLE_USER_SITE", + True, + ) + # Mock sys.prefix == sys.base_prefix to simulate non-venv + mocker.patch( + "cuda.pathfinder._utils.find_sub_dirs.sys.prefix", + str(tmp_path / "system"), + ) + mocker.patch( + "cuda.pathfinder._utils.find_sub_dirs.sys.base_prefix", + str(tmp_path / "system"), + ) + + # Clear cache to ensure mocks take effect + from cuda.pathfinder._utils.find_sub_dirs import find_sub_dirs_cached + + find_sub_dirs_cached.cache_clear() + + result = find_sub_dirs_all_sitepackages(test_subdir) + + # Verify order: user should come first, then system + assert len(result) == 2 + assert result[0] == str(user_site / "nvidia" / "cuda_runtime" / "lib") + assert result[1] == str(system_site / "nvidia" / "cuda_runtime" / "lib") + + +def test_find_sub_dirs_all_sitepackages_venv_first_match(mocker, tmp_path): + """Test that in a venv, venv site-packages is searched first (fixes #1716). + + This test ensures that when a file exists in both venv and user site-packages, + the venv version is found first, matching Python's import behavior. + """ + # Create test directories + venv_site = tmp_path / "venv" / "lib" / "python3.12" / "site-packages" + user_site = tmp_path / "user" / ".local" / "lib" / "python3.12" / "site-packages" + + venv_site.mkdir(parents=True) + user_site.mkdir(parents=True) + + # Create the same subdirectory in both + test_subdir = ("nvidia", "cuda_runtime", "lib") + (venv_site / "nvidia" / "cuda_runtime" / "lib").mkdir(parents=True) + (user_site / "nvidia" / "cuda_runtime" / "lib").mkdir(parents=True) + + # Mock site.getsitepackages() to return venv first + mocker.patch( + "cuda.pathfinder._utils.find_sub_dirs.site.getsitepackages", + return_value=[str(venv_site)], + ) + # Mock user site-packages + mocker.patch( + "cuda.pathfinder._utils.find_sub_dirs.site.getusersitepackages", + return_value=str(user_site), + ) + mocker.patch( + "cuda.pathfinder._utils.find_sub_dirs.site.ENABLE_USER_SITE", + True, + ) + # Mock sys.prefix != sys.base_prefix to simulate venv + mocker.patch( + "cuda.pathfinder._utils.find_sub_dirs.sys.prefix", + str(tmp_path / "venv"), + ) + mocker.patch( + "cuda.pathfinder._utils.find_sub_dirs.sys.base_prefix", + str(tmp_path / "system"), + ) + + # Clear cache to ensure mocks take effect + from cuda.pathfinder._utils.find_sub_dirs import find_sub_dirs_cached + + find_sub_dirs_cached.cache_clear() + + result = find_sub_dirs_all_sitepackages(test_subdir) + + # Verify venv comes first (this would fail with old code that puts user first) + assert len(result) >= 2 + assert result[0] == str(venv_site / "nvidia" / "cuda_runtime" / "lib") + assert result[1] == str(user_site / "nvidia" / "cuda_runtime" / "lib")