PyArgumentList#addArgument fixed and test added

python/psi-api/src/com/jetbrains/python/psi/PyArgumentList.java

@@ -27,12 +27,25 @@ import org.jetbrains.annotations.Nullable;
  */
 public interface PyArgumentList extends PyElement {
 
-  @NotNull PyExpression[] getArguments();
+  @NotNull
+  PyExpression[] getArguments();
 
-  @Nullable PyKeywordArgument getKeywordArgument(String name);
+  @Nullable
+  PyKeywordArgument getKeywordArgument(String name);
+
+  /**
+   * TODO: Copy/Paste with {@link com.jetbrains.python.psi.PyCallExpression#addArgument(PyExpression)} ?
+   * Adds argument to the appropriate place:
+   * {@link com.jetbrains.python.psi.PyKeywordArgument} goes to the end.
+   * All other go before key arguments (if any) but after last non-key arguments.
+   * Commas should be set correctly as well.
+   *
+   * @param arg argument to add
+   */
+  void addArgument(@NotNull PyExpression arg);
 
-  void addArgument(PyExpression arg);
   void addArgumentFirst(PyExpression arg);
+
   void addArgumentAfter(PyExpression argument, PyExpression afterThis);
 
   /**
@@ -43,9 +56,10 @@ public interface PyArgumentList extends PyElement {
 
   /**
    * Tries to map the argument list to callee's idea of parameters.
-   * @return a result object with mappings and diagnostic flags.
+   *
    * @param resolveContext the reference resolution context
    * @param implicitOffset known from the context implicit offset
+   * @return a result object with mappings and diagnostic flags.
    */
   @NotNull
   CallArgumentsMapping analyzeCall(PyResolveContext resolveContext, int implicitOffset);

python/psi-api/src/com/jetbrains/python/psi/PyCallExpression.java

@@ -68,6 +68,10 @@ public interface PyCallExpression extends PyExpression {
   @Nullable
   PyExpression getKeywordArgument(String keyword);
 
+  /**
+   * TODO: Copy/Paste with {@link com.jetbrains.python.psi.PyArgumentList#addArgument(PyExpression)}
+   * @param expression
+   */
   void addArgument(PyExpression expression);
 
   /**

python/psi-api/src/com/jetbrains/python/psi/PyElementGenerator.java

@@ -67,6 +67,21 @@ public abstract class PyElementGenerator {
 
   public abstract PyExpression createExpressionFromText(final LanguageLevel languageLevel, String text);
 
+  /**
+   * Adds elements to list inserting required commas.
+   * Method is like {@link #insertItemIntoList(PyElement, PyExpression, PyExpression)} but does not add unneeded commas.
+   *
+   * @param list where to add
+   * @param afterThis after which element it should be added (null for add to the head)
+   * @param toInsert what to insert
+   * @return newly inserted element
+   */
+  @NotNull
+  public abstract PsiElement insertItemIntoListRemoveRedundantCommas(
+    @NotNull PyElement list,
+    @Nullable PyExpression afterThis,
+    @NotNull PyExpression toInsert);
+
   public abstract PsiElement insertItemIntoList(PyElement list, @Nullable PyExpression afterThis, PyExpression toInsert)
     throws IncorrectOperationException;
 

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

@@ -15,13 +15,15 @@
  */
 package com.jetbrains.python.psi.impl;
 
+import com.google.common.collect.Collections2;
+import com.google.common.collect.Queues;
 import com.intellij.lang.ASTFactory;
 import com.intellij.lang.ASTNode;
-import com.intellij.psi.PsiElement;
 import com.intellij.psi.tree.IElementType;
 import com.intellij.psi.tree.TokenSet;
 import com.intellij.psi.util.PsiTreeUtil;
 import com.intellij.util.IncorrectOperationException;
+import com.jetbrains.NotNullPredicate;
 import com.jetbrains.python.PyElementTypes;
 import com.jetbrains.python.PyTokenTypes;
 import com.jetbrains.python.PythonDialectsTokenSetProvider;
@@ -30,9 +32,13 @@ import com.jetbrains.python.psi.resolve.PyResolveContext;
 import org.jetbrains.annotations.NotNull;
 import org.jetbrains.annotations.Nullable;
 
-import java.util.Arrays;
+import java.util.*;
 
 public class PyArgumentListImpl extends PyElementImpl implements PyArgumentList {
+
+  // Filters all expressions but keyword arguments
+  private static final NoKeyArguments NO_KEY_ARGUMENTS = new NoKeyArguments();
+
   public PyArgumentListImpl(ASTNode astNode) {
     super(astNode);
   }
@@ -61,47 +67,64 @@ public class PyArgumentListImpl extends PyElementImpl implements PyArgumentList
     return null;
   }
 
-  public void addArgument(PyExpression arg) {
-    // it should find the comma after the argument to add after, and add after
-    // that. otherwise it won't deal with comments nicely
+  @Override
+  public void addArgument(@NotNull final PyExpression arg) {
+    final PyElementGenerator generator = new PyElementGeneratorImpl(getProject());
+
+    // Adds param to appropriate place
+    final Deque<PyKeywordArgument> keywordArguments = getKeyWordArguments();
+    final Deque<PyExpression> parameters = getParameters();
+
+    if (keywordArguments.isEmpty() && parameters.isEmpty()) {
+      generator.insertItemIntoListRemoveRedundantCommas(this, null, arg);
+      return;
+    }
+
+
     if (arg instanceof PyKeywordArgument) {
-      PyKeywordArgument keywordArgument = (PyKeywordArgument)arg;
-      PyKeywordArgument lastKeyArg = null;
-      PyExpression firstNonKeyArg = null;
-      for (PsiElement element : getChildren()) {
-        if (element instanceof PyKeywordArgument) {
-          lastKeyArg = (PyKeywordArgument)element;
-        }
-        else if (element instanceof PyExpression && firstNonKeyArg == null) {
-          firstNonKeyArg = (PyExpression)element;
-        }
-      }
-      if (lastKeyArg != null) {
-        // add after last key arg
-        addArgumentNode(keywordArgument, lastKeyArg.getNode().getTreeNext(), true);
-
-      }
-      else if (firstNonKeyArg != null) {
-        // add before first non key arg
-        addArgumentNode(keywordArgument, firstNonKeyArg.getNode(), true);
-
+      if (parameters.isEmpty()) {
+        generator.insertItemIntoListRemoveRedundantCommas(this, keywordArguments.getLast(), arg);
       }
       else {
-        // add as only argument
-        addArgumentLastWithoutComma(arg);
+        if (keywordArguments.isEmpty()) {
+          generator.insertItemIntoListRemoveRedundantCommas(this, parameters.getLast(), arg);
+        }
+        else {
+          generator.insertItemIntoListRemoveRedundantCommas(this, keywordArguments.getLast(), arg);
+        }
       }
     }
     else {
-      final PyExpression[] args = getArguments();
-      if (args.length > 0) {
-        addArgumentAfter(arg, args [args.length-1]);
+      if (parameters.isEmpty()) {
+        generator.insertItemIntoListRemoveRedundantCommas(this, null, arg);
       }
       else {
-        addArgumentLastWithoutComma(arg);
+        generator.insertItemIntoListRemoveRedundantCommas(this, parameters.getLast(), arg);
       }
     }
   }
 
+
+  /**
+   * @return parameters (as opposite to keyword arguments)
+   */
+  @NotNull
+  private Deque<PyExpression> getParameters() {
+    final PyExpression[] childrenOfType = PsiTreeUtil.getChildrenOfType(this, PyExpression.class);
+    if (childrenOfType == null) {
+      return new ArrayDeque<PyExpression>(0);
+    }
+    return Queues.newArrayDeque(Collections2.filter(Arrays.asList(childrenOfType), NO_KEY_ARGUMENTS));
+  }
+
+  /**
+   * @return keyword arguments (as opposite to parameters)
+   */
+  @NotNull
+  private Deque<PyKeywordArgument> getKeyWordArguments() {
+    return Queues.newArrayDeque(PsiTreeUtil.findChildrenOfType(this, PyKeywordArgument.class));
+  }
+
   public void addArgumentFirst(PyExpression arg) {
     ASTNode node = getNode();
     ASTNode[] pars = node.getChildren(TokenSet.create(PyTokenTypes.LPAR));
@@ -113,13 +136,12 @@ public class PyArgumentListImpl extends PyElementImpl implements PyArgumentList
       catch (IncorrectOperationException e1) {
         throw new IllegalStateException(e1);
       }
-
     }
     else {
       ASTNode before = PyUtil.getNextNonWhitespace(pars[0]);
       ASTNode anchorBefore;
       if (before != null && elementPrecedesElementsOfType(before, PythonDialectsTokenSetProvider.INSTANCE.getExpressionTokens())) {
-        ASTNode comma = PyElementGenerator.getInstance(getProject()).createComma();
+        ASTNode comma = createComma();
         node.addChild(comma, before);
         node.addChild(ASTFactory.whitespace(" "), before);
         anchorBefore = comma;
@@ -137,6 +159,14 @@ public class PyArgumentListImpl extends PyElementImpl implements PyArgumentList
     }
   }
 
+  /**
+   * @return newly created comma
+   */
+  @NotNull
+  private ASTNode createComma() {
+    return PyElementGenerator.getInstance(getProject()).createComma();
+  }
+
   private static boolean elementPrecedesElementsOfType(ASTNode before, TokenSet expressions) {
     ASTNode node = before;
     while (node != null) {
@@ -279,4 +309,11 @@ public class PyArgumentListImpl extends PyElementImpl implements PyArgumentList
     }
     return ret;
   }
+
+  private static class NoKeyArguments extends NotNullPredicate<PyExpression> {
+    @Override
+    protected boolean applyNotNull(@NotNull final PyExpression input) {
+      return (PsiTreeUtil.getParentOfType(input, PyKeywordArgument.class) == null) && !(input instanceof PyKeywordArgument);
+    }
+  }
 }

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

@@ -43,6 +43,7 @@ public class PyCallExpressionHelper {
   }
 
   /**
+   * TODO: Copy/Paste with {@link com.jetbrains.python.psi.PyArgumentList#addArgument(com.jetbrains.python.psi.PyExpression)}
    * Adds an argument to the end of argument list.
    * @param us the arg list
    * @param expression what to add

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

@@ -15,6 +15,8 @@
  */
 package com.jetbrains.python.psi.impl;
 
+import com.google.common.collect.Collections2;
+import com.google.common.collect.Queues;
 import com.intellij.lang.ASTNode;
 import com.intellij.openapi.project.Project;
 import com.intellij.openapi.util.Pair;
@@ -23,9 +25,12 @@ import com.intellij.psi.PsiElement;
 import com.intellij.psi.PsiFile;
 import com.intellij.psi.PsiFileFactory;
 import com.intellij.psi.impl.PsiFileFactoryImpl;
+import com.intellij.psi.impl.source.tree.LeafPsiElement;
 import com.intellij.psi.tree.TokenSet;
+import com.intellij.psi.util.PsiTreeUtil;
 import com.intellij.testFramework.LightVirtualFile;
 import com.intellij.util.IncorrectOperationException;
+import com.jetbrains.NotNullPredicate;
 import com.jetbrains.python.PyTokenTypes;
 import com.jetbrains.python.PythonFileType;
 import com.jetbrains.python.PythonLanguage;
@@ -37,12 +42,15 @@ import org.jetbrains.annotations.Nullable;
 import java.nio.charset.Charset;
 import java.nio.charset.CharsetEncoder;
 import java.util.Arrays;
+import java.util.Collection;
+import java.util.Deque;
 import java.util.Formatter;
 
 /**
  * @author yole
  */
 public class PyElementGeneratorImpl extends PyElementGenerator {
+  private static final CommasOnly COMMAS_ONLY = new CommasOnly();
   private final Project myProject;
 
   public PyElementGeneratorImpl(Project project) {
@@ -173,6 +181,30 @@ public class PyElementGeneratorImpl extends PyElementGenerator {
     return dot.copyElement();
   }
 
+  @Override
+  @NotNull
+  public PsiElement insertItemIntoListRemoveRedundantCommas(
+    @NotNull final PyElement list,
+    @Nullable final PyExpression afterThis,
+    @NotNull final PyExpression toInsert) {
+    // TODO: #insertItemIntoList is probably buggy. In such case, fix it and get rid of this method
+    final PsiElement result = insertItemIntoList(list, afterThis, toInsert);
+    final LeafPsiElement[] leafs = PsiTreeUtil.getChildrenOfType(list, LeafPsiElement.class);
+    if (leafs != null) {
+      final Deque<LeafPsiElement> commas = Queues.newArrayDeque(Collections2.filter(Arrays.asList(leafs), COMMAS_ONLY));
+      if (! commas.isEmpty()) {
+        final LeafPsiElement lastComma = commas.getLast();
+        if (PsiTreeUtil.getNextSiblingOfType(lastComma, PyExpression.class) == null) { //Comma has no expression after it
+          lastComma.delete();
+        }
+      }
+    }
+
+    return result;
+  }
+
+  // TODO: Adds comma to empty list: adding "foo" to () will create (foo,). That is why "insertItemIntoListRemoveRedundantCommas" was created.
+  // We probably need to fix this method and delete insertItemIntoListRemoveRedundantCommas
   public PsiElement insertItemIntoList(PyElement list, @Nullable PyExpression afterThis, PyExpression toInsert)
     throws IncorrectOperationException {
     ASTNode add = toInsert.getNode().copyElement();
@@ -225,6 +257,7 @@ public class PyElementGeneratorImpl extends PyElementGenerator {
     return createExpressionFromText(LanguageLevel.getDefault(), text);
   }
 
+  @NotNull
   public PyExpression createExpressionFromText(final LanguageLevel languageLevel, final String text) {
     final PsiFile dummyFile = createDummyFile(languageLevel, text);
     final PsiElement element = dummyFile.getFirstChild();
@@ -364,4 +397,11 @@ public class PyElementGeneratorImpl extends PyElementGenerator {
     return createFromText(LanguageLevel.getDefault(),
                           PyExpressionStatement.class, content + "\n");
   }
+
+  private static class CommasOnly extends NotNullPredicate<LeafPsiElement> {
+    @Override
+    protected boolean applyNotNull(@NotNull final LeafPsiElement input) {
+      return input.getNode().getElementType().equals(PyTokenTypes.COMMA);
+    }
+  }
 }

python/testData/refactoring/argumentList/addKeyArgument/addArgumentFile.after.py

@@ -0,0 +1,32 @@
+from stub import *
+import stub
+
+
+class MyOldClass(metaclass=ABCMeta):
+    pass
+
+
+class MyNewClass(object,metaclass=ABCMeta):
+    pass
+
+
+class MyNewClass_2(object, datetime,metaclass=ABCMeta):
+    pass
+
+class NewClass_3(stub.object,metaclass=ABCMeta):
+    pass
+
+class NewClass_4(stub.object, stub.datetime,metaclass=ABCMeta):
+    pass
+
+class NewClass_5(stub.datetime, foo=stub.object,metaclass=ABCMeta):
+    pass
+
+
+spam = "new_param"
+
+my_function(new_param=spam)
+my_function_1("some_param",new_param=spam)
+my_function_2(named_param="ham",new_param=spam)
+my_function_3("some_param", "some_param_2", named_param="ham",new_param=spam)
+my_function_4("some_param", "some_param_2", named_param=stub.object, named_param_2="eggs",new_param=spam)

python/testData/refactoring/argumentList/addKeyArgument/addArgumentFile.py

@@ -0,0 +1,32 @@
+from stub import *
+import stub
+
+
+class MyOldClass():
+    pass
+
+
+class MyNewClass(object):
+    pass
+
+
+class MyNewClass_2(object, datetime):
+    pass
+
+class NewClass_3(stub.object):
+    pass
+
+class NewClass_4(stub.object, stub.datetime):
+    pass
+
+class NewClass_5(stub.datetime, foo=stub.object):
+    pass
+
+
+spam = "new_param"
+
+my_function()
+my_function_1("some_param")
+my_function_2(named_param="ham")
+my_function_3("some_param", "some_param_2", named_param="ham")
+my_function_4("some_param", "some_param_2", named_param=stub.object, named_param_2="eggs")

python/testData/refactoring/argumentList/addKeyArgument/stub.py

@@ -0,0 +1,13 @@
+class object: pass
+
+class datetime: pass
+
+
+class ABCMeta: pass
+
+
+def my_function(new_param="spam"): pass
+def my_function_1(first_param,some_param, new_param="spam"): pass
+def my_function_2(first_param,named_param="ham", new_param="spam"): pass
+def my_function_3(first_param,some_param, some_param_2, named_param="ham", new_param="spam"): pass
+def my_function_4(first_param,some_param, some_param_2, named_param="ham", named_param_2="eggs", new_param="spam"): pass

python/testData/refactoring/argumentList/addParam/addArgumentFile.after.py

@@ -0,0 +1,27 @@
+from stub import *
+import stub
+
+
+class MyOldClass(SuperClass):
+    pass
+
+
+class MyNewClass(object,SuperClass):
+    pass
+
+
+class MyNewClass_2(object, datetime,SuperClass):
+    pass
+
+class NewClass_3(stub.object,SuperClass):
+    pass
+
+class NewClass_4(stub.object, stub.datetime,SuperClass):
+    pass
+
+new_param = "ogg"
+
+my_function(new_param,some_param="spam")
+my_function(new_param)
+my_function_2("some_param",new_param)
+my_function_3(new_param,some_param="spam",some_another_param=stub.object)

python/testData/refactoring/argumentList/addParam/addArgumentFile.py

@@ -0,0 +1,27 @@
+from stub import *
+import stub
+
+
+class MyOldClass():
+    pass
+
+
+class MyNewClass(object):
+    pass
+
+
+class MyNewClass_2(object, datetime):
+    pass
+
+class NewClass_3(stub.object):
+    pass
+
+class NewClass_4(stub.object, stub.datetime):
+    pass
+
+new_param = "ogg"
+
+my_function(some_param="spam")
+my_function()
+my_function_2("some_param")
+my_function_3(some_param="spam",some_another_param=stub.object)

python/testData/refactoring/argumentList/addParam/stub.py

@@ -0,0 +1,13 @@
+class object: pass
+
+class datetime: pass
+
+
+class ABCMeta: pass
+
+class SuperClass: pass
+
+
+def my_function(new_param,some_param="spam"): pass
+def my_function_2(new_param,some_param): pass
+def my_function_3(new_param,some_param="spam",some_another_param="eggs"): pass

python/testSrc/com/jetbrains/python/psi/impl/PyArgumentListImplTest.java

@@ -0,0 +1,89 @@
+package com.jetbrains.python.psi.impl;
+
+import com.intellij.openapi.command.WriteCommandAction;
+import com.intellij.psi.impl.source.PostprocessReformattingAspect;
+import com.intellij.psi.util.PsiTreeUtil;
+import com.jetbrains.python.psi.*;
+import com.jetbrains.python.refactoring.classes.PyClassRefactoringTest;
+import org.jetbrains.annotations.NotNull;
+
+/**
+ * Tests {@link com.jetbrains.python.psi.impl.PyArgumentListImpl#addArgument(com.jetbrains.python.psi.PyExpression)}
+ *
+ * @author Ilya.Kazakevich
+ */
+public class PyArgumentListImplTest extends PyClassRefactoringTest {
+  private PyElementGeneratorImpl myGenerator;
+  private LanguageLevel myLanguagelevel;
+
+  public PyArgumentListImplTest() {
+    super("argumentList");
+  }
+
+  @Override
+  public void setUp() throws Exception {
+    super.setUp();
+    myGenerator = new PyElementGeneratorImpl(myFixture.getProject());
+    myLanguagelevel = LanguageLevel.PYTHON34;
+    setLanguageLevel(myLanguagelevel);
+  }
+
+  /**
+   * Ensures new keyword argument is set into appropriate place
+   */
+  public void testAddKeyArgument() throws Exception {
+    final PyKeywordArgument classKeyword = myGenerator.createKeywordArgument(myLanguagelevel, "metaclass", "ABCMeta");
+    final PyKeywordArgument functionKeyword = myGenerator.createKeywordArgument(myLanguagelevel, "new_param", "spam");
+
+
+    doTest(classKeyword, functionKeyword);
+  }
+
+  /**
+   * Ensures new param (NOT keyword argument) is set into appropriate place
+   */
+  public void testAddParam() throws Exception {
+    final PyExpression classParameter = myGenerator.createParameter("SuperClass");
+    final PyExpression functionParameter = myGenerator.createParameter("new_param");
+
+
+    doTest(classParameter, functionParameter);
+  }
+
+  /**
+   * Adds expressions to the superclass list and to the function calls in file
+   *
+   * @param superClassExpression expressions to add to the list of superclasses to any class found on file
+   * @param functionExpression   expressions to add to any function call found in file
+   */
+  private void doTest(@NotNull final PyExpression superClassExpression, @NotNull final PyExpression functionExpression) {
+    configureMultiFile("addArgumentFile", "stub");
+    myFixture.configureByFile(getMultiFileBaseName() + "/addArgumentFile.py");
+
+    //TODO: newly created expressions has no indent info, it leads to errors in postprocessing formatting. Need to investigate.
+    PostprocessReformattingAspect.getInstance(myFixture.getProject()).disablePostprocessFormattingInside(new Runnable() {
+      @Override
+      public void run() {
+        WriteCommandAction.runWriteCommandAction(myFixture.getProject(), new Runnable() {
+          @Override
+          public void run() {
+
+
+            for (final PyClass aClass : PsiTreeUtil.findChildrenOfType(myFixture.getFile(), PyClass.class)) {
+              final PyArgumentList superClassExpressionList = aClass.getSuperClassExpressionList();
+              assert superClassExpressionList != null : "Class has no expression list!";
+              superClassExpressionList.addArgument(superClassExpression);
+            }
+
+            for (final PyCallExpression expression : PsiTreeUtil.findChildrenOfType(myFixture.getFile(), PyCallExpression.class)) {
+              final PyArgumentList list = expression.getArgumentList();
+              assert list != null : "Callable statement has no argument list?";
+              list.addArgument(functionExpression);
+            }
+          }
+        });
+      }
+    });
+    myFixture.checkResultByFile(getMultiFileBaseName() + "/addArgumentFile.after.py");
+  }
+}