if the same name is declared in the 'try' and 'except' parts of a try/except statement, prefer the one in the 'try' part (fixes PyGame: PY-2197)

python/src/com/jetbrains/python/PyElementTypes.java

@@ -31,7 +31,7 @@ public interface PyElementTypes {
   PyStubElementType<PyImportElementStub, PyImportElement> IMPORT_ELEMENT = new PyImportElementElementType();
   
   PyElementType STAR_IMPORT_ELEMENT = new PyElementType("STAR_IMPORT_ELEMENT", PyStarImportElementImpl.class);
-  PyElementType EXCEPT_PART = new PyElementType("EXCEPT_PART", PyExceptPartImpl.class);
+  PyStubElementType<PyExceptPartStub, PyExceptPart> EXCEPT_PART = new PyExceptPartElementType();
   PyElementType PRINT_TARGET = new PyElementType("PRINT_TARGET", PyPrintTargetImpl.class);
   PyElementType DECORATOR = new PyElementType("DECORATOR", PyDecoratorImpl.class);
 

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

@@ -1,13 +1,18 @@
 package com.jetbrains.python.psi;
 
+import com.intellij.psi.StubBasedPsiElement;
+import com.jetbrains.python.psi.stubs.PyExceptPartStub;
 import org.jetbrains.annotations.Nullable;
 
 /**
  * @author dcheryasov
  */
-public interface PyExceptPart extends PyElement, NameDefiner, PyStatementPart {
-    PyExceptPart[] EMPTY_ARRAY = new PyExceptPart[0];
+public interface PyExceptPart extends PyElement, StubBasedPsiElement<PyExceptPartStub>, NameDefiner, PyStatementPart {
+  PyExceptPart[] EMPTY_ARRAY = new PyExceptPart[0];
 
-    @Nullable PyExpression getExceptClass();
-    @Nullable PyExpression getTarget();
+  @Nullable
+  PyExpression getExceptClass();
+
+  @Nullable
+  PyExpression getTarget();
 }

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

@@ -41,7 +41,7 @@ public class PyFileElementType extends IStubFileElementType<PyFileStub> {
 
   @Override
   public int getStubVersion() {
-    return 31;
+    return 32;
   }
 
   @Override

python/src/com/jetbrains/python/psi/impl/PyCallExpressionHelper.java

@@ -544,7 +544,10 @@ public class PyCallExpressionHelper {
         // check length of myTupleArg
         PyType tuple_arg_type = null;
         if (type_context != null) {
-          tuple_arg_type = type_context.getType(PsiTreeUtil.getChildOfType(myTupleArg, PyExpression.class));
+          final PyExpression expression = PsiTreeUtil.getChildOfType(myTupleArg, PyExpression.class);
+          if (expression != null) {
+            tuple_arg_type = type_context.getType(expression);
+          }
         }
         int tuple_length;
         boolean tuple_length_known;

python/src/com/jetbrains/python/psi/impl/PyExceptPartImpl.java

@@ -3,6 +3,7 @@ package com.jetbrains.python.psi.impl;
 import com.intellij.lang.ASTNode;
 import com.jetbrains.python.PyElementTypes;
 import com.jetbrains.python.psi.*;
+import com.jetbrains.python.psi.stubs.PyExceptPartStub;
 import org.jetbrains.annotations.NotNull;
 import org.jetbrains.annotations.Nullable;
 
@@ -11,11 +12,15 @@ import java.util.ArrayList;
 /**
  * @author dcheryasov
  */
-public class PyExceptPartImpl extends PyElementImpl implements PyExceptPart {
+public class PyExceptPartImpl extends PyBaseElementImpl<PyExceptPartStub> implements PyExceptPart {
   public PyExceptPartImpl(ASTNode astNode) {
     super(astNode);
   }
 
+  public PyExceptPartImpl(PyExceptPartStub stub) {
+    super(stub, PyElementTypes.EXCEPT_PART);
+  }
+
   @Override
   protected void acceptPyVisitor(PyElementVisitor pyVisitor) {
     pyVisitor.visitPyExceptBlock(this);

python/src/com/jetbrains/python/psi/impl/stubs/PyExceptPartElementType.java

@@ -0,0 +1,46 @@
+package com.jetbrains.python.psi.impl.stubs;
+
+import com.intellij.lang.ASTNode;
+import com.intellij.psi.PsiElement;
+import com.intellij.psi.stubs.StubElement;
+import com.intellij.psi.stubs.StubInputStream;
+import com.intellij.psi.stubs.StubOutputStream;
+import com.jetbrains.python.psi.PyExceptPart;
+import com.jetbrains.python.psi.PyStubElementType;
+import com.jetbrains.python.psi.impl.PyExceptPartImpl;
+import com.jetbrains.python.psi.stubs.PyExceptPartStub;
+
+import java.io.IOException;
+
+/**
+ * @author yole
+ */
+public class PyExceptPartElementType extends PyStubElementType<PyExceptPartStub, PyExceptPart> {
+  public PyExceptPartElementType() {
+    super("EXCEPT_PART");
+  }
+
+  @Override
+  public PsiElement createElement(ASTNode node) {
+    return new PyExceptPartImpl(node);
+  }
+
+  @Override
+  public PyExceptPart createPsi(PyExceptPartStub stub) {
+    return new PyExceptPartImpl(stub);
+  }
+
+  @Override
+  public PyExceptPartStub createStub(PyExceptPart psi, StubElement parentStub) {
+    return new PyExceptPartStubImpl(parentStub);
+  }
+
+  @Override
+  public void serialize(PyExceptPartStub stub, StubOutputStream dataStream) throws IOException {
+  }
+
+  @Override
+  public PyExceptPartStub deserialize(StubInputStream dataStream, StubElement parentStub) throws IOException {
+    return new PyExceptPartStubImpl(parentStub);
+  }
+}

python/src/com/jetbrains/python/psi/impl/stubs/PyExceptPartStubImpl.java

@@ -0,0 +1,16 @@
+package com.jetbrains.python.psi.impl.stubs;
+
+import com.intellij.psi.stubs.StubBase;
+import com.intellij.psi.stubs.StubElement;
+import com.jetbrains.python.PyElementTypes;
+import com.jetbrains.python.psi.PyExceptPart;
+import com.jetbrains.python.psi.stubs.PyExceptPartStub;
+
+/**
+ * @author yole
+ */
+public class PyExceptPartStubImpl extends StubBase<PyExceptPart> implements PyExceptPartStub {
+  protected PyExceptPartStubImpl(final StubElement parent) {
+    super(parent, PyElementTypes.EXCEPT_PART);
+  }
+}

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

@@ -581,7 +581,10 @@ public class ResolveImportUtil {
       // OTOH, quite often a module named foo exports a class or function named foo, which is used as a fallback
       // by a module one level higher (e.g. curses.set_key). Prefer it to submodule if possible.
       ret = ((PyFile)parent).getElementNamed(referencedName);
-      if (ret != null && !PyUtil.instanceOf(ret, PsiFile.class, PsiDirectory.class)) return ret;
+      if (ret != null && !PyUtil.instanceOf(ret, PsiFile.class, PsiDirectory.class) &&
+          PsiTreeUtil.getStubOrPsiParentOfType(ret, PyExceptPart.class) == null) {
+        return ret;
+      }
       if (possible_ret != null) return possible_ret;
     }
     else if (parent instanceof PsiDirectory) {

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

@@ -53,26 +53,26 @@ public class ResolveProcessor implements PsiScopeProcessor {
       final VirtualFile file = ((PyFile)element).getVirtualFile();
       if (file != null) {
         if (myName.equals(file.getNameWithoutExtension())) {
-          return setResult(element);
+          return setResult(element, null);
         }
         else if (PyNames.INIT_DOT_PY.equals(file.getName())) {
           VirtualFile dir = file.getParent();
           if ((dir != null) && myName.equals(dir.getName())) {
-            return setResult(element);
+            return setResult(element, null);
           }
         }
       }
     }
     else if (element instanceof PsiNamedElement) {
       if (myName.equals(((PsiNamedElement)element).getName())) {
-        return setResult(element);
+        return setResult(element, null);
       }
     }
     else if (element instanceof PyReferenceExpression) {
       PyReferenceExpression expr = (PyReferenceExpression)element;
       String referencedName = expr.getReferencedName();
       if (referencedName != null && referencedName.equals(myName)) {
-        return setResult(element);
+        return setResult(element, null);
       }
     }
     else if (element instanceof NameDefiner) {
@@ -85,13 +85,15 @@ public class ResolveProcessor implements PsiScopeProcessor {
           return false;
         }
 
-        setResult(by_name);
+        setResult(by_name, definer);
         if (!PsiTreeUtil.isAncestor(element, by_name, true)) {
           addNameDefiner(definer);
         }
         // we can have same module imported directly and as part of chain (import os; import os.path)
         // direct imports always take precedence over imported modules
-        if (!(myResult instanceof PyImportedModule)) {
+        // also, if some name is defined both in 'try' and 'except' parts of the same try/except statement,
+        // we prefer the declaration in the 'try' part
+        if (!(myResult instanceof PyImportedModule) && PsiTreeUtil.getParentOfType(element, PyExceptPart.class) == null) {
           return false;
         }
       }
@@ -122,8 +124,8 @@ public class ResolveProcessor implements PsiScopeProcessor {
   public void handleEvent(Event event, Object associated) {
   }
 
-  private boolean setResult(PsiElement result) {
-    if (myResult == null || getScope(myResult) == getScope(result)) {
+  private boolean setResult(PsiElement result, @Nullable PsiElement definer) {
+    if (myResult == null || getScope(myResult) == getScope(result) || (definer != null && getScope(myResult) == getScope(definer))) {
       myResult = result;
     }
     return false;

python/src/com/jetbrains/python/psi/stubs/PyExceptPartStub.java

@@ -0,0 +1,10 @@
+package com.jetbrains.python.psi.stubs;
+
+import com.intellij.psi.stubs.StubElement;
+import com.jetbrains.python.psi.PyExceptPart;
+
+/**
+ * @author yole
+ */
+public interface PyExceptPartStub extends StubElement<PyExceptPart> {
+}

python/testData/resolve/ImportInTryExcept.py

@@ -0,0 +1,7 @@
+try:
+    import sys
+except:
+    sys = None
+
+print sys
+#      <ref>

python/testData/stubs/ImportInTryExcept.py

@@ -0,0 +1,4 @@
+try:
+    import sys
+except:
+    sys = None

python/testSrc/com/jetbrains/python/PyMultiFileResolveTest.java

@@ -215,6 +215,10 @@ public class PyMultiFileResolveTest extends PyResolveTestCase {
     assertResolvesTo(PyFunction.class, "do_stuff", "/src/mypackage1.py");
   }
 
+  public void testImportPackageIntoSelf() {
+    assertResolvesTo(PyFunction.class, "foo", "/src/mygame/display.py");
+  }
+
   private PsiFile prepareFile() throws Exception {
     String testName = getTestName(true);
     String fileName = getTestName(false) + ".py";
@@ -236,7 +240,7 @@ public class PyMultiFileResolveTest extends PyResolveTestCase {
   protected PsiElement doResolve() throws Exception {
     PsiFile psiFile = prepareFile();
     int offset = findMarkerOffset(psiFile);
-    final PsiReference ref = psiFile.findReferenceAt(offset);
+    final PsiPolyVariantReference ref = (PsiPolyVariantReference) psiFile.findReferenceAt(offset);
     final PsiManagerImpl psiManager = (PsiManagerImpl)myFixture.getPsiManager();
     psiManager.setAssertOnFileLoadingFilter(new VirtualFileFilter() {
       @Override
@@ -246,7 +250,11 @@ public class PyMultiFileResolveTest extends PyResolveTestCase {
       }
     });
     try {
-      return ref.resolve();
+      final ResolveResult[] resolveResults = ref.multiResolve(false);
+      if (resolveResults.length == 0) {
+        return null;
+      }
+      return resolveResults[0].isValidResult() ? resolveResults [0].getElement() : null;
     }
     finally {
       psiManager.setAssertOnFileLoadingFilter(VirtualFileFilter.NONE);

python/testSrc/com/jetbrains/python/PyResolveTest.java

@@ -375,4 +375,8 @@ public class PyResolveTest extends PyResolveTestCase {
   public void testLambdaToClass() {  // PY-2182
     assertResolvesTo(PyClass.class, "TestTwo");
   }
+
+  public void testImportInTryExcept() {  // PY-2197
+    assertResolvesTo(PyFile.class, "sys.py");
+  }
 }

python/testSrc/com/jetbrains/python/PyStubsTest.java

@@ -248,6 +248,13 @@ public class PyStubsTest extends PyLightFixtureTestCase {
     assertNotParsed(file);
   }
 
+  public void testImportInTryExcept() {
+    final PyFileImpl file = (PyFileImpl) getTestFile();
+    final PsiElement element = file.findExportedName("sys");
+    assertTrue(element != null ? element.toString() : "null", element instanceof PyImportElement);
+    assertNotParsed(file);
+  }
+
   private PyFile getTestFile() {
     return getTestFile(getTestName(false) + ".py");
   }