Modified import resolution logic to handle namespace stub packages. T…

…his case isn't explicitly covered by PEP 561, but there is a proposal to amend the PEP to clarify how to handle this case.
microsoft · Oct 2, 2021 · 5e38fab · 5e38fab
1 parent cdbe961
commit 5e38fab
Show file tree

Hide file tree

Showing 2 changed files with 31 additions and 1 deletion.
diff --git a/packages/pyright-internal/src/analyzer/importResolver.ts b/packages/pyright-internal/src/analyzer/importResolver.ts
@@ -793,7 +793,11 @@ export class ImportResolver {
 
             // We found fully typed stub packages.
             if (importResult.packageDirectory) {
-                return importResult;
+                // If this is a namespace package that wasn't resolved, assume that
+                // it's a partial stub package and continue looking for a real package.
+                if (!importResult.isNamespacePackage || importResult.isImportFound) {
+                    return importResult;
+                }
             }
         }
 

diff --git a/packages/pyright-internal/src/tests/importResolver.test.ts b/packages/pyright-internal/src/tests/importResolver.test.ts
@@ -137,6 +137,10 @@ test('stub package', () => {
             path: combinePaths(libraryRoot, 'myLib-stubs', 'stub.pyi'),
             content: '# empty',
         },
+        {
+            path: combinePaths(libraryRoot, 'myLib-stubs', '__init__.pyi'),
+            content: '# empty',
+        },
         {
             path: combinePaths(libraryRoot, 'myLib', 'partialStub.py'),
             content: 'def test(): pass',
@@ -148,6 +152,28 @@ test('stub package', () => {
     assert(!importResult.isImportFound);
 });
 
+test('stub namespace package', () => {
+    const files = [
+        {
+            path: combinePaths(libraryRoot, 'myLib-stubs', 'stub.pyi'),
+            content: '# empty',
+        },
+        {
+            path: combinePaths(libraryRoot, 'myLib', 'partialStub.py'),
+            content: 'def test(): pass',
+        },
+    ];
+
+    // If fully typed stub package exists, that wins over the real package.
+    const importResult = getImportResult(files, ['myLib', 'partialStub']);
+    assert(importResult.isImportFound);
+    assert(!importResult.isStubFile);
+    assert.strictEqual(
+        1,
+        importResult.resolvedPaths.filter((f) => f === combinePaths(libraryRoot, 'myLib', 'partialStub.py')).length
+    );
+});
+
 test('stub in typing folder over partial stub package', () => {
     const typingFolder = combinePaths(normalizeSlashes('/'), 'typing');
     const files = [