Look up each module name without a package version only once

Author: iacobdaniel

aikido_zen/sinks/builtins_import.py

@@ -6,6 +6,9 @@
 
 running_import_scan = contextvars.ContextVar("running_import_scan", default=False)
 
+# version() scans all of sys.path, so names it cannot find (re, json) are looked up once.
+names_without_version = set()
+
 
 @after
 def _import(func, instance, args, kwargs, return_value):
@@ -30,7 +33,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):
+        if name in names_without_version or PackagesStore.get_package(name):
             return
 
         version = None
@@ -40,6 +43,8 @@ def _import(func, instance, args, kwargs, return_value):
             pass
         if version:
             PackagesStore.add_package(name, version)
+        else:
+            names_without_version.add(name)
     finally:
         running_import_scan.set(False)
 

aikido_zen/sinks/tests/builtins_import_test.py

@@ -1,15 +1,20 @@
+import importlib.metadata
+from unittest.mock import patch
+
 import pytest
 
 import aikido_zen
 
 aikido_zen.protect()
 
+import aikido_zen.sinks.builtins_import as builtins_import
 from aikido_zen.background_process.packages import PackagesStore
 
 
 @pytest.fixture(autouse=True)
 def run_around_tests():
     PackagesStore.clear()
+    builtins_import.names_without_version.clear()
 
 
 def test_flask_import():
@@ -24,6 +29,17 @@ 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."""