python/src/com/jetbrains/python/psi/impl/PyArgumentListImpl.java - platform/tools/idea - Git at Google

 /*
  * Copyright 2000-2013 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.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.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.FunctionParameter;
 import com.jetbrains.python.PyElementTypes;
 import com.jetbrains.python.PyTokenTypes;
 import com.jetbrains.python.PythonDialectsTokenSetProvider;
 import com.jetbrains.python.psi.*;
 import com.jetbrains.python.psi.resolve.PyResolveContext;
 import org.jetbrains.annotations.NotNull;
 import org.jetbrains.annotations.Nullable;

 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);
   }

   @Override
   protected void acceptPyVisitor(PyElementVisitor pyVisitor) {
     pyVisitor.visitPyArgumentList(this);
   }

   @Override
   @NotNull
   public Collection<PyExpression> getArgumentExpressions() {
     final PyExpression[] arguments = getArguments();
     final Collection<PyExpression> result = new ArrayList<PyExpression>(arguments.length);
     for (final PyExpression expression : arguments) {
       if (expression instanceof PyKeywordArgument) {
         final PyExpression valueExpression = ((PyKeywordArgument)expression).getValueExpression();
         result.add(valueExpression);
       }
       if (expression instanceof PyReferenceExpression) {
         result.add(expression);
       }
     }
     return result;
   }

   @NotNull
   public PyExpression[] getArguments() {
     return childrenToPsi(PythonDialectsTokenSetProvider.INSTANCE.getExpressionTokens(), PyExpression.EMPTY_ARRAY);
   }

   @Nullable
   public PyKeywordArgument getKeywordArgument(String name) {
     ASTNode node = getNode().getFirstChildNode();
     while (node != null) {
       if (node.getElementType() == PyElementTypes.KEYWORD_ARGUMENT_EXPRESSION) {
         PyKeywordArgument arg = (PyKeywordArgument)node.getPsi();
         String keyword = arg.getKeyword();
         if (keyword != null && keyword.equals(name)) return arg;
       }
       node = node.getTreeNext();
     }
     return null;
   }

   @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) {
       if (parameters.isEmpty()) {
         generator.insertItemIntoListRemoveRedundantCommas(this, keywordArguments.getLast(), arg);
       }
       else {
         if (keywordArguments.isEmpty()) {
           generator.insertItemIntoListRemoveRedundantCommas(this, parameters.getLast(), arg);
         }
         else {
           generator.insertItemIntoListRemoveRedundantCommas(this, keywordArguments.getLast(), arg);
         }
       }
     }
     else {
       if (parameters.isEmpty()) {
         generator.insertItemIntoListRemoveRedundantCommas(this, null, arg);
       }
       else {
         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));
     if (pars.length == 0) {
       // there's no starting paren
       try {
         add(arg);
       }
       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 = createComma();
         node.addChild(comma, before);
         node.addChild(ASTFactory.whitespace(" "), before);
         anchorBefore = comma;
       }
       else {
         anchorBefore = before;
       }
       ASTNode argNode = arg.getNode();
       if (anchorBefore == null) {
         node.addChild(argNode);
       }
       else {
         node.addChild(argNode, anchorBefore);
       }
     }
   }

   /**
    * @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) {
       if (expressions.contains(node.getElementType())) return true;
       node = node.getTreeNext();
     }
     return false;
   }

   private void addArgumentLastWithoutComma(PyExpression arg) {
     ASTNode par = getClosingParen();
     if (par == null) {
       // there's no ending paren
       try {
         add(arg);
       }
       catch (IncorrectOperationException e1) {
         throw new IllegalStateException(e1);
       }
     }
     else {
       getNode().addChild(arg.getNode(), par);
     }
   }

   @Nullable
   public ASTNode getClosingParen() {
     ASTNode node = getNode();
     final ASTNode[] children = node.getChildren(TokenSet.create(PyTokenTypes.RPAR));
     return children.length == 0 ? null : children[children.length - 1];
   }

   private void addArgumentNode(PyExpression arg, ASTNode beforeThis, boolean commaFirst) {
     ASTNode comma = PyElementGenerator.getInstance(getProject()).createComma();
     ASTNode node = getNode();
     ASTNode argNode = arg.getNode();
     if (commaFirst) {
       node.addChild(comma, beforeThis);
       node.addChild(ASTFactory.whitespace(" "), beforeThis);
       node.addChild(argNode, beforeThis);
     }
     else {
       node.addChild(argNode, beforeThis);
       node.addChild(comma, beforeThis);
       node.addChild(ASTFactory.whitespace(" "), beforeThis);
     }
   }

   public void addArgumentAfter(PyExpression argument, @Nullable PyExpression afterThis) {
     if (afterThis == null) {
       addArgumentFirst(argument);
       return;
     }
     boolean good = false;
     for (PyExpression expression : getArguments()) {
       if (expression == afterThis) {
         good = true;
         break;
       }
     }
     if (!good) {
       throw new IllegalArgumentException("Expression " + afterThis + " is not an argument (" + Arrays.toString(getArguments()) + ")");
     }
     // CASES:
     ASTNode node = afterThis.getNode().getTreeNext();
     while (node != null) {
       IElementType type = node.getElementType();
       if (type == PyTokenTypes.RPAR) {
         // 1: Nothing, just add
         addArgumentNode(argument, node, true);
         break;
       }
       else if (PythonDialectsTokenSetProvider.INSTANCE.getExpressionTokens().contains(type)) {
         // 2: After some argument followed by comma: after comma, add element, add comma
         // 3: After some argument not followed by comma: add comma, add element
         addArgumentNode(argument, node, true);
         break;
       }
       else if (type == PyTokenTypes.COMMA) {
         ASTNode next = PyUtil.getNextNonWhitespace(node);
         if (next == null) {
           addArgumentLastWithoutComma(argument);
         }
         else if (next.getElementType() == PyTokenTypes.RPAR) {
           addArgumentNode(argument, next, false);
         }
         else {
           addArgumentNode(argument, next, false);
         }
         break;
       }
       node = node.getTreeNext();
     }
   }

   @Nullable
   public PyCallExpression getCallExpression() {
     return PsiTreeUtil.getParentOfType(this, PyCallExpression.class);
   }

   @Override
   public void deleteChildInternal(@NotNull ASTNode node) {
     //noinspection SuspiciousMethodCalls
     if (Arrays.asList(getArguments()).contains(node.getPsi())) {
       ASTNode next = PyPsiUtils.getNextComma(node);
       if (next == null) {
         next = PyPsiUtils.getPrevComma(node);
       }
       if (next != null) {
         deleteChildInternal(next);
       }
     }
     super.deleteChildInternal(node);
   }

   @NotNull
   public CallArgumentsMapping analyzeCall(PyResolveContext resolveContext) {
     return analyzeCall(resolveContext, 0);
   }

   @NotNull
   public CallArgumentsMapping analyzeCall(PyResolveContext resolveContext, int offset) {
     final CallArgumentsMappingImpl ret = new CallArgumentsMappingImpl(this);
     // declaration-based checks
     // proper arglist is: [positional,...][name=value,...][*tuple,][**dict]
     // where "positional" may be a tuple of nested "positional" parameters, too.
     // following the spec: http://docs.python.org/ref/calls.html
     PyCallExpression call = getCallExpression();
     if (call != null) {
       PyCallExpression.PyMarkedCallee resolvedCallee = call.resolveCallee(resolveContext, offset);
       if (resolvedCallee != null) {
         ret.mapArguments(resolvedCallee, resolveContext.getTypeEvalContext());
       }
       else {
         ret.verifyArguments();
       }
     }
     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);
     }
   }

   @Nullable
   @Override
   public PyExpression getValueExpressionForParam(@NotNull final FunctionParameter parameter) {
     final String parameterName = parameter.getName();
     if (parameterName != null) {
       final PyKeywordArgument kwarg = getKeywordArgument(parameterName);
       if (kwarg != null) {
         return kwarg.getValueExpression();
       }
     }

     final PyExpression[] arguments = getArguments();
     final int position = parameter.getPosition();
     if ((position != FunctionParameter.POSITION_NOT_SUPPORTED) && (arguments.length > position)) {
       final PyExpression result = arguments[position];
       if (result instanceof PyKeywordArgument) {
         ((PyKeywordArgument)result).getValueExpression();
       }
       else {
         return result;
       }
     }

     return null;
   }
 }
	/*
	* Copyright 2000-2013 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.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.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.FunctionParameter;
	import com.jetbrains.python.PyElementTypes;
	import com.jetbrains.python.PyTokenTypes;
	import com.jetbrains.python.PythonDialectsTokenSetProvider;
	import com.jetbrains.python.psi.*;
	import com.jetbrains.python.psi.resolve.PyResolveContext;
	import org.jetbrains.annotations.NotNull;
	import org.jetbrains.annotations.Nullable;

	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);
	}

	@Override
	protected void acceptPyVisitor(PyElementVisitor pyVisitor) {
	pyVisitor.visitPyArgumentList(this);
	}

	@Override
	@NotNull
	public Collection<PyExpression> getArgumentExpressions() {
	final PyExpression[] arguments = getArguments();
	final Collection<PyExpression> result = new ArrayList<PyExpression>(arguments.length);
	for (final PyExpression expression : arguments) {
	if (expression instanceof PyKeywordArgument) {
	final PyExpression valueExpression = ((PyKeywordArgument)expression).getValueExpression();
	result.add(valueExpression);
	}
	if (expression instanceof PyReferenceExpression) {
	result.add(expression);
	}
	}
	return result;
	}

	@NotNull
	public PyExpression[] getArguments() {
	return childrenToPsi(PythonDialectsTokenSetProvider.INSTANCE.getExpressionTokens(), PyExpression.EMPTY_ARRAY);
	}

	@Nullable
	public PyKeywordArgument getKeywordArgument(String name) {
	ASTNode node = getNode().getFirstChildNode();
	while (node != null) {
	if (node.getElementType() == PyElementTypes.KEYWORD_ARGUMENT_EXPRESSION) {
	PyKeywordArgument arg = (PyKeywordArgument)node.getPsi();
	String keyword = arg.getKeyword();
	if (keyword != null && keyword.equals(name)) return arg;
	}
	node = node.getTreeNext();
	}
	return null;
	}

	@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) {
	if (parameters.isEmpty()) {
	generator.insertItemIntoListRemoveRedundantCommas(this, keywordArguments.getLast(), arg);
	}
	else {
	if (keywordArguments.isEmpty()) {
	generator.insertItemIntoListRemoveRedundantCommas(this, parameters.getLast(), arg);
	}
	else {
	generator.insertItemIntoListRemoveRedundantCommas(this, keywordArguments.getLast(), arg);
	}
	}
	}
	else {
	if (parameters.isEmpty()) {
	generator.insertItemIntoListRemoveRedundantCommas(this, null, arg);
	}
	else {
	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));
	if (pars.length == 0) {
	// there's no starting paren
	try {
	add(arg);
	}
	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 = createComma();
	node.addChild(comma, before);
	node.addChild(ASTFactory.whitespace(" "), before);
	anchorBefore = comma;
	}
	else {
	anchorBefore = before;
	}
	ASTNode argNode = arg.getNode();
	if (anchorBefore == null) {
	node.addChild(argNode);
	}
	else {
	node.addChild(argNode, anchorBefore);
	}
	}
	}

	/**
	* @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) {
	if (expressions.contains(node.getElementType())) return true;
	node = node.getTreeNext();
	}
	return false;
	}

	private void addArgumentLastWithoutComma(PyExpression arg) {
	ASTNode par = getClosingParen();
	if (par == null) {
	// there's no ending paren
	try {
	add(arg);
	}
	catch (IncorrectOperationException e1) {
	throw new IllegalStateException(e1);
	}
	}
	else {
	getNode().addChild(arg.getNode(), par);
	}
	}

	@Nullable
	public ASTNode getClosingParen() {
	ASTNode node = getNode();
	final ASTNode[] children = node.getChildren(TokenSet.create(PyTokenTypes.RPAR));
	return children.length == 0 ? null : children[children.length - 1];
	}

	private void addArgumentNode(PyExpression arg, ASTNode beforeThis, boolean commaFirst) {
	ASTNode comma = PyElementGenerator.getInstance(getProject()).createComma();
	ASTNode node = getNode();
	ASTNode argNode = arg.getNode();
	if (commaFirst) {
	node.addChild(comma, beforeThis);
	node.addChild(ASTFactory.whitespace(" "), beforeThis);
	node.addChild(argNode, beforeThis);
	}
	else {
	node.addChild(argNode, beforeThis);
	node.addChild(comma, beforeThis);
	node.addChild(ASTFactory.whitespace(" "), beforeThis);
	}
	}

	public void addArgumentAfter(PyExpression argument, @Nullable PyExpression afterThis) {
	if (afterThis == null) {
	addArgumentFirst(argument);
	return;
	}
	boolean good = false;
	for (PyExpression expression : getArguments()) {
	if (expression == afterThis) {
	good = true;
	break;
	}
	}
	if (!good) {
	throw new IllegalArgumentException("Expression " + afterThis + " is not an argument (" + Arrays.toString(getArguments()) + ")");
	}
	// CASES:
	ASTNode node = afterThis.getNode().getTreeNext();
	while (node != null) {
	IElementType type = node.getElementType();
	if (type == PyTokenTypes.RPAR) {
	// 1: Nothing, just add
	addArgumentNode(argument, node, true);
	break;
	}
	else if (PythonDialectsTokenSetProvider.INSTANCE.getExpressionTokens().contains(type)) {
	// 2: After some argument followed by comma: after comma, add element, add comma
	// 3: After some argument not followed by comma: add comma, add element
	addArgumentNode(argument, node, true);
	break;
	}
	else if (type == PyTokenTypes.COMMA) {
	ASTNode next = PyUtil.getNextNonWhitespace(node);
	if (next == null) {
	addArgumentLastWithoutComma(argument);
	}
	else if (next.getElementType() == PyTokenTypes.RPAR) {
	addArgumentNode(argument, next, false);
	}
	else {
	addArgumentNode(argument, next, false);
	}
	break;
	}
	node = node.getTreeNext();
	}
	}

	@Nullable
	public PyCallExpression getCallExpression() {
	return PsiTreeUtil.getParentOfType(this, PyCallExpression.class);
	}

	@Override
	public void deleteChildInternal(@NotNull ASTNode node) {
	//noinspection SuspiciousMethodCalls
	if (Arrays.asList(getArguments()).contains(node.getPsi())) {
	ASTNode next = PyPsiUtils.getNextComma(node);
	if (next == null) {
	next = PyPsiUtils.getPrevComma(node);
	}
	if (next != null) {
	deleteChildInternal(next);
	}
	}
	super.deleteChildInternal(node);
	}

	@NotNull
	public CallArgumentsMapping analyzeCall(PyResolveContext resolveContext) {
	return analyzeCall(resolveContext, 0);
	}

	@NotNull
	public CallArgumentsMapping analyzeCall(PyResolveContext resolveContext, int offset) {
	final CallArgumentsMappingImpl ret = new CallArgumentsMappingImpl(this);
	// declaration-based checks
	// proper arglist is: [positional,...][name=value,...][tuple,][*dict]
	// where "positional" may be a tuple of nested "positional" parameters, too.
	// following the spec: http://docs.python.org/ref/calls.html
	PyCallExpression call = getCallExpression();
	if (call != null) {
	PyCallExpression.PyMarkedCallee resolvedCallee = call.resolveCallee(resolveContext, offset);
	if (resolvedCallee != null) {
	ret.mapArguments(resolvedCallee, resolveContext.getTypeEvalContext());
	}
	else {
	ret.verifyArguments();
	}
	}
	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);
	}
	}

	@Nullable
	@Override
	public PyExpression getValueExpressionForParam(@NotNull final FunctionParameter parameter) {
	final String parameterName = parameter.getName();
	if (parameterName != null) {
	final PyKeywordArgument kwarg = getKeywordArgument(parameterName);
	if (kwarg != null) {
	return kwarg.getValueExpression();
	}
	}

	final PyExpression[] arguments = getArguments();
	final int position = parameter.getPosition();
	if ((position != FunctionParameter.POSITION_NOT_SUPPORTED) && (arguments.length > position)) {
	final PyExpression result = arguments[position];
	if (result instanceof PyKeywordArgument) {
	((PyKeywordArgument)result).getValueExpression();
	}
	else {
	return result;
	}
	}

	return null;
	}
	}