From 45bb462960919d27b14779c70ac8b4607c43af78 Mon Sep 17 00:00:00 2001 From: Stefan van der Walt Date: Fri, 15 Mar 2024 15:47:02 -0700 Subject: [PATCH] Refuse star imports in stub loader Block these types of imports until #94 is resolved. --- lazy_loader/__init__.py | 8 +++++++- lazy_loader/tests/test_lazy_loader.py | 5 +++++ 2 files changed, 12 insertions(+), 1 deletion(-) diff --git a/lazy_loader/__init__.py b/lazy_loader/__init__.py index ae593c8..7c9dbb9 100644 --- a/lazy_loader/__init__.py +++ b/lazy_loader/__init__.py @@ -296,7 +296,13 @@ def visit_ImportFrom(self, node: ast.ImportFrom): ) if node.module: attrs: list = self._submod_attrs.setdefault(node.module, []) - attrs.extend(alias.name for alias in node.names) + aliases = [alias.name for alias in node.names] + if "*" in aliases: + raise ValueError( + "lazy stub loader does not support star import " + f"`from {node.module} import *`" + ) + attrs.extend(aliases) else: self._submodules.update(alias.name for alias in node.names) diff --git a/lazy_loader/tests/test_lazy_loader.py b/lazy_loader/tests/test_lazy_loader.py index f28439f..19187ba 100644 --- a/lazy_loader/tests/test_lazy_loader.py +++ b/lazy_loader/tests/test_lazy_loader.py @@ -153,6 +153,11 @@ def test_stub_loading_errors(tmp_path): with pytest.raises(ValueError, match="Cannot load imports from non-existent stub"): lazy.attach_stub("name", "not a file") + stub2 = tmp_path / "stub2.pyi" + stub2.write_text("from .mod import *\n") + with pytest.raises(ValueError, match=".*does not support star import"): + lazy.attach_stub("name", str(stub2)) + def test_require_kwarg(): have_importlib_metadata = importlib.util.find_spec("importlib.metadata") is not None