com.jetbrains.python.inspections.quickfix.AddCallSuperQuickFix Maven / Gradle / Ivy

Go to download
Show more of this group Show more artifacts with this name
Show all versions of python-community Show documentation
A packaging of the IntelliJ Community Edition python-community library. This is release number 1 of trunk branch 142.
The newest version!
/*
 * 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.inspections.quickfix;

import com.intellij.codeInspection.LocalQuickFix;
import com.intellij.codeInspection.ProblemDescriptor;
import com.intellij.openapi.project.Project;
import com.intellij.openapi.util.Couple;
import com.intellij.openapi.util.text.StringUtil;
import com.intellij.psi.util.PsiTreeUtil;
import com.intellij.util.containers.ContainerUtil;
import com.jetbrains.python.PyBundle;
import com.jetbrains.python.PyNames;
import com.jetbrains.python.psi.*;
import com.jetbrains.python.psi.impl.PyPsiUtils;
import org.jetbrains.annotations.NonNls;
import org.jetbrains.annotations.NotNull;
import org.jetbrains.annotations.Nullable;

import java.util.*;

/**
 * For:
 * class B(A):
 * def __init__(self):
 * A.__init__(self)           #  inserted
 * print "Constructor B was called"
 * 
 * User: catherine
 */
public class AddCallSuperQuickFix implements LocalQuickFix {

  @NotNull
  public String getName() {
    return PyBundle.message("QFIX.add.super");
  }

  @NonNls
  @NotNull
  public String getFamilyName() {
    return getName();
  }

  public void applyFix(@NotNull final Project project, @NotNull final ProblemDescriptor descriptor) {
    final PyFunction problemFunction = PsiTreeUtil.getParentOfType(descriptor.getPsiElement(), PyFunction.class);
    if (problemFunction == null) return;
    final StringBuilder superCall = new StringBuilder();
    final PyClass klass = problemFunction.getContainingClass();
    if (klass == null) return;
    final PyClass[] superClasses = klass.getSuperClasses();
    if (superClasses.length == 0) return;

    final PyClass superClass = superClasses[0];
    final PyFunction superInit = superClass.findMethodByName(PyNames.INIT, true);
    if (superInit == null) return;

    final ParametersInfo origInfo = new ParametersInfo(problemFunction.getParameterList());
    final ParametersInfo superInfo = new ParametersInfo(superInit.getParameterList());
    final boolean addSelfToCall;

    if (klass.isNewStyleClass()) {
      addSelfToCall = false;
      if (LanguageLevel.forElement(klass).isPy3K()) {
        superCall.append("super().__init__(");
      }
      else {
        superCall.append("super(").append(klass.getName()).append(", ").append(getSelfParameterName(origInfo)).append(").__init__(");
      }
    }
    else {
      addSelfToCall = true;
      superCall.append(superClass.getName()).append(".__init__(");
    }
    final StringBuilder newFunction = new StringBuilder("def __init__(");

    final Couple> couple = buildNewFunctionParamsAndSuperInitCallArgs(origInfo, superInfo, addSelfToCall);
    StringUtil.join(couple.getFirst(), ", ", newFunction);
    newFunction.append(")");

    if (problemFunction.getAnnotation() != null) {
      newFunction.append(problemFunction.getAnnotation().getText());
    }
    newFunction.append(": pass");

    StringUtil.join(couple.getSecond(), ", ", superCall);
    superCall.append(")");

    final PyElementGenerator generator = PyElementGenerator.getInstance(project);
    final LanguageLevel languageLevel = LanguageLevel.forElement(problemFunction);
    final PyStatement callSuperStatement = generator.createFromText(languageLevel, PyStatement.class, superCall.toString());
    final PyParameterList newParameterList = generator.createFromText(languageLevel,
                                                                      PyParameterList.class,
                                                                      newFunction.toString(),
                                                                      new int[]{0, 3});
    problemFunction.getParameterList().replace(newParameterList);
    final PyStatementList statementList = problemFunction.getStatementList();
    PyUtil.addElementToStatementList(callSuperStatement, statementList, true);
    PyPsiUtils.removeRedundantPass(statementList);
  }

  @NotNull
  private static String getSelfParameterName(@NotNull ParametersInfo info) {
    final PyParameter selfParameter = info.getSelfParameter();
    if (selfParameter == null) {
      return PyNames.CANONICAL_SELF;
    }
    return StringUtil.defaultIfEmpty(selfParameter.getName(), PyNames.CANONICAL_SELF);
  }

  @NotNull
  private static Couple> buildNewFunctionParamsAndSuperInitCallArgs(@NotNull ParametersInfo origInfo,
                                                                                 @NotNull ParametersInfo superInfo,
                                                                                 boolean addSelfToCall) {
    final List newFunctionParams = new ArrayList();
    final List superCallArgs = new ArrayList();

    final PyParameter selfParameter = origInfo.getSelfParameter();
    if (selfParameter != null && StringUtil.isNotEmpty(selfParameter.getName())) {
      newFunctionParams.add(selfParameter.getText());
    }
    else {
      newFunctionParams.add(PyNames.CANONICAL_SELF);
    }

    if (addSelfToCall) {
      superCallArgs.add(getSelfParameterName(origInfo));
    }

    // Required parameters (not-keyword)
    for (PyParameter param : origInfo.getRequiredParameters()) {
      newFunctionParams.add(param.getText());
    }
    for (PyParameter param : superInfo.getRequiredParameters()) {
      // Special case as if base class has constructor __init__((a, b), c) and
      // subclass has constructor __init__(a, (b, c))
      final PyTupleParameter tupleParam = param.getAsTuple();
      if (tupleParam != null) {
        final List uniqueNames = collectParameterNames(tupleParam);
        final boolean hasDuplicates = uniqueNames.removeAll(origInfo.getAllParameterNames());
        if (hasDuplicates) {
          newFunctionParams.addAll(uniqueNames);
        }
        else {
          newFunctionParams.add(param.getText());
        }
        // Retain original structure of tuple parameter.
        // Note that tuple parameters cannot have annotations or nested default values, so it's syntactically safe
        superCallArgs.add(param.getText());
      }
      else {
        if (!origInfo.getAllParameterNames().contains(param.getName())) {
          newFunctionParams.add(param.getText());
        }
        superCallArgs.add(param.getName());
      }
    }

    // Optional parameters (not-keyword)
    for (PyParameter param : origInfo.getOptionalParameters()) {
      newFunctionParams.add(param.getText());
    }

    // Positional vararg
    PyParameter starredParam = null;
    if (origInfo.getPositionalContainerParameter() != null) {
      starredParam = origInfo.getPositionalContainerParameter();
    }
    else if (superInfo.getPositionalContainerParameter() != null) {
      starredParam = superInfo.getPositionalContainerParameter();
    }
    else if (origInfo.getSingleStarParameter() != null) {
      starredParam = origInfo.getSingleStarParameter();
    }
    else if (superInfo.getSingleStarParameter() != null) {
      starredParam = superInfo.getSingleStarParameter();
    }
    if (starredParam != null) {
      newFunctionParams.add(starredParam.getText());
      if (superInfo.getPositionalContainerParameter() != null) {
        superCallArgs.add("*" + starredParam.getName());
      }
    }

    // Required keyword-only parameters
    boolean hasKeywordOnlyParams = false;
    for (PyParameter param : origInfo.getRequiredKeywordOnlyParameters()) {
      newFunctionParams.add(param.getText());
      hasKeywordOnlyParams = true;
    }
    for (PyParameter param : superInfo.getRequiredKeywordOnlyParameters()) {
      if (!origInfo.getAllParameterNames().contains(param.getName())) {
        newFunctionParams.add(param.getText());
        hasKeywordOnlyParams = true;
      }
      superCallArgs.add(param.getName() + "=" + param.getName());
    }
    if (starredParam instanceof PySingleStarParameter && !hasKeywordOnlyParams) {
      newFunctionParams.remove(newFunctionParams.size() - 1);
    }

    // Optional keyword-only parameters
    for (PyParameter param : origInfo.getOptionalKeywordOnlyParameters()) {
      newFunctionParams.add(param.getText());
    }

    // Keyword vararg
    PyParameter doubleStarredParam = null;
    if (origInfo.getKeywordContainerParameter() != null) {
      doubleStarredParam = origInfo.getKeywordContainerParameter();
    }
    else if (superInfo.getKeywordContainerParameter() != null) {
      doubleStarredParam = superInfo.getKeywordContainerParameter();
    }
    if (doubleStarredParam != null) {
      newFunctionParams.add(doubleStarredParam.getText());
      if (superInfo.getKeywordContainerParameter() != null) {
        superCallArgs.add("**" + doubleStarredParam.getName());
      }
    }
    return Couple.of(newFunctionParams, superCallArgs);
  }

  private static class ParametersInfo {

    private final PyParameter mySelfParam;
    /**
     * Parameters without default value that come before first "*..." parameter.
     */
    private final List myRequiredParams = new ArrayList();
    /**
     * Parameters with default value that come before first "*..." parameter.
     */
    private final List myOptionalParams = new ArrayList();
    /**
     * Parameter of form "*args" (positional vararg), not the same as single "*".
     */
    private final PyParameter myPositionalContainerParam;
    /**
     * Parameter "*", that is used to delimit normal and keyword-only parameters.
     */
    private final PyParameter mySingleStarParam;
    /**
     * Parameters without default value that come after first "*..." parameter.
     */
    private final List myRequiredKwOnlyParams = new ArrayList();
    /**
     * Parameters with default value that come after first "*..." parameter.
     */
    private final List myOptionalKwOnlyParams = new ArrayList();
    /**
     * Parameter of form "**kwargs" (keyword vararg).
     */
    private final PyParameter myKeywordContainerParam;

    private final Set myAllParameterNames = new LinkedHashSet();

    public ParametersInfo(@NotNull PyParameterList parameterList) {
      PyParameter positionalContainer = null;
      PyParameter singleStarParam = null;
      PyParameter keywordContainer = null;
      PyParameter selfParam = null;

      for (PyParameter param : parameterList.getParameters()) {
        myAllParameterNames.addAll(collectParameterNames(param));

        if (param.isSelf()) {
          selfParam = param;
        }
        else if (param instanceof PySingleStarParameter) {
          singleStarParam = param;
        }
        else if (param.getAsNamed() != null && param.getAsNamed().isKeywordContainer()) {
          keywordContainer = param;
        }
        else if (param.getAsNamed() != null && param.getAsNamed().isPositionalContainer()) {
          positionalContainer = param;
        }
        else if (param.getAsNamed() == null || !param.getAsNamed().isKeywordOnly()) {
          if (param.hasDefaultValue()) {
            myOptionalParams.add(param);
          }
          else {
            myRequiredParams.add(param);
          }
        }
        else {
          if (param.hasDefaultValue()) {
            myOptionalKwOnlyParams.add(param);
          }
          else {
            myRequiredKwOnlyParams.add(param);
          }
        }
      }

      mySelfParam = selfParam;
      myPositionalContainerParam = positionalContainer;
      mySingleStarParam = singleStarParam;
      myKeywordContainerParam = keywordContainer;
    }

    @Nullable
    public PyParameter getSelfParameter() {
      return mySelfParam;
    }

    @NotNull
    public List getRequiredParameters() {
      return Collections.unmodifiableList(myRequiredParams);
    }

    @NotNull
    public List getOptionalParameters() {
      return Collections.unmodifiableList(myOptionalParams);
    }

    @Nullable
    public PyParameter getPositionalContainerParameter() {
      return myPositionalContainerParam;
    }

    @Nullable
    public PyParameter getSingleStarParameter() {
      return mySingleStarParam;
    }

    @NotNull
    public List getRequiredKeywordOnlyParameters() {
      return Collections.unmodifiableList(myRequiredKwOnlyParams);
    }

    @NotNull
    public List getOptionalKeywordOnlyParameters() {
      return Collections.unmodifiableList(myOptionalKwOnlyParams);
    }

    @Nullable
    public PyParameter getKeywordContainerParameter() {
      return myKeywordContainerParam;
    }

    @NotNull
    public Set getAllParameterNames() {
      return Collections.unmodifiableSet(myAllParameterNames);
    }
  }

  @NotNull
  private static List collectParameterNames(@NotNull PyParameter param) {
    final List result = new ArrayList();
    collectParameterNames(param, result);
    return result;
  }


  private static void collectParameterNames(@NotNull PyParameter param, @NotNull Collection acc) {
    final PyTupleParameter tupleParam = param.getAsTuple();
    if (tupleParam != null) {
      for (PyParameter subParam : tupleParam.getContents()) {
        collectParameterNames(subParam, acc);
      }
    }
    else {
      ContainerUtil.addIfNotNull(acc, param.getName());
    }
  }
}