Merge pull request #67 from joerunde/faster-imports

⚡ Remove inspect.stack() for performance

import_tracker/lazy_import_errors.py

@@ -11,7 +11,6 @@
 from typing import Callable, Optional, Set
 import importlib.abc
 import importlib.util
-import inspect
 import sys
 
 ## Public ######################################################################
@@ -90,8 +89,8 @@ def _make_extras_import_error(
 
     # Look through frames in the stack to see if there's an extras module
     extras_module = None
-    for frame in inspect.stack():
-        frame_module = frame.frame.f_globals["__name__"]
+    for frame in _FastFrameGenerator():
+        frame_module = frame.f_globals.get("__name__", "")
         if frame_module in extras_modules:
             extras_module = frame_module
             break
@@ -514,23 +513,6 @@ def find_spec(self, fullname, path, *args, **kwargs):
 
     ## Implementation Details ######################################################
 
-    # Custom iterable that uses the low-level sys._getframe to get frames
-    # one-at-a-time
-    class _FrameGenerator:
-        def __init__(self):
-            self._depth = -1
-
-        def __iter__(self):
-            return self
-
-        def __next__(self):
-            self._depth += 1
-            try:
-                return sys._getframe(self._depth)
-            except ValueError:
-                self._depth = -1
-                raise StopIteration
-
     @classmethod
     def _get_non_import_modules(cls):
 
@@ -539,12 +521,33 @@ def _get_non_import_modules(cls):
         return filter(
             lambda x: x != "importlib",
             (
-                frame.f_globals["__name__"].split(".")[0]
-                for frame in cls._FrameGenerator()
+                frame.f_globals.get("__name__", "").split(".")[0]
+                for frame in _FastFrameGenerator()
             ),
         )
 
 
+class _FastFrameGenerator:
+    """Custom iterable that uses the low-level sys._getframe to get frames
+    one-at-a-time.
+    Iterating over this is way faster than using `inspect.stack()`
+    """
+
+    def __init__(self):
+        self._depth = -1
+
+    def __iter__(self):
+        return self
+
+    def __next__(self):
+        self._depth += 1
+        try:
+            return sys._getframe(self._depth)
+        except ValueError:
+            self._depth = -1
+            raise StopIteration
+
+
 def _is_import_time() -> bool:
     """Function to detect if the execution is being called at import
     time by detecting the presence of `importlib._bootstrap` in stack
@@ -554,5 +557,5 @@ def _is_import_time() -> bool:
             True if the execution is at import time otherwise, False
     """
     return "importlib._bootstrap" in [
-        frame.frame.f_globals["__name__"] for frame in inspect.stack()
+        frame.f_globals.get("__name__", "") for frame in _FastFrameGenerator()
     ]

test/test_lazy_import_errors.py

@@ -15,7 +15,7 @@
 import pytest
 
 # Local
-from import_tracker.lazy_import_errors import _LazyErrorMetaFinder
+from import_tracker.lazy_import_errors import _FastFrameGenerator, _LazyErrorMetaFinder
 import import_tracker
 
 
@@ -376,7 +376,7 @@ def test_frame_generator_stop():
     """For completeness, we need to ensure that the FrameGenerator will stop
     correctly if iterated to the end
     """
-    list(_LazyErrorMetaFinder._FrameGenerator())
+    list(_FastFrameGenerator())
 
 
 def test_lazy_import_error_nested():