From dc682c3e062156ae8e9d0c96a5b689e727a64182 Mon Sep 17 00:00:00 2001 From: Jacob Walls Date: Wed, 15 May 2024 19:56:23 -0400 Subject: [PATCH 1/2] Add prefer_stubs configuration --- ChangeLog | 5 +++++ astroid/interpreter/_import/spec.py | 8 ++------ astroid/manager.py | 14 +++++++++++++- astroid/modutils.py | 15 +++++++++------ 4 files changed, 29 insertions(+), 13 deletions(-) diff --git a/ChangeLog b/ChangeLog index 17e69ccb6f..d3b37e696b 100644 --- a/ChangeLog +++ b/ChangeLog @@ -17,6 +17,11 @@ Release date: TBA Closes pylint-dev/pylint#9139 +* Add ``AstroidManager.prefer_stubs`` attribute to control the astroid 3.2.0 feature that prefers stubs. + + Refs pylint-dev/#9626 + Refs pylint-dev/#9623 + What's New in astroid 3.2.0? ============================ diff --git a/astroid/interpreter/_import/spec.py b/astroid/interpreter/_import/spec.py index 59546a016a..09e98c888b 100644 --- a/astroid/interpreter/_import/spec.py +++ b/astroid/interpreter/_import/spec.py @@ -161,14 +161,10 @@ def find_module( pass submodule_path = sys.path - # We're looping on pyi first because if a pyi exists there's probably a reason - # (i.e. the code is hard or impossible to parse), so we take pyi into account - # But we're not quite ready to do this for numpy, see https://github.com/pylint-dev/astroid/pull/2375 - suffixes = (".pyi", ".py", importlib.machinery.BYTECODE_SUFFIXES[0]) - numpy_suffixes = (".py", ".pyi", importlib.machinery.BYTECODE_SUFFIXES[0]) + suffixes = (".py", ".pyi", importlib.machinery.BYTECODE_SUFFIXES[0]) for entry in submodule_path: package_directory = os.path.join(entry, modname) - for suffix in numpy_suffixes if "numpy" in entry else suffixes: + for suffix in suffixes: package_file_name = "__init__" + suffix file_path = os.path.join(package_directory, package_file_name) if os.path.isfile(file_path): diff --git a/astroid/manager.py b/astroid/manager.py index fc30bf9e2f..ade31c0e91 100644 --- a/astroid/manager.py +++ b/astroid/manager.py @@ -61,6 +61,7 @@ class AstroidManager: "extension_package_whitelist": set(), "module_denylist": set(), "_transform": TransformVisitor(), + "prefer_stubs": False, } def __init__(self) -> None: @@ -73,6 +74,7 @@ def __init__(self) -> None: ] self.module_denylist = AstroidManager.brain["module_denylist"] self._transform = AstroidManager.brain["_transform"] + self.prefer_stubs = AstroidManager.brain["prefer_stubs"] @property def always_load_extensions(self) -> bool: @@ -111,6 +113,14 @@ def unregister_transform(self): def builtins_module(self) -> nodes.Module: return self.astroid_cache["builtins"] + @property + def prefer_stubs(self) -> bool: + return AstroidManager.brain["prefer_stubs"] + + @prefer_stubs.setter + def prefer_stubs(self, value: bool) -> None: + AstroidManager.brain["prefer_stubs"] = value + def visit_transforms(self, node: nodes.NodeNG) -> InferenceResult: """Visit the transforms and apply them to the given *node*.""" return self._transform.visit(node) @@ -136,7 +146,9 @@ def ast_from_file( # Call get_source_file() only after a cache miss, # since it calls os.path.exists(). try: - filepath = get_source_file(filepath, include_no_ext=True) + filepath = get_source_file( + filepath, include_no_ext=True, prefer_stubs=self.prefer_stubs + ) source = True except NoSourceFile: pass diff --git a/astroid/modutils.py b/astroid/modutils.py index 6f67d1ab9a..bcb602c6b8 100644 --- a/astroid/modutils.py +++ b/astroid/modutils.py @@ -44,10 +44,12 @@ if sys.platform.startswith("win"): - PY_SOURCE_EXTS = ("pyi", "pyw", "py") + PY_SOURCE_EXTS = ("py", "pyw", "pyi") + PY_SOURCE_EXTS_STUBS_FIRST = ("pyi", "pyw", "py") PY_COMPILED_EXTS = ("dll", "pyd") else: - PY_SOURCE_EXTS = ("pyi", "py") + PY_SOURCE_EXTS = ("py", "pyi") + PY_SOURCE_EXTS_STUBS_FIRST = ("pyi", "py") PY_COMPILED_EXTS = ("so",) @@ -484,7 +486,9 @@ def get_module_files( return files -def get_source_file(filename: str, include_no_ext: bool = False) -> str: +def get_source_file( + filename: str, include_no_ext: bool = False, prefer_stubs: bool = False +) -> str: """Given a python module's file name return the matching source file name (the filename will be returned identically if it's already an absolute path to a python source file). @@ -499,7 +503,7 @@ def get_source_file(filename: str, include_no_ext: bool = False) -> str: base, orig_ext = os.path.splitext(filename) if orig_ext == ".pyi" and os.path.exists(f"{base}{orig_ext}"): return f"{base}{orig_ext}" - for ext in PY_SOURCE_EXTS if "numpy" not in filename else reversed(PY_SOURCE_EXTS): + for ext in PY_SOURCE_EXTS_STUBS_FIRST if prefer_stubs else PY_SOURCE_EXTS: source_path = f"{base}.{ext}" if os.path.exists(source_path): return source_path @@ -671,8 +675,7 @@ def _has_init(directory: str) -> str | None: else return None. """ mod_or_pack = os.path.join(directory, "__init__") - exts = reversed(PY_SOURCE_EXTS) if "numpy" in directory else PY_SOURCE_EXTS - for ext in (*exts, "pyc", "pyo"): + for ext in (*PY_SOURCE_EXTS, "pyc", "pyo"): if os.path.exists(mod_or_pack + "." + ext): return mod_or_pack + "." + ext return None From 7aeebe6e7d5afc9cce9da6076c2e42718796a74c Mon Sep 17 00:00:00 2001 From: Jacob Walls Date: Wed, 15 May 2024 20:04:49 -0400 Subject: [PATCH 2/2] Update test result --- tests/test_modutils.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/tests/test_modutils.py b/tests/test_modutils.py index 3741428bcc..85452b0f77 100644 --- a/tests/test_modutils.py +++ b/tests/test_modutils.py @@ -300,7 +300,8 @@ def test_pyi_preferred(self) -> None: package = resources.find("pyi_data/find_test") module = os.path.join(package, "__init__.py") self.assertEqual( - modutils.get_source_file(module), os.path.normpath(module) + "i" + modutils.get_source_file(module, prefer_stubs=True), + os.path.normpath(module) + "i", )