optimize _add_imports_to_all

Author: LarsKue

bayesflow/utils/_docs/_populate_all.py

@@ -1,23 +1,31 @@
-import inspect
+import sys
+import types
 
 
 def _add_imports_to_all(include_modules: bool | list[str] = False, exclude: list[str] | None = None):
     """Add all global variables to __all__"""
     if not isinstance(include_modules, (bool, list)):
         raise ValueError("include_modules must be a boolean or a list of strings")
 
-    exclude = exclude or []
-    calling_module = inspect.stack()[1]
-    local_stack = calling_module[0]
-    global_vars = local_stack.f_globals
-    all_vars = global_vars["__all__"] if "__all__" in global_vars else []
-    included_vars = []
-    for var_name in set(global_vars.keys()):
-        if inspect.ismodule(global_vars[var_name]):
-            if include_modules is True and var_name not in exclude and not var_name.startswith("_"):
-                included_vars.append(var_name)
-            elif isinstance(include_modules, list) and var_name in include_modules:
-                included_vars.append(var_name)
-        elif var_name not in exclude and not var_name.startswith("_"):
-            included_vars.append(var_name)
-    global_vars["__all__"] = sorted(list(set(all_vars).union(included_vars)))
+    exclude_set = set(exclude or [])
+    contains = exclude_set.__contains__
+    mod_type = types.ModuleType
+    frame = sys._getframe(1)
+    g: dict = frame.f_globals
+    existing = set(g.get("__all__", []))
+
+    to_add = []
+    include_list = include_modules if isinstance(include_modules, list) else ()
+    inc_all = include_modules is True
+
+    for name, val in g.items():
+        if name.startswith("_") or contains(name):
+            continue
+
+        if isinstance(val, mod_type):
+            if inc_all or name in include_list:
+                to_add.append(name)
+        else:
+            to_add.append(name)
+
+    g["__all__"] = sorted(existing.union(to_add))