diff --git a/aikido_zen/background_process/packages.py b/aikido_zen/background_process/packages.py index 6b744bb9d..991ba7823 100644 --- a/aikido_zen/background_process/packages.py +++ b/aikido_zen/background_process/packages.py @@ -56,9 +56,22 @@ def is_version_supported(version, required_version): # packages store, uses python's built in GlobalInterpreterLock (GIL) packages_store = {} +# version() scans all of sys.path, so names it cannot find (re, json) are looked up once. +names_without_version = set() class PackagesStore: + @staticmethod + def add_installed_package(name): + if name in names_without_version or name in packages_store: + return + version = None + try: + version = importlib_metadata.version(name) + except importlib_metadata.PackageNotFoundError: + pass + PackagesStore.add_package(name, version) + @staticmethod def export(): result = [] @@ -70,6 +83,9 @@ def export(): @staticmethod def add_package(package, version): + if not version: + names_without_version.add(package) + return packages_store[package] = { "name": package, "version": version, @@ -91,9 +107,13 @@ def clear(): for package in packages_store.values(): package["cleared"] = True + @staticmethod + def clear_names_without_version(): + names_without_version.clear() + @staticmethod def import_list(imported_packages): for package in imported_packages: - if PackagesStore.get_package(package["name"]): + if package["name"] in packages_store: continue packages_store[package["name"]] = package diff --git a/aikido_zen/sinks/builtins_import.py b/aikido_zen/sinks/builtins_import.py index fe6a00a2c..2a0481624 100644 --- a/aikido_zen/sinks/builtins_import.py +++ b/aikido_zen/sinks/builtins_import.py @@ -1,7 +1,5 @@ from aikido_zen.sinks import on_import, patch_function, after import contextvars -import importlib.metadata -from importlib.metadata import PackageNotFoundError from aikido_zen.background_process.packages import PackagesStore running_import_scan = contextvars.ContextVar("running_import_scan", default=False) @@ -30,16 +28,7 @@ def _import(func, instance, args, kwargs, return_value): # Avoid circular dependencies, this is a double safety-check for if contextvar check fails. return - if PackagesStore.get_package(name): - return - - version = None - try: - version = importlib.metadata.version(name) - except PackageNotFoundError: - pass - if version: - PackagesStore.add_package(name, version) + PackagesStore.add_installed_package(name) finally: running_import_scan.set(False) diff --git a/aikido_zen/sinks/tests/builtins_import_test.py b/aikido_zen/sinks/tests/builtins_import_test.py index 9953417c3..889c95198 100644 --- a/aikido_zen/sinks/tests/builtins_import_test.py +++ b/aikido_zen/sinks/tests/builtins_import_test.py @@ -1,3 +1,6 @@ +import importlib.metadata +from unittest.mock import patch + import pytest import aikido_zen @@ -10,6 +13,7 @@ @pytest.fixture(autouse=True) def run_around_tests(): PackagesStore.clear() + PackagesStore.clear_names_without_version() def test_flask_import(): @@ -24,20 +28,37 @@ def test_django_import(): assert PackagesStore.get_package("django")["version"] == "4.0" +def test_name_without_a_distribution_is_looked_up_once(): + with patch( + "importlib.metadata.version", wraps=importlib.metadata.version + ) as version: + for _ in range(10): + import json + + assert [c.args for c in version.call_args_list] == [("json",)] + assert PackagesStore.get_package("json") is None + + def test_recursive_package_store(monkeypatch): """Test that recursive imports during package scanning don't cause max recursion depth errors.""" - def recursive_get_package(name): + calls = [] + + def recursive_add_installed_package(name): """Recursively add package and its dependencies to PackagesStore.""" + calls.append(name) import flask PackagesStore.clear() - monkeypatch.setattr(PackagesStore, "get_package", recursive_get_package) + monkeypatch.setattr( + PackagesStore, "add_installed_package", recursive_add_installed_package + ) import flask # Restore the original method after the test monkeypatch.undo() + assert calls == ["flask"] def test_recursive_package_store_2(monkeypatch):