diff --git a/python/src/com/jetbrains/python/psi/impl/PyImportElementImpl.java b/python/src/com/jetbrains/python/psi/impl/PyImportElementImpl.java index 308bb256bd8c..a54a712585eb 100644 --- a/python/src/com/jetbrains/python/psi/impl/PyImportElementImpl.java +++ b/python/src/com/jetbrains/python/psi/impl/PyImportElementImpl.java @@ -219,7 +219,10 @@ public class PyImportElementImpl extends PyBaseElementImpl if (qName == null || qName.getComponentCount() == 0 || !qName.getComponents().get(0).equals(the_name)) { return null; } - return ResolveImportUtil.resolveImportElement(this, PyQualifiedName.fromComponents(the_name)); + if (qName.getComponentCount() == 1) { + return ResolveImportUtil.resolveImportElement(this, PyQualifiedName.fromComponents(the_name)); + } + return new PyImportedModule((PyFile) getContainingFile(), PyQualifiedName.fromComponents(the_name)); } } diff --git a/python/src/com/jetbrains/python/psi/impl/PyImportedModule.java b/python/src/com/jetbrains/python/psi/impl/PyImportedModule.java index 276bbd815589..e81836a9c4b5 100644 --- a/python/src/com/jetbrains/python/psi/impl/PyImportedModule.java +++ b/python/src/com/jetbrains/python/psi/impl/PyImportedModule.java @@ -4,11 +4,10 @@ import com.intellij.psi.PsiElement; import com.intellij.psi.PsiElementVisitor; import com.intellij.psi.impl.light.LightElement; import com.jetbrains.python.PythonLanguage; -import com.jetbrains.python.psi.NameDefiner; -import com.jetbrains.python.psi.PyElement; -import com.jetbrains.python.psi.PyFile; -import com.jetbrains.python.psi.PyImportElement; +import com.jetbrains.python.psi.*; +import com.jetbrains.python.psi.resolve.ResolveImportUtil; import org.jetbrains.annotations.NotNull; +import org.jetbrains.annotations.Nullable; import java.util.List; @@ -25,6 +24,14 @@ public class PyImportedModule extends LightElement implements NameDefiner { myImportedPrefix = importedPrefix; } + public PyFile getContainingFile() { + return myContainingFile; + } + + public PyQualifiedName getImportedPrefix() { + return myImportedPrefix; + } + @NotNull public Iterable iterateNames() { throw new UnsupportedOperationException(); @@ -32,29 +39,27 @@ public class PyImportedModule extends LightElement implements NameDefiner { public PsiElement getElementNamed(String the_name) { PyQualifiedName prefix = myImportedPrefix.append(the_name); - final List imports = ((PyFileImpl)myContainingFile).getImportTargets(); - for (PyImportElement anImport : imports) { - final PyQualifiedName qName = anImport.getImportedQName(); - if (qName != null && matchesPrefix(qName, prefix)) { - if (qName.getComponentCount() == prefix.getComponentCount()) { - return anImport; - } - return new PyImportedModule(myContainingFile, prefix); + final PyImportElement importElement = findMatchingImportElement(prefix); + if (importElement != null) { + final PyQualifiedName qName = importElement.getImportedQName(); + if (qName != null && qName.getComponentCount() == prefix.getComponentCount()) { + return resolve(importElement, prefix); } + return new PyImportedModule(myContainingFile, prefix); } return null; } - private static boolean matchesPrefix(PyQualifiedName qName, PyQualifiedName prefix) { - if (qName.getComponentCount() < prefix.getComponentCount()) { - return false; - } - for (int i = 0; i < prefix.getComponentCount(); i++) { - if (!qName.getComponents().get(i).equals(prefix.getComponents().get(i))) { - return false; + @Nullable + private PyImportElement findMatchingImportElement(PyQualifiedName prefix) { + final List imports = ((PyFileImpl)myContainingFile).getImportTargets(); + for (PyImportElement anImport : imports) { + final PyQualifiedName qName = anImport.getImportedQName(); + if (qName != null && qName.matchesPrefix(prefix)) { + return anImport; } } - return true; + return null; } public boolean mustResolveOutside() { @@ -77,4 +82,22 @@ public class PyImportedModule extends LightElement implements NameDefiner { public String toString() { return "PyImportedModule:" + myImportedPrefix; } + + @NotNull + @Override + public PsiElement getNavigationElement() { + final PyImportElement importElement = findMatchingImportElement(myImportedPrefix); + if (importElement != null) { + final PsiElement element = resolve(importElement, myImportedPrefix); + if (element != null) { + return element; + } + } + return super.getNavigationElement(); + } + + @Nullable + private PsiElement resolve(PyImportElement importElement, final PyQualifiedName prefix) { + return PyUtil.turnDirIntoInit(ResolveImportUtil.resolveImportElement(importElement, prefix)); + } } diff --git a/python/src/com/jetbrains/python/psi/impl/PyQualifiedName.java b/python/src/com/jetbrains/python/psi/impl/PyQualifiedName.java index 22ebc448de5f..dfc50b79796f 100644 --- a/python/src/com/jetbrains/python/psi/impl/PyQualifiedName.java +++ b/python/src/com/jetbrains/python/psi/impl/PyQualifiedName.java @@ -1,5 +1,6 @@ package com.jetbrains.python.psi.impl; +import com.intellij.openapi.util.text.StringUtil; import com.intellij.psi.stubs.StubInputStream; import com.intellij.psi.stubs.StubOutputStream; import com.intellij.util.io.StringRef; @@ -61,6 +62,18 @@ public class PyQualifiedName { return true; } + public boolean matchesPrefix(PyQualifiedName prefix) { + if (getComponentCount() < prefix.getComponentCount()) { + return false; + } + for (int i = 0; i < prefix.getComponentCount(); i++) { + if (!getComponents().get(i).equals(prefix.getComponents().get(i))) { + return false; + } + } + return true; + } + public static void serialize(@Nullable PyQualifiedName qName, StubOutputStream dataStream) throws IOException { if (qName == null) { dataStream.writeVarInt(0); @@ -97,4 +110,9 @@ public class PyQualifiedName { } return myComponents.get(myComponents.size()-1); } + + @Override + public String toString() { + return StringUtil.join(myComponents, "."); + } } diff --git a/python/src/com/jetbrains/python/psi/impl/PyReferenceExpressionImpl.java b/python/src/com/jetbrains/python/psi/impl/PyReferenceExpressionImpl.java index f6bcdc32572c..ceb8b515a0b6 100644 --- a/python/src/com/jetbrains/python/psi/impl/PyReferenceExpressionImpl.java +++ b/python/src/com/jetbrains/python/psi/impl/PyReferenceExpressionImpl.java @@ -241,6 +241,9 @@ public class PyReferenceExpressionImpl extends PyElementImpl implements PyRefere if (target instanceof PyFile) { return new PyModuleType((PyFile) target); } + if (target instanceof PyImportedModule) { + return new PyImportedModuleType((PyImportedModule) target); + } if (target instanceof PyTargetExpression && context.allowDataFlow() && anchor != null) { final ScopeOwner scopeOwner = PsiTreeUtil.getParentOfType(anchor, ScopeOwner.class); if (scopeOwner != null && scopeOwner == PsiTreeUtil.getParentOfType(target, ScopeOwner.class)) { diff --git a/python/src/com/jetbrains/python/psi/resolve/ResolveImportUtil.java b/python/src/com/jetbrains/python/psi/resolve/ResolveImportUtil.java index beb2d3be0dda..58a967d996f7 100644 --- a/python/src/com/jetbrains/python/psi/resolve/ResolveImportUtil.java +++ b/python/src/com/jetbrains/python/psi/resolve/ResolveImportUtil.java @@ -68,6 +68,9 @@ public class ResolveImportUtil { */ @Nullable public static PsiDirectory stepBackFrom(PsiFile base, int depth) { + if (depth == 0) { + return base.getContainingDirectory(); + } PsiDirectory result; if (base != null) { base = base.getOriginalFile(); // just to make sure diff --git a/python/src/com/jetbrains/python/psi/types/PyImportedModuleType.java b/python/src/com/jetbrains/python/psi/types/PyImportedModuleType.java new file mode 100644 index 000000000000..f18b8acc8cbb --- /dev/null +++ b/python/src/com/jetbrains/python/psi/types/PyImportedModuleType.java @@ -0,0 +1,50 @@ +package com.jetbrains.python.psi.types; + +import com.intellij.codeInsight.lookup.LookupElement; +import com.intellij.codeInsight.lookup.LookupElementBuilder; +import com.intellij.psi.PsiElement; +import com.intellij.util.ProcessingContext; +import com.jetbrains.python.psi.PyImportElement; +import com.jetbrains.python.psi.PyQualifiedExpression; +import com.jetbrains.python.psi.impl.PyFileImpl; +import com.jetbrains.python.psi.impl.PyImportedModule; +import com.jetbrains.python.psi.impl.PyQualifiedName; + +import java.util.ArrayList; +import java.util.List; + +/** + * @author yole + */ +public class PyImportedModuleType implements PyType { + private PyImportedModule myImportedModule; + + public PyImportedModuleType(PyImportedModule importedModule) { + myImportedModule = importedModule; + } + + public PsiElement resolveMember(String name) { + return myImportedModule.getElementNamed(name); + } + + public Object[] getCompletionVariants(PyQualifiedExpression referenceExpression, ProcessingContext context) { + List result = new ArrayList(); + final List importTargets = ((PyFileImpl)myImportedModule.getContainingFile()).getImportTargets(); + final int imported = myImportedModule.getImportedPrefix().getComponentCount(); + for (PyImportElement importTarget : importTargets) { + final PyQualifiedName qName = importTarget.getImportedQName(); + if (qName != null && qName.matchesPrefix(myImportedModule.getImportedPrefix())) { + final List components = qName.getComponents(); + if (components.size() > imported) { + String module = components.get(imported); + result.add(LookupElementBuilder.create(module)); + } + } + } + return result.toArray(new Object[result.size()]); + } + + public String getName() { + return "PyImportedModuleType:" + myImportedModule.toString(); + } +} diff --git a/python/testData/completion/py874.after.py b/python/testData/completion/py874.after.py new file mode 100644 index 000000000000..39dd3f1cfa93 --- /dev/null +++ b/python/testData/completion/py874.after.py @@ -0,0 +1,2 @@ +import root.nested_mod +root.nested_mod diff --git a/python/testData/completion/py874.py b/python/testData/completion/py874.py new file mode 100644 index 000000000000..75468149c486 --- /dev/null +++ b/python/testData/completion/py874.py @@ -0,0 +1,2 @@ +import root.nested_mod +root. diff --git a/python/testData/completion/root/__init__.py b/python/testData/completion/root/__init__.py new file mode 100644 index 000000000000..e69de29bb2d1 diff --git a/python/testData/completion/root/nested_mod.py b/python/testData/completion/root/nested_mod.py new file mode 100644 index 000000000000..e69de29bb2d1 diff --git a/python/testSrc/com/jetbrains/python/PyMultiFileResolveTest.java b/python/testSrc/com/jetbrains/python/PyMultiFileResolveTest.java index 28a3c73c35e2..7610bae14c50 100644 --- a/python/testSrc/com/jetbrains/python/PyMultiFileResolveTest.java +++ b/python/testSrc/com/jetbrains/python/PyMultiFileResolveTest.java @@ -168,7 +168,8 @@ public class PyMultiFileResolveTest extends PyLightFixtureTestCase { public void testNestedPackageElement() throws Exception { PsiElement element = doResolve(); - assertTrue(element instanceof PyFile); + element = element.getNavigationElement(); + assertInstanceOf(element, PyFile.class); assertEquals("__init__.py", ((PyFile) element).getName()); } diff --git a/python/testSrc/com/jetbrains/python/PythonCompletionTest.java b/python/testSrc/com/jetbrains/python/PythonCompletionTest.java index 0f3023400bc4..9795702315d2 100644 --- a/python/testSrc/com/jetbrains/python/PythonCompletionTest.java +++ b/python/testSrc/com/jetbrains/python/PythonCompletionTest.java @@ -95,4 +95,13 @@ public class PythonCompletionTest extends PyLightFixtureTestCase { myFixture.completeBasic(); myFixture.checkResultByFile(testName + ".after.py"); } + + public void testPy874() throws Exception { + final String dirname = "completion/"; + final String testName = dirname + "py874"; + myFixture.configureByFile(testName + ".py"); + myFixture.copyDirectoryToProject(dirname + "root", dirname + "root"); + myFixture.completeBasic(); + myFixture.checkResultByFile(testName + ".after.py"); + } } \ No newline at end of file