From 60fdbc241da3925292afcd996f720ccbe1838408 Mon Sep 17 00:00:00 2001
From: Andrey Vlasovskikh <andrey.vlasovskikh@jetbrains.com>
Date: Mon, 13 Jul 2015 14:19:05 +0300
Subject: [PATCH] Resolve module attributes in Python stub files (PY-16402)

---
 python/src/META-INF/python-core.xml           |  1 +
 .../python/pyi/PyiModuleMembersProvider.java  | 66 +++++++++++++++++++
 .../moduleAttribute/ModuleAttribute.py        |  5 ++
 .../pyi/resolve/moduleAttribute/m1.py         |  0
 .../pyi/resolve/moduleAttribute/m1.pyi        |  1 +
 .../jetbrains/python/pyi/PyiResolveTest.java  |  5 ++
 6 files changed, 78 insertions(+)
 create mode 100644 python/src/com/jetbrains/python/pyi/PyiModuleMembersProvider.java
 create mode 100644 python/testData/pyi/resolve/moduleAttribute/ModuleAttribute.py
 create mode 100644 python/testData/pyi/resolve/moduleAttribute/m1.py
 create mode 100644 python/testData/pyi/resolve/moduleAttribute/m1.pyi

diff --git a/python/src/META-INF/python-core.xml b/python/src/META-INF/python-core.xml
index 6ca83b0cf881..419813822649 100644
--- a/python/src/META-INF/python-core.xml
+++ b/python/src/META-INF/python-core.xml
@@ -630,6 +630,7 @@
     <!-- typing -->
     <typeProvider implementation="com.jetbrains.python.codeInsight.PyTypingTypeProvider"/>
     <typeProvider implementation="com.jetbrains.python.pyi.PyiTypeProvider"/>
+    <pyModuleMembersProvider implementation="com.jetbrains.python.pyi.PyiModuleMembersProvider"/>
 
     <typeProvider implementation="com.jetbrains.python.debugger.PyCallSignatureTypeProvider"/>
     <pyReferenceResolveProvider implementation="com.jetbrains.python.psi.resolve.PythonBuiltinReferenceResolveProvider"/>
diff --git a/python/src/com/jetbrains/python/pyi/PyiModuleMembersProvider.java b/python/src/com/jetbrains/python/pyi/PyiModuleMembersProvider.java
new file mode 100644
index 000000000000..28fb341f06b1
--- /dev/null
+++ b/python/src/com/jetbrains/python/pyi/PyiModuleMembersProvider.java
@@ -0,0 +1,66 @@
+/*
+ * 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.psi.PsiElement;
+import com.jetbrains.python.codeInsight.PyCustomMember;
+import com.jetbrains.python.psi.NameDefiner;
+import com.jetbrains.python.psi.PyElement;
+import com.jetbrains.python.psi.PyFile;
+import com.jetbrains.python.psi.resolve.PointInImport;
+import com.jetbrains.python.psi.types.PyModuleMembersProvider;
+import org.jetbrains.annotations.Nullable;
+
+import java.util.ArrayList;
+import java.util.Collection;
+import java.util.Collections;
+import java.util.List;
+
+/**
+ * @author vlan
+ */
+public class PyiModuleMembersProvider extends PyModuleMembersProvider {
+  @Nullable
+  @Override
+  public PsiElement resolveMember(PyFile module, String name) {
+    final PsiElement pythonStub = PyiUtil.getPythonStub(module);
+    if (pythonStub instanceof NameDefiner) {
+      return ((NameDefiner)pythonStub).getElementNamed(name);
+    }
+    return null;
+  }
+
+  @Override
+  public Collection<PyCustomMember> getMembers(PyFile module, PointInImport point) {
+    final PsiElement pythonStub = PyiUtil.getPythonStub(module);
+    if (pythonStub instanceof NameDefiner) {
+      final List<PyCustomMember> results = new ArrayList<PyCustomMember>();
+      for (PyElement element : ((NameDefiner)pythonStub).iterateNames()) {
+        final String name = element.getName();
+        if (name != null) {
+          results.add(new PyCustomMember(name, element));
+        }
+      }
+      return results;
+    }
+    return Collections.emptyList();
+  }
+
+  @Override
+  protected Collection<PyCustomMember> getMembersByQName(PyFile module, String qName) {
+    return null;
+  }
+}
diff --git a/python/testData/pyi/resolve/moduleAttribute/ModuleAttribute.py b/python/testData/pyi/resolve/moduleAttribute/ModuleAttribute.py
new file mode 100644
index 000000000000..609db17ed3de
--- /dev/null
+++ b/python/testData/pyi/resolve/moduleAttribute/ModuleAttribute.py
@@ -0,0 +1,5 @@
+import m1
+
+
+print(m1.foo)
+#        <ref>
diff --git a/python/testData/pyi/resolve/moduleAttribute/m1.py b/python/testData/pyi/resolve/moduleAttribute/m1.py
new file mode 100644
index 000000000000..e69de29bb2d1
diff --git a/python/testData/pyi/resolve/moduleAttribute/m1.pyi b/python/testData/pyi/resolve/moduleAttribute/m1.pyi
new file mode 100644
index 000000000000..7fa904692fbb
--- /dev/null
+++ b/python/testData/pyi/resolve/moduleAttribute/m1.pyi
@@ -0,0 +1 @@
+foo = ...  # int
diff --git a/python/testSrc/com/jetbrains/python/pyi/PyiResolveTest.java b/python/testSrc/com/jetbrains/python/pyi/PyiResolveTest.java
index dc9732b68823..2c0b08f6983d 100644
--- a/python/testSrc/com/jetbrains/python/pyi/PyiResolveTest.java
+++ b/python/testSrc/com/jetbrains/python/pyi/PyiResolveTest.java
@@ -18,6 +18,7 @@ package com.jetbrains.python.pyi;
 import com.jetbrains.python.PythonTestUtil;
 import com.jetbrains.python.fixtures.PyMultiFileResolveTestCase;
 import com.jetbrains.python.psi.PyClass;
+import com.jetbrains.python.psi.PyTargetExpression;
 
 /**
  * @author vlan
@@ -39,4 +40,8 @@ public class PyiResolveTest extends PyMultiFileResolveTestCase {
   public void testFromPyiToClassInPy() {
     assertResolvesTo(PyClass.class, "C");
   }
+
+  public void testModuleAttribute() {
+    assertResolvesTo(PyTargetExpression.class, "foo");
+  }
 }