PY-21231 Resolve to Python stub even if target is present in Python file

python/psi-api/src/com/jetbrains/python/PyNames.java

@@ -69,7 +69,9 @@ public class PyNames {
   public static final String INIT = "__init__";
   public static final String DICT = "__dict__";
   public static final String DOT_PY = ".py";
+  public static final String DOT_PYI = ".pyi";
   public static final String INIT_DOT_PY = INIT + DOT_PY;
+  public static final String INIT_DOT_PYI = INIT + DOT_PYI;
 
   public static final String SETUP_DOT_PY = "setup" + DOT_PY;
 

python/src/com/jetbrains/python/psi/PyUtil.java

@@ -931,9 +931,13 @@ public class PyUtil {
   public static PsiElement turnDirIntoInit(@Nullable PsiElement target) {
     if (target instanceof PsiDirectory) {
       final PsiDirectory dir = (PsiDirectory)target;
-      final PsiFile file = dir.findFile(PyNames.INIT_DOT_PY);
-      if (file != null) {
-        return file; // ResolveImportUtil will extract directory part as needed, everyone else are better off with a file.
+      final PsiFile initStub = dir.findFile(PyNames.INIT_DOT_PYI);
+      if (initStub != null) {
+        return initStub;
+      }
+      final PsiFile initFile = dir.findFile(PyNames.INIT_DOT_PY);
+      if (initFile != null) {
+        return initFile; // ResolveImportUtil will extract directory part as needed, everyone else are better off with a file.
       }
       else {
         return null;
@@ -1007,6 +1011,9 @@ public class PyUtil {
         return true;
       }
     }
+    if (directory.findFile(PyNames.INIT_DOT_PYI) != null) {
+      return true;
+    }
     if (directory.findFile(PyNames.INIT_DOT_PY) != null) {
       return true;
     }
@@ -1025,7 +1032,7 @@ public class PyUtil {
         return true;
       }
     }
-    return PyNames.INIT_DOT_PY.equals(file.getName());
+    return PyNames.INIT_DOT_PY.equals(file.getName()) || PyNames.INIT_DOT_PYI.equals(name);
   }
 
   private static boolean isSetuptoolsNamespacePackage(@NotNull PsiDirectory directory) {

python/src/com/jetbrains/python/psi/resolve/PyResolveImportUtil.kt

@@ -47,6 +47,8 @@ import com.jetbrains.python.sdk.PythonSdkType
 /**
  * Python resolve utilities for qualified names.
  *
+ * TODO: Merge with ResolveImportUtil, maybe make these functions the methods of PyQualifiedNameResolveContext.
+ *
  * @author vlan
  */
 

python/src/com/jetbrains/python/psi/resolve/ResolveImportUtil.java

@@ -397,13 +397,12 @@ public class ResolveImportUtil {
     }
 
     if (!isFileOnly) {
-      // not a subdir, not a file; could be a name in parent/__init__.py
-      final PsiFile initPy = dir.findFile(PyNames.INIT_DOT_PY);
-      if (initPy == containingFile) {
+      final PsiElement packageElement = PyUtil.getPackageElement(dir, containingFile);
+      if (packageElement == containingFile) {
         return Collections.emptyList(); // don't dive into the file we're in
       }
-      if (initPy instanceof PyFile) {
-        return ((PyFile)initPy).multiResolveName(referencedName);
+      if (packageElement instanceof PyFile) {
+        return ((PyFile)packageElement).multiResolveName(referencedName);
       }
     }
     return Collections.emptyList();
@@ -411,7 +410,10 @@ public class ResolveImportUtil {
 
   @Nullable
   private static PsiFile findPyFileInDir(PsiDirectory dir, String referencedName) {
-    PsiFile file = dir.findFile(referencedName + PyNames.DOT_PY);
+    PsiFile file = dir.findFile(referencedName + PyNames.DOT_PYI);
+    if (file == null) {
+      file = dir.findFile(referencedName + PyNames.DOT_PY);
+    }
     if (file == null) {
       final List<FileNameMatcher> associations = FileTypeManager.getInstance().getAssociations(PythonFileType.INSTANCE);
       for (FileNameMatcher association : associations) {

python/src/com/jetbrains/python/pyi/PyiFile.java

@@ -1,47 +0,0 @@
-/*
- * Copyright 2000-2015 JetBrains s.r.o.
- *
- * Licensed under the Apache License, Version 2.0 (the "License");
- * you may not use this file except in compliance with the License.
- * You may obtain a copy of the License at
- *
- * http://www.apache.org/licenses/LICENSE-2.0
- *
- * Unless required by applicable law or agreed to in writing, software
- * distributed under the License is distributed on an "AS IS" BASIS,
- * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
- * See the License for the specific language governing permissions and
- * limitations under the License.
- */
-package com.jetbrains.python.pyi;
-
-import com.intellij.openapi.fileTypes.FileType;
-import com.intellij.psi.FileViewProvider;
-import com.jetbrains.python.psi.LanguageLevel;
-import com.jetbrains.python.psi.impl.PyFileImpl;
-import org.jetbrains.annotations.NotNull;
-
-/**
- * @author vlan
- */
-public class PyiFile extends PyFileImpl {
-  public PyiFile(FileViewProvider viewProvider) {
-    super(viewProvider, PyiLanguageDialect.getInstance());
-  }
-
-  @NotNull
-  @Override
-  public FileType getFileType() {
-    return PyiFileType.INSTANCE;
-  }
-
-  @Override
-  public String toString() {
-    return "PyiFile:" + getName();
-  }
-
-  @Override
-  public LanguageLevel getLanguageLevel() {
-    return LanguageLevel.PYTHON36;
-  }
-}

python/src/com/jetbrains/python/pyi/PyiFile.kt

@@ -0,0 +1,51 @@
+/*
+ * Copyright 2000-2016 JetBrains s.r.o.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package com.jetbrains.python.pyi
+
+import com.intellij.psi.FileViewProvider
+import com.jetbrains.python.psi.LanguageLevel
+import com.jetbrains.python.psi.PyCallable
+import com.jetbrains.python.psi.impl.PyFileImpl
+import com.jetbrains.python.psi.resolve.RatedResolveResult
+import com.jetbrains.python.psi.types.TypeEvalContext
+
+/**
+ * @author vlan
+ */
+class PyiFile(viewProvider: FileViewProvider) : PyFileImpl(viewProvider, PyiLanguageDialect.getInstance()) {
+  override fun getFileType() = PyiFileType.INSTANCE
+
+  override fun toString() = "PyiFile:" + name
+
+  override fun getLanguageLevel() = LanguageLevel.PYTHON36
+
+  override fun multiResolveName(name: String): List<RatedResolveResult> {
+    val exportedResults = super.multiResolveName(name)
+        // XXX: This filter rules out names that are needed in PyClassImpl.getElementQNamed()
+//        .filter {
+//          val importedResult = it as? ImportedResolveResult
+//          val importElement = importedResult?.definer as? PyImportElement
+//          if (importElement != null) importElement.asName != null else true
+//        }
+    val firstOverload = exportedResults.firstOrNull {
+      val element = it.element
+      element is PyCallable &&
+          element.containingFile is PyiFile &&
+          PyiTypeProvider.isOverload(element, TypeEvalContext.deepCodeInsight(project))
+    }
+    return if (firstOverload != null) listOf(firstOverload) else exportedResults
+  }
+}

python/src/com/jetbrains/python/pyi/PyiFileType.java

@@ -22,7 +22,8 @@ import org.jetbrains.annotations.NotNull;
  * @author vlan
  */
 public class PyiFileType extends PythonFileType {
-  public static PythonFileType INSTANCE = new PyiFileType();
+  @NotNull
+  public static final PythonFileType INSTANCE = new PyiFileType();
 
   protected PyiFileType() {
     super(new PyiLanguageDialect());

python/src/com/jetbrains/python/pyi/PyiTypeProvider.java

@@ -191,7 +191,7 @@ public class PyiTypeProvider extends PyTypeProviderBase {
     return overloads;
   }
 
-  private static boolean isOverload(@NotNull PyCallable callable, @NotNull TypeEvalContext context) {
+  public static boolean isOverload(@NotNull PyCallable callable, @NotNull TypeEvalContext context) {
     if (callable instanceof PyDecoratable) {
       final PyDecoratable decorated = (PyDecoratable)callable;
       final ImmutableSet<PyKnownDecoratorUtil.KnownDecorator> decorators =

python/testData/completion/instanceFromProvidedCallAttr/a.after.py

@@ -1,7 +0,0 @@
-class FooMaker(object):
-    pass
-
-fm = FooMaker()
-f3 = fm()
-
-f3.bit_length()

python/testData/completion/instanceFromProvidedCallAttr/a.py

@@ -1,7 +0,0 @@
-class FooMaker(object):
-    pass
-
-fm = FooMaker()
-f3 = fm()
-
-f3.bi<caret>

python/testData/completion/instanceFromProvidedCallAttr/a.pyi

@@ -1,2 +0,0 @@
-class FooMaker:
-    def __call__(self) -> int: ...

python/testData/pyi/inspections/hiddenPyiImports/HiddenPyiImports.py

@@ -0,0 +1,7 @@
+from m1 import <error descr="Unresolved reference 'foo'">foo</error>
+from m1 import <error descr="Unresolved reference 'bar'">bar</error>
+from m1 import bar_imported
+from m1 import <error descr="Unresolved reference 'm2'">m2</error>
+from m1 import m2_imported
+
+print(foo, bar, bar_imported, m2, m2_imported)

python/testData/pyi/inspections/hiddenPyiImports/m1.pyi

@@ -0,0 +1,4 @@
+from m2 import foo
+from m2 import bar as bar_imported
+import m2
+import m2 as m2_imported

python/testData/pyi/inspections/hiddenPyiImports/m2.py

@@ -0,0 +1,3 @@
+foo = 1
+bar = 2
+baz = 3

python/testData/pyi/inspections/unresolvedClassAttributes/UnresolvedClassAttributes.py

@@ -2,9 +2,9 @@ from m1 import C
 
 c = C()
 
-print(c.class_field)
-print(c.instance_field)
-print(c.method())
+print(c.<warning descr="Unresolved attribute reference 'original_class_field' for class 'C'">original_class_field</warning>)
+print(c.<warning descr="Unresolved attribute reference 'original_instance_field' for class 'C'">original_instance_field</warning>)
+print(c.<warning descr="Unresolved attribute reference 'original_method' for class 'C'">original_method</warning>)
 
 print(c.provided_class_field)
 print(c.provided_instance_field)

python/testData/pyi/inspections/unresolvedClassAttributes/m1.py

@@ -1,8 +1,8 @@
 class C:
-    class_field = 0
+    original_class_field = 0
 
     def __init__(self):
-        self.instance_field = 1
+        self.original_instance_field = 1
 
-    def method(self):
+    def original_method(self):
         pass

python/testData/pyi/inspections/unresolvedModuleAttributes/UnresolvedModuleAttributes.py

@@ -1,11 +1,6 @@
 import m1
 
 
-print(m1.module_attr)
+print(m1.<warning descr="Cannot find reference 'module_only_attr' in 'm1.pyi'">module_only_attr</warning>)
 print(m1.provided_attr)
-print(m1.<warning descr="Cannot find reference 'not_provided_attr' in 'm1.py'">not_provided_attr</warning>)
-
-
-print(m1.<warning descr="Cannot find reference 'm2' in 'm1.py'">m2</warning>)
-print(m1.<warning descr="Cannot find reference 'm3' in 'm1.py'">m3</warning>)
-print(m1.m3_imported_as)
+print(m1.<warning descr="Cannot find reference 'not_provided_attr' in 'm1.pyi'">not_provided_attr</warning>)

python/testData/pyi/inspections/unresolvedModuleAttributes/m1.py

@@ -1 +1 @@
-module_attr = 0
+module_only_attr = 0

python/testData/pyi/inspections/unresolvedModuleAttributes/m1.pyi

@@ -1,5 +1 @@
-import m2
-import m3 as m3_imported_as
-
-
 provided_attr = ...  # int

python/testData/pyi/inspections/unresolvedModuleAttributes/m2.py

@@ -1 +0,0 @@
-

python/testData/pyi/resolve/genericAttribute/GenericAttribute.py

@@ -0,0 +1,4 @@
+from m1 import C, Concrete
+
+C(Concrete()).get().foo
+#                   <ref>

python/testData/pyi/resolve/genericAttribute/m1.pyi

@@ -0,0 +1,13 @@
+from typing import Generic, TypeVar, Any
+
+
+T = TypeVar('T')
+
+
+class Concrete:
+    foo = ...  # type: Any
+
+
+class C(Generic[T]):
+    def __init__(self, x: T): ...
+    def get(self) -> T: ...

python/testData/pyi/resolve/genericAttribute/typing.pyi

@@ -0,0 +1,7 @@
+class Generic: ...
+
+
+def TypeVar(x): ...
+
+
+class Any: ...

python/testData/pyi/resolve/moduleAttributePyiOverPy/ModuleAttributePyiOverPy.py

@@ -0,0 +1,2 @@
+from m1 import foo
+#              <ref>

python/testData/pyi/resolve/moduleAttributePyiOverPy/m1.py

@@ -0,0 +1 @@
+foo = 0

python/testData/pyi/resolve/moduleAttributePyiOverPy/m1.pyi

@@ -0,0 +1 @@
+foo = ...  # type: int

python/testSrc/com/jetbrains/python/PythonCompletionTest.java

@@ -1002,11 +1002,6 @@ public class PythonCompletionTest extends PyTestCase {
     doTest();
   }
 
-  // PY-12425
-  public void testInstanceFromProvidedCallAttr() {
-    doMultiFileTest();
-  }
-
   // PY-18684
   public void testRPowSignature() {
     doTest();

python/testSrc/com/jetbrains/python/pyi/PyiInspectionsTest.java

@@ -48,6 +48,12 @@ public class PyiInspectionsTest extends PyTestCase {
     doPyTest(PyUnresolvedReferencesInspection.class);
   }
 
+  // TODO: Enable when the issue with the visibility of imports in Python stubs is fixed
+  @SuppressWarnings("unused")
+  public void _testHiddenPyiImports() {
+    doPyTest(PyUnresolvedReferencesInspection.class);
+  }
+
   public void testUnresolvedClassAttributes() {
     doPyTest(PyUnresolvedReferencesInspection.class);
   }

python/testSrc/com/jetbrains/python/pyi/PyiResolveTest.java

@@ -15,6 +15,8 @@
  */
 package com.jetbrains.python.pyi;
 
+import com.intellij.psi.PsiElement;
+import com.intellij.psi.PsiFile;
 import com.jetbrains.python.PythonTestUtil;
 import com.jetbrains.python.fixtures.PyMultiFileResolveTestCase;
 import com.jetbrains.python.psi.PyClass;
@@ -44,4 +46,19 @@ public class PyiResolveTest extends PyMultiFileResolveTestCase {
   public void testModuleAttribute() {
     assertResolvesTo(PyTargetExpression.class, "foo");
   }
+
+  // PY-21231
+  public void testModuleAttributePyiOverPy() {
+    final PsiElement result = doResolve();
+    assertInstanceOf(result, PyTargetExpression.class);
+    final PyTargetExpression target = (PyTargetExpression)result;
+    assertEquals("foo", target.getName());
+    final PsiFile containingFile = target.getContainingFile();
+    assertInstanceOf(containingFile, PyiFile.class);
+  }
+
+  // PY-21231
+  public void testGenericAttribute() {
+    assertResolvesTo(PyTargetExpression.class, "foo");
+  }
 }