Recognize stub pyi Python files. (#2182)

Recognize stub ``pyi`` Python files.

Refs pylint-dev/pylint#4987

Co-authored-by: Jacob Walls <[email protected]>

Author: mbyrnepr2

ChangeLog

@@ -156,6 +156,10 @@ Release date: TBA
 
   Refs pylint-dev/pylint#8554
 
+* Recognize stub ``pyi`` Python files.
+
+  Refs pylint-dev/pylint#4987
+
 
 What's New in astroid 2.15.4?
 =============================

astroid/interpreter/_import/spec.py

@@ -163,7 +163,7 @@ def find_module(
 
         for entry in submodule_path:
             package_directory = os.path.join(entry, modname)
-            for suffix in (".py", importlib.machinery.BYTECODE_SUFFIXES[0]):
+            for suffix in (".py", ".pyi", importlib.machinery.BYTECODE_SUFFIXES[0]):
                 package_file_name = "__init__" + suffix
                 file_path = os.path.join(package_directory, package_file_name)
                 if os.path.isfile(file_path):

astroid/modutils.py

@@ -44,10 +44,10 @@
 
 
 if sys.platform.startswith("win"):
-    PY_SOURCE_EXTS = ("py", "pyw")
+    PY_SOURCE_EXTS = ("py", "pyw", "pyi")
     PY_COMPILED_EXTS = ("dll", "pyd")
 else:
-    PY_SOURCE_EXTS = ("py",)
+    PY_SOURCE_EXTS = ("py", "pyi")
     PY_COMPILED_EXTS = ("so",)
 
 
@@ -274,9 +274,6 @@ def _get_relative_base_path(filename: str, path_to_check: str) -> list[str] | No
     if os.path.normcase(real_filename).startswith(path_to_check):
         importable_path = real_filename
 
-    # if "var" in path_to_check:
-    #     breakpoint()
-
     if importable_path:
         base_path = os.path.splitext(importable_path)[0]
         relative_base_path = base_path[len(path_to_check) :]
@@ -476,7 +473,7 @@ def get_module_files(
             continue
         _handle_blacklist(blacklist, dirnames, filenames)
         # check for __init__.py
-        if not list_all and "__init__.py" not in filenames:
+        if not list_all and {"__init__.py", "__init__.pyi"}.isdisjoint(filenames):
             dirnames[:] = ()
             continue
         for filename in filenames:
@@ -499,6 +496,8 @@ def get_source_file(filename: str, include_no_ext: bool = False) -> str:
     """
     filename = os.path.abspath(_path_from_filename(filename))
     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:
         source_path = f"{base}.{ext}"
         if os.path.exists(source_path):
@@ -663,7 +662,7 @@ def _is_python_file(filename: str) -> bool:
 
     .pyc and .pyo are ignored
     """
-    return filename.endswith((".py", ".so", ".pyd", ".pyw"))
+    return filename.endswith((".py", ".pyi", ".so", ".pyd", ".pyw"))
 
 
 def _has_init(directory: str) -> str | None:

tests/test_modutils.py

@@ -287,6 +287,11 @@ def test(self) -> None:
     def test_raise(self) -> None:
         self.assertRaises(modutils.NoSourceFile, modutils.get_source_file, "whatever")
 
+    def test_(self) -> None:
+        package = resources.find("pyi_data")
+        module = os.path.join(package, "__init__.pyi")
+        self.assertEqual(modutils.get_source_file(module), os.path.normpath(module))
+
 
 class IsStandardModuleTest(resources.SysPathSetup, unittest.TestCase):
     """
@@ -417,8 +422,12 @@ def test_success(self) -> None:
         assert modutils.module_in_path("data.module", datadir)
         assert modutils.module_in_path("data.module", (datadir,))
         assert modutils.module_in_path("data.module", os.path.abspath(datadir))
+        assert modutils.module_in_path("pyi_data.module", datadir)
+        assert modutils.module_in_path("pyi_data.module", (datadir,))
+        assert modutils.module_in_path("pyi_data.module", os.path.abspath(datadir))
         # "" will evaluate to cwd
         assert modutils.module_in_path("data.module", "")
+        assert modutils.module_in_path("pyi_data.module", "")
 
     def test_bad_import(self) -> None:
         datadir = resources.find("")
@@ -496,6 +505,19 @@ def test_get_module_files_1(self) -> None:
         ]
         self.assertEqual(modules, {os.path.join(package, x) for x in expected})
 
+    def test_get_module_files_2(self) -> None:
+        package = resources.find("pyi_data/find_test")
+        modules = set(modutils.get_module_files(package, []))
+        expected = [
+            "__init__.py",
+            "__init__.pyi",
+            "module.py",
+            "module2.py",
+            "noendingnewline.py",
+            "nonregr.py",
+        ]
+        self.assertEqual(modules, {os.path.join(package, x) for x in expected})
+
     def test_get_all_files(self) -> None:
         """Test that list_all returns all Python files from given location."""
         non_package = resources.find("data/notamodule")