Recognize stub ``pyi`` Python files. (#2182)HEADmain

Recognize stub ``pyi`` Python files.

Refs pylint-dev/pylint#4987

Co-authored-by: Jacob Walls <jacobtylerwalls@gmail.com>

12 files changed, 33 insertions, 8 deletions

diff --git a/ChangeLog b/ChangeLog
index b5e97c31..58587dc7 100644
--- a/ChangeLog
+++ b/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?
 =============================
diff --git a/astroid/interpreter/_import/spec.py b/astroid/interpreter/_import/spec.py
index b1f8e8db..3c21fd73 100644
--- a/astroid/interpreter/_import/spec.py
+++ b/astroid/interpreter/_import/spec.py
@@ -163,7 +163,7 @@ class ImportlibFinder(Finder):
 
         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):
diff --git a/astroid/modutils.py b/astroid/modutils.py
index b4f3b6e3..33fd3eeb 100644
--- a/astroid/modutils.py
+++ b/astroid/modutils.py
@@ -44,10 +44,10 @@ logger = logging.getLogger(__name__)
 
 
 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:
diff --git a/tests/test_modutils.py b/tests/test_modutils.py
index 34f7132f..0da6ce21 100644
--- a/tests/test_modutils.py
+++ b/tests/test_modutils.py
@@ -287,6 +287,11 @@ class GetSourceFileTest(unittest.TestCase):
     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 @@ class ModuleInPathTest(resources.SysPathSetup, unittest.TestCase):
         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 @@ class GetModuleFilesTest(unittest.TestCase):
         ]
         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")
diff --git a/tests/testdata/python3/pyi_data/__init__.pyi b/tests/testdata/python3/pyi_data/__init__.pyi
new file mode 100644
index 00000000..e69de29b
--- /dev/null
+++ b/tests/testdata/python3/pyi_data/__init__.pyi
diff --git a/tests/testdata/python3/pyi_data/find_test/__init__.py b/tests/testdata/python3/pyi_data/find_test/__init__.py
new file mode 100644
index 00000000..e69de29b
--- /dev/null
+++ b/tests/testdata/python3/pyi_data/find_test/__init__.py
diff --git a/tests/testdata/python3/pyi_data/find_test/__init__.pyi b/tests/testdata/python3/pyi_data/find_test/__init__.pyi
new file mode 100644
index 00000000..e69de29b
--- /dev/null
+++ b/tests/testdata/python3/pyi_data/find_test/__init__.pyi
diff --git a/tests/testdata/python3/pyi_data/find_test/module.py b/tests/testdata/python3/pyi_data/find_test/module.py
new file mode 100644
index 00000000..e69de29b
--- /dev/null
+++ b/tests/testdata/python3/pyi_data/find_test/module.py
diff --git a/tests/testdata/python3/pyi_data/find_test/module2.py b/tests/testdata/python3/pyi_data/find_test/module2.py
new file mode 100644
index 00000000..e69de29b
--- /dev/null
+++ b/tests/testdata/python3/pyi_data/find_test/module2.py
diff --git a/tests/testdata/python3/pyi_data/find_test/noendingnewline.py b/tests/testdata/python3/pyi_data/find_test/noendingnewline.py
new file mode 100644
index 00000000..e69de29b
--- /dev/null
+++ b/tests/testdata/python3/pyi_data/find_test/noendingnewline.py
diff --git a/tests/testdata/python3/pyi_data/find_test/nonregr.py b/tests/testdata/python3/pyi_data/find_test/nonregr.py
new file mode 100644
index 00000000..e69de29b
--- /dev/null
+++ b/tests/testdata/python3/pyi_data/find_test/nonregr.py
diff --git a/tests/testdata/python3/pyi_data/module.py b/tests/testdata/python3/pyi_data/module.py
new file mode 100644
index 00000000..e69de29b
--- /dev/null
+++ b/tests/testdata/python3/pyi_data/module.py