Merge branch 'python-fixes'

python/src/com/jetbrains/python/codeInsight/imports/AddImportHelper.java

@@ -87,8 +87,9 @@ public class AddImportHelper {
   }
 
   @Nullable
-  private static PsiElement getInsertPosition(final PsiFile file, @Nullable String nameToImport, @Nullable ImportPriority priority) {
-    PsiElement feeler = file.getFirstChild();
+  private static PsiElement getInsertPosition(final PsiElement insertParent, @Nullable String nameToImport,
+                                              @Nullable ImportPriority priority) {
+    PsiElement feeler = insertParent.getFirstChild();
     if (feeler == null) return null;
     // skip initial comments and whitespace and try to get just below the last import stmt
     boolean skippedOverImports = false;
@@ -97,7 +98,7 @@ public class AddImportHelper {
     final boolean isInjected = InjectedLanguageManager.getInstance(feeler.getProject()).isInjectedFragment(feeler.getContainingFile());
     do {
       if (feeler instanceof PyImportStatementBase && !isInjected) {
-        if (nameToImport != null && priority != null && shouldInsertBefore(file, (PyImportStatementBase)feeler, nameToImport, priority)) {
+        if (nameToImport != null && priority != null && shouldInsertBefore(insertParent, (PyImportStatementBase)feeler, nameToImport, priority)) {
           break;
         }
         seeker = feeler;
@@ -109,10 +110,9 @@ public class AddImportHelper {
         feeler = feeler.getNextSibling();
       }
       // maybe we arrived at the doc comment stmt; skip over it, too
-      else if (!skippedOverImports && !skippedOverDoc && file instanceof PyFile) {
+      else if (!skippedOverImports && !skippedOverDoc && insertParent instanceof PyFile) {
         PsiElement doc_elt =
-          DocStringUtil
-            .findDocStringExpression((PyElement)file); // this gives the literal; its parent is the expr seeker may have encountered
+          DocStringUtil.findDocStringExpression((PyElement)insertParent); // this gives the literal; its parent is the expr seeker may have encountered
         if (doc_elt != null && doc_elt.getParent() == feeler) {
           feeler = feeler.getNextSibling();
           seeker = feeler; // skip over doc even if there's nothing below it
@@ -130,7 +130,8 @@ public class AddImportHelper {
     return seeker;
   }
 
-  private static boolean shouldInsertBefore(PsiFile file, PyImportStatementBase relativeTo, String nameToImport, ImportPriority priority) {
+  private static boolean shouldInsertBefore(PsiElement element, PyImportStatementBase relativeTo, String nameToImport,
+                                            ImportPriority priority) {
     QualifiedName relativeToName;
     PsiElement source;
     if (relativeTo instanceof PyFromImportStatement) {
@@ -161,7 +162,7 @@ public class AddImportHelper {
     }
     ImportPriority relativeToPriority = source == null || containingFile == null
                                         ? ImportPriority.BUILTIN
-                                        : getImportPriority(file, containingFile);
+                                        : getImportPriority(element, containingFile);
     final int rc = priority.compareTo(relativeToPriority);
     if (rc < 0) {
       return true;
@@ -213,22 +214,32 @@ public class AddImportHelper {
    * @param from   name of the module
    * @param name   imported name
    * @param asName optional name for 'as' clause
+   * @param anchor place where the imported name was used
    */
-  public static void addImportFromStatement(PsiFile file, String from, String name, @Nullable String asName, ImportPriority priority) {
+  public static void addImportFromStatement(PsiFile file, String from, String name, @Nullable String asName, ImportPriority priority,
+                                            @Nullable PsiElement anchor) {
     final PyElementGenerator generator = PyElementGenerator.getInstance(file.getProject());
     final LanguageLevel languageLevel = LanguageLevel.forElement(file);
     final PyFromImportStatement nodeToInsert = generator.createFromImportStatement(languageLevel, from, name, asName);
     try {
+      final PyImportStatementBase importStatement = PsiTreeUtil.getParentOfType(anchor, PyImportStatementBase.class, false);
+      final PsiElement insertParent;
+      if (importStatement != null && importStatement.getContainingFile() == file) {
+        insertParent = importStatement.getParent();
+      }
+      else {
+        insertParent = file;
+      }
       if (InjectedLanguageManager.getInstance(file.getProject()).isInjectedFragment(file)) {
-        final PsiElement element = file.addBefore(nodeToInsert, getInsertPosition(file, from, priority));
+        final PsiElement element = insertParent.addBefore(nodeToInsert, getInsertPosition(insertParent, from, priority));
         PsiElement whitespace = element.getNextSibling();
         if (!(whitespace instanceof PsiWhiteSpace)) {
           whitespace = PsiParserFacade.SERVICE.getInstance(file.getProject()).createWhiteSpaceFromText("  >>> ");
         }
-        file.addBefore(whitespace, element);
+        insertParent.addBefore(whitespace, element);
       }
       else {
-        file.addBefore(nodeToInsert, getInsertPosition(file, from, priority));
+        insertParent.addBefore(nodeToInsert, getInsertPosition(insertParent, from, priority));
       }
     }
     catch (IncorrectOperationException e) {
@@ -237,7 +248,7 @@ public class AddImportHelper {
   }
 
   public static boolean addImportFrom(PsiFile file, @Nullable PsiElement target, String path, final String name,
-                                      @Nullable String asName, ImportPriority priority) {
+                                      @Nullable String asName, ImportPriority priority, @Nullable PsiElement anchor) {
     final List<PyFromImportStatement> existingImports = ((PyFile)file).getFromImports();
     for (PyFromImportStatement existingImport : existingImports) {
       if (target != null && existingImport.getTextRange().getStartOffset() > target.getTextRange().getStartOffset()) {
@@ -259,7 +270,7 @@ public class AddImportHelper {
         return true;
       }
     }
-    addImportFromStatement(file, path, name, asName, priority);
+    addImportFromStatement(file, path, name, asName, priority, anchor);
     return true;
   }
 
@@ -283,7 +294,7 @@ public class AddImportHelper {
         element.replace(elementGenerator.createExpressionFromText(LanguageLevel.forElement(target), toImportQName + "." + targetName));
       }
       else {
-        addImportFrom(file, null, toImportQName.toString(), target.getName(), null, priority);
+        addImportFrom(file, null, toImportQName.toString(), target.getName(), null, priority, element);
       }
     }
   }

python/src/com/jetbrains/python/codeInsight/imports/ImportFromExistingAction.java

@@ -181,7 +181,7 @@ public class ImportFromExistingAction implements QuestionAction {
           AddImportHelper.addLocalFromImportStatement(myTarget, qualifiedName, myName);
         }
         else {
-          AddImportHelper.addImportFromStatement(file, qualifiedName, myName, null, priority);
+          AddImportHelper.addImportFromStatement(file, qualifiedName, myName, null, priority, null);
         }
       }
     }

python/src/com/jetbrains/python/codeInsight/testIntegration/PyTestCreator.java

@@ -85,7 +85,7 @@ public class PyTestCreator implements TestCreator {
 
               PsiFile psiFile = PyUtil.getOrCreateFile(
                 dialog.getTargetDir() + "/" + fileName, project);
-              AddImportHelper.addImportFrom(psiFile, null, "unittest", "TestCase", null, AddImportHelper.ImportPriority.BUILTIN);
+              AddImportHelper.addImportFrom(psiFile, null, "unittest", "TestCase", null, AddImportHelper.ImportPriority.BUILTIN, null);
 
               PyElement createdClass = PyElementGenerator.getInstance(project).createFromText(
                 LanguageLevel.forElement(psiFile), PyClass.class,

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

@@ -161,15 +161,33 @@ public class PyBuiltinCache {
   }
 
   @Nullable
-  static PyType createLiteralCollectionType(final PySequenceExpression sequence, final String name) {
-    final PyBuiltinCache builtinCache = getInstance(sequence);
-    final PyClass setClass = builtinCache.getClass(name);
-    if (setClass != null) {
-      return new PyLiteralCollectionType(setClass, false, sequence);
+  public PyType createLiteralCollectionType(final PySequenceExpression sequence, final String name, @NotNull TypeEvalContext context) {
+    final PyClass cls = getClass(name);
+    if (cls != null) {
+      return new PyCollectionTypeImpl(cls, false, getSequenceElementType(sequence, context));
     }
     return null;
   }
 
+  @Nullable
+  private static PyType getSequenceElementType(@NotNull PySequenceExpression sequence, @NotNull TypeEvalContext context) {
+    final PyExpression[] elements = sequence.getElements();
+    if (elements.length == 0 || elements.length > 10 /* performance */) {
+      return null;
+    }
+    final PyType result = context.getType(elements[0]);
+    if (result == null) {
+      return null;
+    }
+    for (int i = 1; i < elements.length; i++) {
+      final PyType elementType = context.getType(elements[i]);
+      if (elementType == null || !elementType.equals(result)) {
+        return null;
+      }
+    }
+    return result;
+  }
+
   @Nullable
   public PyFile getBuiltinsFile() {
     return myBuiltinsFile;

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

@@ -38,7 +38,7 @@ public class PyDictLiteralExpressionImpl extends PyElementImpl implements PyDict
   }
 
   public PyType getType(@NotNull TypeEvalContext context, @NotNull TypeEvalContext.Key key) {
-    return PyBuiltinCache.createLiteralCollectionType(this, "dict");
+    return PyBuiltinCache.getInstance(this).createLiteralCollectionType(this, "dict", context);
   }
 
   @Override

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

@@ -270,6 +270,16 @@ public class PyFunctionImpl extends PyBaseElementImpl<PyFunctionStub> implements
       public String getPresentableText() {
         return notNullize(getName(), PyNames.UNNAMED_ELEMENT) + getParameterList().getPresentableText(true);
       }
+
+      @Nullable
+      @Override
+      public String getLocationString() {
+        final PyClass containingClass = getContainingClass();
+        if (containingClass != null) {
+          return "(" + containingClass.getName() + " in " + getPackageForFile(getContainingFile()) + ")";
+        }
+        return super.getLocationString();
+      }
     };
   }
 

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

@@ -86,6 +86,6 @@ public class PyListLiteralExpressionImpl extends PyElementImpl implements PyList
   }
 
   public PyType getType(@NotNull TypeEvalContext context, @NotNull TypeEvalContext.Key key) {
-    return PyBuiltinCache.createLiteralCollectionType(this, "list");
+    return PyBuiltinCache.getInstance(this).createLiteralCollectionType(this, "list", context);
   }
 }

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

@@ -31,7 +31,7 @@ public class PySetLiteralExpressionImpl extends PyElementImpl implements PySetLi
   }
 
   public PyType getType(@NotNull TypeEvalContext context, @NotNull TypeEvalContext.Key key) {
-    return PyBuiltinCache.createLiteralCollectionType(this, "set");
+    return PyBuiltinCache.getInstance(this).createLiteralCollectionType(this, "set", context);
   }
 
   @Override

python/src/com/jetbrains/python/psi/types/PyLiteralCollectionType.java

@@ -1,52 +0,0 @@
-/*
- * Copyright 2000-2014 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.psi.types;
-
-import com.jetbrains.python.psi.PyClass;
-import com.jetbrains.python.psi.PyExpression;
-import com.jetbrains.python.psi.PySequenceExpression;
-import org.jetbrains.annotations.NotNull;
-
-/**
- * @author yole
- */
-public class PyLiteralCollectionType extends PyClassTypeImpl implements PyCollectionType {
-  private final PySequenceExpression mySequence;
-
-  public PyLiteralCollectionType(@NotNull PyClass source, boolean isDefinition, PySequenceExpression sequence) {
-    super(source, isDefinition);
-    mySequence = sequence;
-  }
-
-  @Override
-  public PyType getElementType(@NotNull TypeEvalContext context) {
-    final PyExpression[] elements = mySequence.getElements();
-    if (elements.length == 0 || elements.length > 10 /* performance */) {
-      return null;
-    }
-    PyType result = context.getType(elements [0]);
-    if (result == null) {
-      return null;
-    }
-    for (int i = 1; i < elements.length; i++) {
-      PyType elementType = context.getType(elements[i]);
-      if (elementType == null || !elementType.equals(result)) {
-        return null;
-      }
-    }
-    return result;
-  }
-}

python/src/com/jetbrains/python/refactoring/classes/PyClassRefactoringUtil.java

@@ -318,7 +318,7 @@ public final class PyClassRefactoringUtil {
     }
     final AddImportHelper.ImportPriority priority = AddImportHelper.getImportPriority(anchor, newFile);
     if (preferFromImport && !containingQName.getComponents().isEmpty()) {
-      return AddImportHelper.addImportFrom(file, null, containingQName.toString(), importedName, asName, priority);
+      return AddImportHelper.addImportFrom(file, null, containingQName.toString(), importedName, asName, priority, anchor);
     }
     else {
       return AddImportHelper.addImportStatement(file, containingQName.append(importedName).toString(), asName, priority);

python/src/com/jetbrains/python/refactoring/classes/membersManager/MethodsManager.java

@@ -163,7 +163,7 @@ class MethodsManager extends MembersManager<PyFunction> {
    */
   private static void addImportFromAbc(@NotNull final PsiFile file, @NotNull final String nameToImport) {
     AddImportHelper.addImportFromStatement(file, ABC_META_PACKAGE, nameToImport, null,
-                                           AddImportHelper.ImportPriority.BUILTIN);
+                                           AddImportHelper.ImportPriority.BUILTIN, null);
   }
 
   /**

python/testData/hierarchy/call/Static/Constructor/Constructor_callee_verification.xml

@@ -1,8 +1,8 @@
-<node text="A.__init__(self) (hierarchy.call.Static.Constructor.main)" base="true">
+<node text="A.__init__(self) (A in hierarchy.call.Static.Constructor.main)" base="true">
   <node text="invoke1(p) (hierarchy.call.Static.Constructor.main)">
-    <node text="A.method1(self) (hierarchy.call.Static.Constructor.main)"/>
+    <node text="A.method1(self) (A in hierarchy.call.Static.Constructor.main)"/>
   </node>
   <node text="invoke2(p) (hierarchy.call.Static.Constructor.main)">
-    <node text="A.method2(self) (hierarchy.call.Static.Constructor.main)"/>
+    <node text="A.method2(self) (A in hierarchy.call.Static.Constructor.main)"/>
   </node>
 </node>

python/testData/hierarchy/call/Static/Constructor/Constructor_caller_verification.xml

@@ -1,6 +1,6 @@
-<node text="A.__init__(self) (hierarchy.call.Static.Constructor.main)" base="true">
+<node text="A.__init__(self) (A in hierarchy.call.Static.Constructor.main)" base="true">
   <node text="invokeA() (hierarchy.call.Static.Constructor.main)">
-    <node text="C.bar(self) (hierarchy.call.Static.Constructor.main)"/>
+    <node text="C.bar(self) (C in hierarchy.call.Static.Constructor.main)"/>
   </node>
-  <node text="C.bar(self) (hierarchy.call.Static.Constructor.main)"/>
+  <node text="C.bar(self) (C in hierarchy.call.Static.Constructor.main)"/>
 </node>

python/testData/hierarchy/call/Static/Inheritance/Inheritance_callee_verification.xml

@@ -1 +1 @@
-<node text="A.target_func(self) (hierarchy.call.Static.Inheritance.main)" base="true"/>
+<node text="A.target_func(self) (A in hierarchy.call.Static.Inheritance.main)" base="true"/>

python/testData/hierarchy/call/Static/Inheritance/Inheritance_caller_verification.xml

@@ -1,4 +1,4 @@
-<node text="A.target_func(self) (hierarchy.call.Static.Inheritance.main)" base="true">
-  <node text="C.func(self, a) (hierarchy.call.Static.Inheritance.main)"/>
+<node text="A.target_func(self) (A in hierarchy.call.Static.Inheritance.main)" base="true">
+  <node text="C.func(self, a) (C in hierarchy.call.Static.Inheritance.main)"/>
   <node text="foo2(b) (hierarchy.call.Static.Inheritance.main)"/>
 </node>

python/testData/hierarchy/call/Static/OverriddenMethod/OverriddenMethod_callee_verification.xml

@@ -1,3 +1,3 @@
-<node text="B.target_func(self, p) (hierarchy.call.Static.OverriddenMethod.main)" base="true">
-  <node text="A.another_func(self) (hierarchy.call.Static.OverriddenMethod.file_1)"/>
+<node text="B.target_func(self, p) (B in hierarchy.call.Static.OverriddenMethod.main)" base="true">
+  <node text="A.another_func(self) (A in hierarchy.call.Static.OverriddenMethod.file_1)"/>
 </node>

python/testData/hierarchy/call/Static/OverriddenMethod/OverriddenMethod_caller_verification.xml

@@ -1,5 +1,5 @@
-<node text="B.target_func(self, p) (hierarchy.call.Static.OverriddenMethod.main)" base="true">
-  <node text="C.func1(self, a) (hierarchy.call.Static.OverriddenMethod.main)"/>
+<node text="B.target_func(self, p) (B in hierarchy.call.Static.OverriddenMethod.main)" base="true">
+  <node text="C.func1(self, a) (C in hierarchy.call.Static.OverriddenMethod.main)"/>
   <node text="bar1(a) (hierarchy.call.Static.OverriddenMethod.main)"/>
-  <node text="C.func2(self) (hierarchy.call.Static.OverriddenMethod.main)"/>
+  <node text="C.func2(self) (C in hierarchy.call.Static.OverriddenMethod.main)"/>
 </node>

python/testData/inspections/PyTypeCheckerInspection/RecursiveDictAttribute.py

@@ -0,0 +1,5 @@
+class C:
+    def f(self, x):
+        self.foo = x
+        self.foo = {'foo': self.foo}
+        return self.foo['foo'] + 10

python/testData/refactoring/move/conditionalImport/after/src/mod1.py

@@ -0,0 +1,6 @@
+if True:
+    import pkgutil
+else:
+    from pkg1 import mod2 as pkgutil
+
+print(pkgutil)

python/testData/refactoring/move/conditionalImport/before/src/mod1.py

@@ -0,0 +1,6 @@
+if True:
+    import pkgutil
+else:
+    import mod2 as pkgutil
+
+print(pkgutil)

python/testSrc/com/jetbrains/python/PyAddImportTest.java

@@ -53,7 +53,7 @@ public class PyAddImportTest extends PyTestCase {
     new WriteCommandAction(myFixture.getProject(), myFixture.getFile()) {
       @Override
       protected void run(Result result) throws Throwable {
-        AddImportHelper.addImportFrom(myFixture.getFile(), null, path, name, null, AddImportHelper.ImportPriority.BUILTIN);
+        AddImportHelper.addImportFrom(myFixture.getFile(), null, path, name, null, AddImportHelper.ImportPriority.BUILTIN, null);
       }
     }.execute();
     myFixture.checkResultByFile("addImport/" + getTestName(true) + ".after.py");

python/testSrc/com/jetbrains/python/PyTypeTest.java

@@ -505,7 +505,7 @@ public class PyTypeTest extends PyTestCase {
 
   // EA-40207
   public void testRecursion() {
-    doTest("list[list]",
+    doTest("list",
            "def f():\n" +
            "    return [f()]\n" +
            "expr = f()\n");

python/testSrc/com/jetbrains/python/inspections/PyTypeCheckerInspectionTest.java

@@ -237,4 +237,9 @@ public class PyTypeCheckerInspectionTest extends PyTestCase {
   public void testDefaultTupleParameter() {
     doTest();
   }
+
+  // PY-14222
+  public void testRecursiveDictAttribute() {
+    doTest();
+  }
 }

python/testSrc/com/jetbrains/python/refactoring/PyMoveTest.java

@@ -133,6 +133,11 @@ public class PyMoveTest extends PyTestCase {
     doMoveSymbolTest("g", "c.py");
   }
 
+  // PY-13870
+  public void testConditionalImport() {
+    doMoveFileTest("mod2.py", "pkg1");
+  }
+
   private void doMoveFileTest(String fileName, String toDirName)  {
     Project project = myFixture.getProject();
     PsiManager manager = PsiManager.getInstance(project);