From 40f6fbe45c8d73dc7ee7e20d30031d54a56b4a5a Mon Sep 17 00:00:00 2001 From: "Ilya.Kazakevich" Date: Fri, 14 Feb 2014 20:34:05 +0400 Subject: [PATCH] PyArgumentList#addArgument fixed and test added --- .../jetbrains/python/psi/PyArgumentList.java | 22 +++- .../python/psi/PyCallExpression.java | 4 + .../python/psi/PyElementGenerator.java | 15 +++ .../python/psi/impl/PyArgumentListImpl.java | 103 ++++++++++++------ .../psi/impl/PyCallExpressionHelper.java | 1 + .../psi/impl/PyElementGeneratorImpl.java | 40 +++++++ .../addKeyArgument/addArgumentFile.after.py | 32 ++++++ .../addKeyArgument/addArgumentFile.py | 32 ++++++ .../argumentList/addKeyArgument/stub.py | 13 +++ .../addParam/addArgumentFile.after.py | 27 +++++ .../argumentList/addParam/addArgumentFile.py | 27 +++++ .../refactoring/argumentList/addParam/stub.py | 13 +++ .../psi/impl/PyArgumentListImplTest.java | 89 +++++++++++++++ 13 files changed, 381 insertions(+), 37 deletions(-) create mode 100644 python/testData/refactoring/argumentList/addKeyArgument/addArgumentFile.after.py create mode 100644 python/testData/refactoring/argumentList/addKeyArgument/addArgumentFile.py create mode 100644 python/testData/refactoring/argumentList/addKeyArgument/stub.py create mode 100644 python/testData/refactoring/argumentList/addParam/addArgumentFile.after.py create mode 100644 python/testData/refactoring/argumentList/addParam/addArgumentFile.py create mode 100644 python/testData/refactoring/argumentList/addParam/stub.py create mode 100644 python/testSrc/com/jetbrains/python/psi/impl/PyArgumentListImplTest.java diff --git a/python/psi-api/src/com/jetbrains/python/psi/PyArgumentList.java b/python/psi-api/src/com/jetbrains/python/psi/PyArgumentList.java index 55a6506c728f..5ba21bf90556 100644 --- a/python/psi-api/src/com/jetbrains/python/psi/PyArgumentList.java +++ b/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); diff --git a/python/psi-api/src/com/jetbrains/python/psi/PyCallExpression.java b/python/psi-api/src/com/jetbrains/python/psi/PyCallExpression.java index 41d93bf81bba..71e31dd92624 100644 --- a/python/psi-api/src/com/jetbrains/python/psi/PyCallExpression.java +++ b/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); /** diff --git a/python/psi-api/src/com/jetbrains/python/psi/PyElementGenerator.java b/python/psi-api/src/com/jetbrains/python/psi/PyElementGenerator.java index 75a0bbce28b4..568cb95740b5 100644 --- a/python/psi-api/src/com/jetbrains/python/psi/PyElementGenerator.java +++ b/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; diff --git a/python/src/com/jetbrains/python/psi/impl/PyArgumentListImpl.java b/python/src/com/jetbrains/python/psi/impl/PyArgumentListImpl.java index 50ba5419070f..dfe12ca6dc30 100644 --- a/python/src/com/jetbrains/python/psi/impl/PyArgumentListImpl.java +++ b/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 keywordArguments = getKeyWordArguments(); + final Deque 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 getParameters() { + final PyExpression[] childrenOfType = PsiTreeUtil.getChildrenOfType(this, PyExpression.class); + if (childrenOfType == null) { + return new ArrayDeque(0); + } + return Queues.newArrayDeque(Collections2.filter(Arrays.asList(childrenOfType), NO_KEY_ARGUMENTS)); + } + + /** + * @return keyword arguments (as opposite to parameters) + */ + @NotNull + private Deque 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 { + @Override + protected boolean applyNotNull(@NotNull final PyExpression input) { + return (PsiTreeUtil.getParentOfType(input, PyKeywordArgument.class) == null) && !(input instanceof PyKeywordArgument); + } + } } diff --git a/python/src/com/jetbrains/python/psi/impl/PyCallExpressionHelper.java b/python/src/com/jetbrains/python/psi/impl/PyCallExpressionHelper.java index 19fa72a1a97b..781dc28c0d35 100644 --- a/python/src/com/jetbrains/python/psi/impl/PyCallExpressionHelper.java +++ b/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 diff --git a/python/src/com/jetbrains/python/psi/impl/PyElementGeneratorImpl.java b/python/src/com/jetbrains/python/psi/impl/PyElementGeneratorImpl.java index 1a8f123a6aff..c306116156a0 100644 --- a/python/src/com/jetbrains/python/psi/impl/PyElementGeneratorImpl.java +++ b/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 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 { + @Override + protected boolean applyNotNull(@NotNull final LeafPsiElement input) { + return input.getNode().getElementType().equals(PyTokenTypes.COMMA); + } + } } diff --git a/python/testData/refactoring/argumentList/addKeyArgument/addArgumentFile.after.py b/python/testData/refactoring/argumentList/addKeyArgument/addArgumentFile.after.py new file mode 100644 index 000000000000..7a9fbe968399 --- /dev/null +++ b/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) \ No newline at end of file diff --git a/python/testData/refactoring/argumentList/addKeyArgument/addArgumentFile.py b/python/testData/refactoring/argumentList/addKeyArgument/addArgumentFile.py new file mode 100644 index 000000000000..32b4fda7409e --- /dev/null +++ b/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") \ No newline at end of file diff --git a/python/testData/refactoring/argumentList/addKeyArgument/stub.py b/python/testData/refactoring/argumentList/addKeyArgument/stub.py new file mode 100644 index 000000000000..9227a3542e01 --- /dev/null +++ b/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 \ No newline at end of file diff --git a/python/testData/refactoring/argumentList/addParam/addArgumentFile.after.py b/python/testData/refactoring/argumentList/addParam/addArgumentFile.after.py new file mode 100644 index 000000000000..97c18fd7f85b --- /dev/null +++ b/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) \ No newline at end of file diff --git a/python/testData/refactoring/argumentList/addParam/addArgumentFile.py b/python/testData/refactoring/argumentList/addParam/addArgumentFile.py new file mode 100644 index 000000000000..32d0235e1266 --- /dev/null +++ b/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) \ No newline at end of file diff --git a/python/testData/refactoring/argumentList/addParam/stub.py b/python/testData/refactoring/argumentList/addParam/stub.py new file mode 100644 index 000000000000..b98bdb0d1fca --- /dev/null +++ b/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 diff --git a/python/testSrc/com/jetbrains/python/psi/impl/PyArgumentListImplTest.java b/python/testSrc/com/jetbrains/python/psi/impl/PyArgumentListImplTest.java new file mode 100644 index 000000000000..7aa1a5c0864c --- /dev/null +++ b/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"); + } +}