Fixes PY-48, PY-44, PY-28 (for this no test yet).

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

@@ -31,11 +31,9 @@ import org.jetbrains.annotations.Nullable;
 import java.util.*;
 
 /**
- * Created by IntelliJ IDEA.
+ * Ref resolution routines.
  * User: yole
  * Date: 14.06.2005
- * Time: 23:45:32
- * To change this template use File | Settings | File Templates.
  */
 public class PyResolveUtil {
 
@@ -110,7 +108,7 @@ public class PyResolveUtil {
    * Crawls up the PSI tree, checking nodes as if crawling backwards through source lexemes.
    * @param processor a visitor that says when the crawl is done and collects info.
    * @param elt element from which we start (not checked by processor); if null, the search immediately fails. 
-   * @param fromunder if true, search not above elt, but from an [possibly imaginary] node right below elt; so elt gets analyzed, too.
+   * @param fromunder if true, search not above elt, but from a [possibly imaginary] node right below elt; so elt gets analyzed, too.
    * @return first element that the processor accepted.
    */
   @Nullable
@@ -146,9 +144,8 @@ public class PyResolveUtil {
         if ((cap != null) && PsiTreeUtil.isAncestor(local_cap, cap, true)) break; // seeker is in a context above elt's
       }
       // maybe we're capped by a class
-      PsiElement possible_class_cap = getConcealingParent(seeker);
-      if (possible_class_cap instanceof PyClass) continue; // class implicitly qualifies things, and we're looking for unqualified.
-      // check
+      if (refersFromMethodToClass(cap, seeker)) continue;
+      // check what we got
       if (seeker != null) {
         if (!processor.execute(seeker, ResolveState.initial())) {
           if (processor instanceof ResolveProcessor) {
@@ -161,6 +158,45 @@ public class PyResolveUtil {
     return null;
   }
 
+  @Nullable
+  public static PsiElement resolveOffContext(@NotNull PyReferenceExpression refex) {
+    // if we're under a cap, an external object that we want to use might be also defined below us.
+    // look through all contexts, closest first.
+    PsiElement ret = null;
+    PsiElement our_cap = getConcealingParent(refex);
+    ResolveProcessor proc = new ResolveProcessor(refex.getReferencedName()); // processor reusable till first hit
+    if (our_cap != null) {
+      PsiElement cap = our_cap;
+      while (true) {
+        cap = getConcealingParent(cap);
+        if (cap == null) cap = refex.getContainingFile();
+        ret = treeCrawlUp(proc, cap, true);
+        if ((ret != null) && !PsiTreeUtil.isAncestor(our_cap, ret, true)) { // found something and it is below our cap
+          // maybe we're in a method, and what we found is in its class context?
+          if (! refersFromMethodToClass(our_cap, ret)) {
+            break; // not in method -> must be all right
+          }
+        }
+        if (cap instanceof PsiFile) break; // file level, can't try more
+      }
+    }
+    return ret;
+  }
+
+  /**
+   * @param inner an element presumably inside a method within a class, or a method itself.
+   * @param outer an element presumably in the class context.
+   * @return true if an outer element is in a class context, while the cap is a method or function inside it.
+   * @see com.jetbrains.python.psi.PyResolveUtil#getConcealingParent(com.intellij.psi.PsiElement)
+   */
+  protected static boolean refersFromMethodToClass(final PsiElement inner, final PsiElement outer) {
+    return (
+      //(PsiTreeUtil.isAncestor(outer, cap, true)) && // just to make sure
+      (getConcealingParent(outer) instanceof PyClass) && // outer is in a class context
+      (PsiTreeUtil.getParentOfType(inner, PyFunction.class, false) != null) // cap is a function or method within the class
+    );
+  }
+
   /**
    * Returns treeCrawlUp(processor, elt, false). A convenience method.
    * @see com.jetbrains.python.psi.PyResolveUtil#treeCrawlUp(PsiScopeProcessor, PsiElement, boolean)

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

@@ -163,22 +163,7 @@ public class PyReferenceExpressionImpl extends PyElementImpl implements PyRefere
       ret = PyResolveUtil.treeCrawlUp(new PyResolveUtil.ResolveProcessor(referencedName), bfile, true); 
     }
     if (ret == null) {
-      // if we're under a cap, an external object that we want to use might be also defined below us.
-      // look through all contexts, closest first.
-      PsiElement our_cap = PyResolveUtil.getConcealingParent(this);
-      PyResolveUtil.ResolveProcessor proc = new PyResolveUtil.ResolveProcessor(referencedName); // reusable till first hit
-      if (our_cap != null) {
-        PsiElement cap = our_cap;
-        while (true) {
-          cap = PyResolveUtil.getConcealingParent(cap);
-          if (cap == null) cap = this.getContainingFile();
-          ret = PyResolveUtil.treeCrawlUp(proc, cap, true);
-          if ((ret != null) && !PsiTreeUtil.isAncestor(our_cap, ret, true)) {
-            break;
-          }
-          if (cap instanceof PsiFile) break; // file level, can't try more
-        }
-      }
+      ret = PyResolveUtil.resolveOffContext(this);
     }
     return ret;
   }

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

@@ -1,5 +1,8 @@
 package com.jetbrains.python.psi.impl;
 
+import com.intellij.ide.DataManager;
+import com.intellij.openapi.actionSystem.DataContext;
+import com.intellij.openapi.actionSystem.PlatformDataKeys;
 import com.intellij.openapi.extensions.Extensions;
 import com.intellij.openapi.module.Module;
 import com.intellij.openapi.module.ModuleUtil;
@@ -10,8 +13,8 @@ import com.intellij.psi.PsiElement;
 import com.intellij.psi.PsiFile;
 import com.intellij.psi.PsiManager;
 import com.jetbrains.python.psi.*;
-import org.jetbrains.annotations.Nullable;
 import org.jetbrains.annotations.NonNls;
+import org.jetbrains.annotations.Nullable;
 
 import java.util.ArrayList;
 import java.util.Iterator;
@@ -113,7 +116,10 @@ public class ResolveImportUtil {
         @Nullable
         public PsiElement visitJdkOrderEntry(final JdkOrderEntry jdkOrderEntry, final PsiElement value) {
           if (value != null) return value;
-          return resolveInRoots(jdkOrderEntry.getRootFiles(OrderRootType.SOURCES), the_name, importRef);
+          LookupRootVisitor visitor = new LookupRootVisitor(the_name, importRef.getManager());
+          visitRoots(jdkOrderEntry.getRootFiles(OrderRootType.SOURCES), visitor);
+          return visitor.getResult();
+          /*return resolveInRoots(jdkOrderEntry.getRootFiles(OrderRootType.SOURCES), the_name, importRef);*/
         }
       };
       return ModuleRootManager.getInstance(module).processOrder(resolvePolicy, null);
@@ -157,8 +163,14 @@ public class ResolveImportUtil {
     return null;
   }
 
+  public static void visitRoots(final VirtualFile[] roots, SdkRootVisitor visitor) {
+    for (VirtualFile root: roots) {
+      if (! visitor.visitRoot(root)) break;
+    }
+  }
+
   /**
-  Tries to find referencedName under a root. Only used for resolution of import statements.
+  Tries to find referencedName under a root.
   @param root where to look for the referenced name.
   @param referencedName which name to look for.
   @param importRef import reference which resolution led to this call.
@@ -180,6 +192,78 @@ public class ResolveImportUtil {
     return null;
   }
 
+
+  interface SdkRootVisitor {
+    /**
+     * @param root what we're visiting.
+     * @return false when visiting must stop.
+     */
+    boolean visitRoot(VirtualFile root);
+  }
+
+  static class LookupRootVisitor implements SdkRootVisitor {
+    String name;
+    PsiManager psimgr;
+    PsiElement result;
+
+    public LookupRootVisitor(String name, PsiManager psimgr) {
+      this.name = name;
+      this.psimgr = psimgr;
+      this.result = null;
+    }
+
+    public boolean visitRoot(final VirtualFile root) {
+      final VirtualFile childFile = root.findChild(name + PY_SUFFIX);
+      if (childFile != null) {
+        result = psimgr.findFile(childFile);
+        return (result == null);
+      }
+
+      final VirtualFile childDir = root.findChild(name);
+      if (childDir != null) {
+        result = psimgr.findDirectory(childDir);
+        return (result == null);
+      }
+      return true;
+    }
+
+    public PsiElement getResult() {
+      return result;
+    }
+  }
+
+  static class CollectingRootVisitor implements SdkRootVisitor {
+    List<String> result;
+    PsiManager psimgr;
+
+    static String cutExt(String name) {
+      return name.substring(0, Math.max(name.length() - PY_SUFFIX.length(), 0));
+    }
+
+    public CollectingRootVisitor(PsiManager psimgr) {
+      result = new ArrayList<String>(25);
+      this.psimgr = psimgr;
+    }
+
+    public boolean visitRoot(final VirtualFile root) {
+      for (VirtualFile vfile : root.getChildren()) {
+        if (vfile.getName().endsWith(PY_SUFFIX)) {
+          PsiFile pfile = psimgr.findFile(vfile);
+          if (pfile != null) result.add(cutExt(pfile.getName()));
+        }
+        else if (vfile.isDirectory() && (vfile.findChild(INIT_PY) != null)) {
+          PsiDirectory pdir = psimgr.findDirectory(vfile);
+          if (pdir != null) result.add(pdir.getName());
+        }
+      }
+      return true; // continue forever
+    }
+
+    public List<String> getResult() {
+      return result;
+    }
+  }
+
   /**
   Tries to find referencedName under the parent element. Used to resolve any names that look imported.
   Parent might happen to be a PyFile(__init__.py), then it is treated <i>both</i> as a file and as ist base dir.
@@ -240,32 +324,49 @@ public class ResolveImportUtil {
   /**
    * Finds reasonable names to import to complete a patrial name.
    * @param partial_ref reference containing the partial name.
-   * @return an array of names ready for gtVariants().
+   * @return an array of names ready for getVariants().
    */
-  public static String[] suggestImportVariants(PyReferenceExpression partial_ref) {
+  public static String[] suggestImportVariants(final PyReferenceExpression partial_ref) {
+    // look in builtins
     List<String> variants = new ArrayList<String>();
-    String prefix_u = partial_ref.getNode().getText().toUpperCase(); // we try case-insensitively
-    //
+    DataContext dataContext = DataManager.getInstance().getDataContext();
     // look at current dir
-    final VirtualFile pfile = partial_ref.getContainingFile().getVirtualFile();
+    final VirtualFile pfile = PlatformDataKeys.VIRTUAL_FILE.getData(dataContext);
     if (pfile != null) {
       VirtualFile pdir = pfile.getParent();
-      _siftDir(pdir, prefix_u, variants, pfile) ;
+      if (pdir != null) {
+        for (VirtualFile a_file : pdir.getChildren()) {
+          if (a_file != pfile) {
+            if (pfile.isDirectory()) {
+              if (pfile.findChild(INIT_PY) != null) variants.add(a_file.getName());
+            }
+            else { // plain file
+              String fname = a_file.getName();
+              if (fname.endsWith(PY_SUFFIX)) {
+                variants.add(fname.substring(0, fname.length() - PY_SUFFIX.length()));
+              }
+            }
+          }
+        }
+      }
     }
     // look in SDK
-    // TODO: implement, reusing resolver code
-    return variants.toArray(new String[variants.size()]); 
+    final CollectingRootVisitor visitor = new CollectingRootVisitor(partial_ref.getManager());
+    final Module module = ModuleUtil.findModuleForPsiElement(partial_ref);
+    if (module != null) {
+      RootPolicy<PsiElement> resolvePolicy = new RootPolicy<PsiElement>() {
+        @Nullable
+        public PsiElement visitJdkOrderEntry(final JdkOrderEntry jdkOrderEntry, final PsiElement value) {
+          if (value != null) return value;
+          visitRoots(jdkOrderEntry.getRootFiles(OrderRootType.SOURCES), visitor);
+          return null;
+        }
+      };
+      ModuleRootManager.getInstance(module).processOrder(resolvePolicy, null);
+      variants.addAll(visitor.getResult());
+    }
+
+    return variants.toArray(new String[variants.size()]);
   }
 
-   static void _siftDir(VirtualFile pdir, String prefix, List<String> variants, VirtualFile pfile) {
-     if (pdir != null) {
-       for (VirtualFile a_file : pdir.getChildren()) {
-         // TODO: check extensions, chack subdirs with __init__.py
-         if ((a_file != pfile) && (a_file.getName().toUpperCase().startsWith(prefix))) {
-           variants.add(a_file.getName());
-         }
-       }
-     }
-   }
-
 }

python/testData/resolve/ClassPeerMembers.py

@@ -0,0 +1,3 @@
+class A:
+  def foo(self): pass
+  bar = f<ref>oo

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

@@ -106,6 +106,11 @@ public class PyResolveTest extends ResolveTestCase {
     assertTrue(targetElement instanceof PyTargetExpression);
   }
 
+  public void testClassPeerMembers() throws Exception {
+    PsiElement target = resolve();
+    assertTrue(target instanceof PyFunction);
+  }
+
   public void testTuple() throws Exception {
     PsiElement targetElement = resolve();
     assertTrue(targetElement instanceof PyTargetExpression);

python/testSrc/com/jetbrains/python/PythonHighlightingTest.java

@@ -11,6 +11,10 @@ public class PythonHighlightingTest extends DaemonAnalyzerTestCase {
     return PathManager.getHomePath() + "/plugins/python/testData/highlighting/";
   }
 
+  public void testImportInTry() throws Exception {
+    doTest();
+  }
+
   public void testReturnOutsideOfFunction() throws Exception {
     doTest();
   }