added completion resolving kwargs from code usage (PY-1002)

2026-09-27 10:03:11 +07:00 · 2010-06-01 22:35:00 +04:00
parent 2717675c3b
commit 5565b19ea3
7 changed files with 198 additions and 38 deletions
@@ -1,6 +1,9 @@
 package com.jetbrains.python.psi;

 import com.intellij.codeInsight.CodeInsightUtilBase;
+import com.intellij.codeInsight.completion.PrioritizedLookupElement;
+import com.intellij.codeInsight.lookup.LookupElement;
+import com.intellij.codeInsight.lookup.LookupElementBuilder;
 import com.intellij.lang.ASTNode;
 import com.intellij.openapi.project.Project;
 import com.intellij.openapi.ui.MessageType;
@@ -14,6 +17,7 @@ import com.intellij.psi.tree.TokenSet;
 import com.intellij.psi.util.PsiTreeUtil;
 import com.intellij.ui.awt.RelativePoint;
 import com.intellij.util.ArrayUtil;
+import com.intellij.util.Icons;
 import com.intellij.util.IncorrectOperationException;
 import com.intellij.util.SmartList;
 import com.intellij.util.containers.HashSet;
@@ -112,6 +116,7 @@ public class PyUtil {
  }

  // Poor man's catamorhpism :)
+
  /**
   * Flattens the representation of every element in targets, and puts all results together.
   * Elements of every tuple nested in target item are brought to the top level: (a, (b, (c, d))) -> (a, b, c, d)
@@ -127,6 +132,7 @@ public class PyUtil {

  // Poor man's filter
  // TODO: move to a saner place
+
  public static boolean instanceOf(Object obj, Class... possibleClasses) {
    for (Class cls : possibleClasses) {
      if (cls.isInstance(obj)) return true;
@@ -137,7 +143,8 @@ public class PyUtil {

  /**
   * Produce a reasonable representation of a PSI element, good for debugging.
-   * @param elt element to represent; nulls and invalid nodes are ok.
+   *
+   * @param elt      element to represent; nulls and invalid nodes are ok.
   * @param cutAtEOL if true, representation stops at nearest EOL inside the element.
   * @return the representation.
   */
@@ -184,6 +191,7 @@ public class PyUtil {

  /**
   * Shows an information balloon in a reasonable place at the top right of the window.
+   *
   * @param project     our project
   * @param message     the text, HTML markup allowed
   * @param messageType message type, changes the icon and the background.
@@ -209,12 +217,17 @@ public class PyUtil {
   * Returns a quoted string representation, or "null".
   */
  public static String nvl(Object s) {
-    if (s != null) return "'" + s.toString() + "'";
-    else return "null";
+    if (s != null) {
+      return "'" + s.toString() + "'";
+    }
+    else {
+      return "null";
+    }
  }

  /**
-   * Adds an item into a comma-separated list in a PSI tree. E.g. can turn "foo, bar" into "foo, bar, baz", adding commas as needed. 
+   * Adds an item into a comma-separated list in a PSI tree. E.g. can turn "foo, bar" into "foo, bar, baz", adding commas as needed.
+   *
   * @param parent     the element to represent the list; we're adding a child to it.
   * @param newItem    the element we're inserting (the "baz" in the example).
   * @param beforeThis node to mark the insertion point inside the list; must belong to a child of target. Set to null to add first element.
@@ -239,6 +252,7 @@ public class PyUtil {
  /**
   * Removes an element from a a comma-separated list in a PSI tree. E.g. can turn "foo, bar, baz" into "foo, baz",
   * removing commas as needed. It removes a trailing comma if it results from deletion.
+   *
   * @param item what to remove. Its parent is considered the list, and commas must be its peers.
   */
  public static void removeListNode(PsiElement item) {
@@ -250,7 +264,7 @@ public class PyUtil {
    ASTNode binder = parent.getNode();
    assert binder != null : "parent node is null, ensureWritable() lied";
    boolean got_comma_after = eraseWhitespaceAndComma(binder, item, false);
-    if (! got_comma_after) {
+    if (!got_comma_after) {
      // there was not a comma after the item; remove a comma before the item
      eraseWhitespaceAndComma(binder, item, true);
    }
@@ -260,6 +274,7 @@ public class PyUtil {

  /**
   * Removes whitespace and comma(s) that are siblings of the item, up to the first non-whitespace and non-comma.
+   *
   * @param parent_node node of the parent of item.
   * @param item        starting point; we erase left or right of it, but not it.
   * @param backwards   true to erase prev siblings, false to erase next siblings.
@@ -274,14 +289,20 @@ public class PyUtil {
    boolean have_skipped_the_item = false;
    while (current != null) {
      candidate = current;
-      current = backwards? current.getTreePrev() : current.getTreeNext();
+      current = backwards ? current.getTreePrev() : current.getTreeNext();
      if (have_skipped_the_item) {
        is_comma = ",".equals(candidate.getText());
        got_comma |= is_comma;
-        if (is_comma || candidate.getElementType() == TokenType.WHITE_SPACE) parent_node.removeChild(candidate);
-        else break;
+        if (is_comma || candidate.getElementType() == TokenType.WHITE_SPACE) {
+          parent_node.removeChild(candidate);
+        }
+        else {
+          break;
+        }
+      }
+      else {
+        have_skipped_the_item = true;
      }
-      else have_skipped_the_item = true;
    }
    return got_comma;
  }
@@ -320,10 +341,12 @@ public class PyUtil {

  /**
   * Finds the first identifier AST node under target element, and returns its text.
+   *
   * @param target
   * @return identifier text, or null.
   */
-  public static @Nullable
+  public static
+  @Nullable
  String getIdentifier(PsiElement target) {
    ASTNode node = target.getNode();
    if (node != null) {
@@ -335,8 +358,10 @@ public class PyUtil {


  // TODO: move to a more proper place?
+
  /**
   * Determine the type of a special attribute. Currently supported: {@code __class__} and {@code __dict__}.
+   *
   * @param ref reference to a possible attribute; only qualified references make sense.
   * @return type, or null (if type cannot be determined, reference is not to a known attribute, etc.)
   */
@@ -365,6 +390,7 @@ public class PyUtil {

  /**
   * Makes sure that 'thing' is not null; else throws an {@link IncorrectOperationException}.
+   *
   * @param thing what we check.
   * @return thing, if not null.
   */
@@ -376,6 +402,7 @@ public class PyUtil {

  /**
   * Makes sure that the 'thing' is true; else throws an {@link IncorrectOperationException}.
+   *
   * @param thing what we check.
   */
  public static void sure(boolean thing) {
@@ -385,10 +412,13 @@ public class PyUtil {
  /**
   * For cases when a function is decorated with only one decorator, and this is a built-in decorator.
   * <br/> <i>TODO: handle multiple decorators sensibly; then rename and move.</i>
+   *
   * @param node the allegedly decorated function
   * @return name of the only built-in decorator, or null (even if there are multiple or non-built-in decorators!)
   */
-  public static @Nullable String getTheOnlyBuiltinDecorator(@NotNull final PyFunction node) {
+  public static
+  @Nullable
+  String getTheOnlyBuiltinDecorator(@NotNull final PyFunction node) {
    PyDecoratorList decolist = node.getDecoratorList();
    if (decolist != null) {
      PyDecorator[] decos = decolist.getDecorators();
@@ -406,6 +436,7 @@ public class PyUtil {

  /**
   * Looks for two standard decorators to a function, or a wrapping assignment that closely follows it.
+   *
   * @param function what to analyze
   * @return a set of flags describing what was detected.
   */
@@ -413,13 +444,15 @@ public class PyUtil {
  public static Set<PyFunction.Flag> detectDecorationsAndWrappersOf(PyFunction function) {
    Set<PyFunction.Flag> flags = EnumSet.noneOf(PyFunction.Flag.class);
    String deconame = getTheOnlyBuiltinDecorator(function);
-    if (PyNames.CLASSMETHOD.equals(deconame)) flags.add(CLASSMETHOD);
+    if (PyNames.CLASSMETHOD.equals(deconame)) {
+      flags.add(CLASSMETHOD);
+    }
    else if (PyNames.STATICMETHOD.equals(deconame)) flags.add(STATICMETHOD);
    // implicit classmethod __new__
    PyClass cls = function.getContainingClass();
    if (cls != null && PyNames.NEW.equals(function.getName()) && cls.isNewStyleClass()) flags.add(CLASSMETHOD);
    //
-    if (! flags.contains(CLASSMETHOD) && ! flags.contains(STATICMETHOD)) { // not set by decos, look for reassignment
+    if (!flags.contains(CLASSMETHOD) && !flags.contains(STATICMETHOD)) { // not set by decos, look for reassignment
      String func_name = function.getName();
      if (func_name != null) {
        PyAssignmentStatement assignment = PsiTreeUtil.getNextSiblingOfType(function, PyAssignmentStatement.class);
@@ -434,7 +467,9 @@ public class PyUtil {
                  PyFunction original = interpreted.getSecond();
                  if (original == function) {
                    String wrapper_name = interpreted.getFirst();
-                    if (PyNames.CLASSMETHOD.equals(wrapper_name)) flags.add(CLASSMETHOD);
+                    if (PyNames.CLASSMETHOD.equals(wrapper_name)) {
+                      flags.add(CLASSMETHOD);
+                    }
                    else if (PyNames.STATICMETHOD.equals(wrapper_name)) flags.add(STATICMETHOD);
                    flags.add(WRAPPED);
                  }
@@ -461,7 +496,7 @@ public class PyUtil {
    final ASTNode node = element.getNode();
    if (node != null) {
      final ASTNode[] children = node.getChildren(filter);
-      return (0 <= number && number < children.length) ? children [number].getPsi() : null;
+      return (0 <= number && number < children.length) ? children[number].getPsi() : null;
    }
    return null;
  }
@@ -472,6 +507,7 @@ public class PyUtil {
   * Alas, resolve() and multiResolve() can't return anything but a PyFile or PsiFileImpl.isPsiUpToDate() would fail.
   * This is because isPsiUpToDate() relies on identity of objects returned by FileViewProvider.getPsi().
   * If we ever need to exactly tell a dir from __init__.py, that logic has to change.
+   *
   * @param target a resolve candidate.
   * @return a PsiFile if target was a PsiDirectory, or null, or target unchanged.
   */
@@ -484,19 +520,26 @@ public class PyUtil {
        file.putCopyableUserData(PyFile.KEY_IS_DIRECTORY, Boolean.TRUE);
        return file; // ResolveImportUtil will extract directory part as needed, everyone else are better off with a file.
      }
-      else return null; // dir without __init__.py does not resolve
+      else {
+        return null;
+      } // dir without __init__.py does not resolve
    }
-    else return target; // don't touch non-dirs
+    else {
+      return target;
+    } // don't touch non-dirs
  }

  /**
   * Counts initial underscores of an identifier.
+   *
   * @param name identifier
   * @return 0 if no initial underscores found, 1 if there's only one underscore, 2 if there's two or more initial underscores.
   */
  public static int getInitialUnderscores(String name) {
-    int underscores=0;
-    if (name.startsWith("__")) underscores = 2;
+    int underscores = 0;
+    if (name.startsWith("__")) {
+      underscores = 2;
+    }
    else if (name.startsWith("_")) underscores = 1;
    return underscores;
  }
@@ -504,6 +547,7 @@ public class PyUtil {
  /**
   * Tries to find nearest parent that conceals names defined inside it. Such elements are 'class' and 'def':
   * anything defined within it does not seep to the namespace below them, but is concealed within.
+   *
   * @param elt starting point of search.
   * @return 'class' or 'def' element, or null if not found.
   */
@@ -513,7 +557,7 @@ public class PyUtil {
      return null;
    }
    PsiElement parent = elt.getParent();
-    while(parent != null) {
+    while (parent != null) {
      if (parent instanceof PyClass || parent instanceof Callable) {
        return parent;
      }
@@ -533,6 +577,15 @@ public class PyUtil {
    return name.startsWith("__") && !name.endsWith("__");
  }

+  public static boolean isPythonIdentifier(String name) {
+    return PyNames.isIdentifier(name);
+  }
+
+  public static LookupElement createNamedParameterLookup(String name) {
+    LookupElementBuilder lookupElementBuilder = LookupElementBuilder.create(name + "=").setIcon(Icons.PARAMETER_ICON);
+    return PrioritizedLookupElement.withGrouping(lookupElementBuilder, 1);
+  }
+
  public static class UnderscoreFilter implements Condition<String> {
    private int myAllowed; // how many starting underscores is allowed: 0 is none, 1 is only one, 2 is two and more.

@@ -1,12 +1,13 @@
 package com.jetbrains.python.psi.impl;

-import com.intellij.codeInsight.completion.PrioritizedLookupElement;
+import com.intellij.codeInsight.lookup.LookupElement;
 import com.intellij.codeInsight.lookup.LookupElementBuilder;
 import com.intellij.lang.ASTNode;
 import com.intellij.lang.annotation.HighlightSeverity;
 import com.intellij.openapi.util.Comparing;
 import com.intellij.openapi.util.Condition;
 import com.intellij.openapi.util.TextRange;
+import com.intellij.openapi.util.text.StringUtil;
 import com.intellij.psi.*;
 import com.intellij.psi.impl.PsiManagerImpl;
 import com.intellij.psi.impl.source.resolve.ResolveCache;
@@ -14,6 +15,8 @@ import com.intellij.psi.util.PsiTreeUtil;
 import com.intellij.util.Icons;
 import com.intellij.util.IncorrectOperationException;
 import com.intellij.util.containers.SortedList;
+import com.jetbrains.appengine.util.PythonUtil;
+import com.jetbrains.appengine.util.StringUtils;
 import com.jetbrains.python.PyNames;
 import com.jetbrains.python.psi.*;
 import com.jetbrains.python.psi.resolve.*;
@@ -54,16 +57,17 @@ public class PyReferenceImpl implements PsiReferenceEx, PsiPolyVariantReference
   * Imported module names: to module file (or directory for a qualifier).
   * Other identifiers: to most recent definition before this reference.
   * This implementation is cached.
+   *
   * @see #resolveInner().
-  **/
+   */
  @Nullable
  public PsiElement resolve() {
    final ResolveResult[] results = multiResolve(false);
-   return results.length >= 1 && !(results [0] instanceof ImplicitResolveResult) ? results[0].getElement() : null;
+    return results.length >= 1 && !(results[0] instanceof ImplicitResolveResult) ? results[0].getElement() : null;
  }

  // it is *not* final so that it can be changed in debug time. if set to false, caching is off
-  private static boolean USE_CACHE = true; 
+  private static boolean USE_CACHE = true;

  /**
   * Resolves reference to possible referred elements.
@@ -72,8 +76,9 @@ public class PyReferenceImpl implements PsiReferenceEx, PsiPolyVariantReference
   * todo Local identifiers: a list of definitions in the most recent compound statement
   * (e.g. <code>if X: a = 1; else: a = 2</code> has two definitions of <code>a</code>.).
   * todo Identifiers not found locally: similar definitions in imported files and builtins.
+   *
   * @see com.intellij.psi.PsiPolyVariantReference#multiResolve(boolean)
-  **/
+   */
  @NotNull
  public ResolveResult[] multiResolve(final boolean incompleteCode) {
    final PsiManager manager = getElement().getManager();
@@ -87,6 +92,7 @@ public class PyReferenceImpl implements PsiReferenceEx, PsiPolyVariantReference
  }

  // sorts and modifies results of resolveInner
+
  private ResolveResult[] multiResolveInner(boolean incomplete) {
    final String referencedName = myElement.getReferencedName();
    if (referencedName == null) return ResolveResult.EMPTY_ARRAY;
@@ -134,6 +140,7 @@ public class PyReferenceImpl implements PsiReferenceEx, PsiPolyVariantReference

  protected static class ResultList extends ArrayList<RatedResolveResult> {
    // Allows to add non-null elements and discard nulls in a hassle-free way.
+
    public boolean poke(final PsiElement what, final int rate) {
      if (what == null) return false;
      super.add(new RatedResolveResult(rate, what));
@@ -144,6 +151,7 @@ public class PyReferenceImpl implements PsiReferenceEx, PsiPolyVariantReference

  /**
   * Does actual resolution of resolve().
+   *
   * @return resolution result.
   * @see #resolve()
   */
@@ -166,7 +174,8 @@ public class PyReferenceImpl implements PsiReferenceEx, PsiPolyVariantReference
      PsiElement one = myElement;
      do {
        one = PyUtil.getConcealingParent(one);
-      } while (one instanceof PyFunction);
+      }
+      while (one instanceof PyFunction);
      if (one instanceof PyClass) roof = one;
    }
    if (roof == null) roof = realContext.getContainingFile();
@@ -221,6 +230,7 @@ public class PyReferenceImpl implements PsiReferenceEx, PsiPolyVariantReference
  }

  // NOTE: very crude
+
  private static int getRate(PsiElement elt) {
    int rate;
    if (elt instanceof PyImportElement || elt instanceof PyStarImportElement) {
@@ -298,10 +308,10 @@ public class PyReferenceImpl implements PsiReferenceEx, PsiPolyVariantReference
        if (PsiTreeUtil.getParentOfType(myElement, PyKeywordArgument.class) == null) {
          PsiElement def = ((PyReferenceExpression)callee).getReference().resolve();
          if (def instanceof PyFunction) {
-            addKeywordArgumentVariants((PyFunction) def, ret);
+            addKeywordArgumentVariants((PyFunction)def, ret);
          }
          else if (def instanceof PyClass) {
-            PyFunction init = ((PyClass) def).findMethodByName(PyNames.INIT, true);  // search in superclasses
+            PyFunction init = ((PyClass)def).findMethodByName(PyNames.INIT, true);  // search in superclasses
            if (init != null) {
              addKeywordArgumentVariants(init, ret);
            }
@@ -353,9 +363,12 @@ public class PyReferenceImpl implements PsiReferenceEx, PsiPolyVariantReference
    }
    visited.add(def);
    final Set<PyFunction.Flag> flags = def.getContainingClass() != null ? PyUtil.detectDecorationsAndWrappersOf(def) : null;
-    final KeywordArgumentCollector collector = new KeywordArgumentCollector(flags, ret);
+    final KwArgParameterCollector collector = new KwArgParameterCollector(flags, ret);
    def.getParameterList().acceptChildren(collector);
-    if (collector.myCount == 2 && collector.myHasSelf && collector.myHasKwArgs) {
+    if (collector.hasKwArgs()) {
+      def.getStatementList().acceptChildren(new KwArgStatementCallCollector(ret, collector.getKwArgs()));
+    }
+    if (collector.hasOnlySelfAndKwArgs()) {
      // nothing interesting besides self and **kwargs, let's look at superclass (PY-778)
      final PsiElement superMethod = PySuperMethodsSearch.search(def).findFirst();
      if (superMethod instanceof PyFunction) {
@@ -364,14 +377,15 @@ public class PyReferenceImpl implements PsiReferenceEx, PsiPolyVariantReference
    }
  }

-  private static class KeywordArgumentCollector extends PyElementVisitor {
+  private static class KwArgParameterCollector extends PyElementVisitor {
    private int myCount;
    private final Set<PyFunction.Flag> myFlags;
    private final List<Object> myRet;
    private boolean myHasSelf = false;
    private boolean myHasKwArgs = false;
+    private PyParameter kwArgsParam = null;

-    public KeywordArgumentCollector(Set<PyFunction.Flag> flags, List<Object> ret) {
+    public KwArgParameterCollector(Set<PyFunction.Flag> flags, List<Object> ret) {
      myFlags = flags;
      myRet = ret;
    }
@@ -383,14 +397,74 @@ public class PyReferenceImpl implements PsiReferenceEx, PsiPolyVariantReference
        myHasSelf = true;
        return;
      }
-      PyNamedParameter n_param = par.getAsNamed();
-      assert n_param != null;
-      if (! n_param.isKeywordContainer() && ! n_param.isPositionalContainer()) {
-        final LookupElementBuilder item = LookupElementBuilder.create(n_param.getName() + "=").setIcon(n_param.getIcon(0));
-        myRet.add(PrioritizedLookupElement.withGrouping(item, 1));
+      PyNamedParameter namedParam = par.getAsNamed();
+      assert namedParam != null;
+      if (!namedParam.isKeywordContainer() && !namedParam.isPositionalContainer()) {
+        final LookupElement item = PyUtil.createNamedParameterLookup(namedParam.getName());
+        myRet.add(item);
      }
-      else if (n_param.isKeywordContainer()) {
+      else if (namedParam.isKeywordContainer()) {
        myHasKwArgs = true;
+        kwArgsParam = namedParam;
+      }
+    }
+
+    public PyParameter getKwArgs() {
+      return kwArgsParam;
+    }
+
+    public boolean hasKwArgs() {
+      return myHasKwArgs;
+    }
+
+    public boolean hasOnlySelfAndKwArgs() {
+      return myCount == 2 && myHasSelf && myHasKwArgs;
+    }
+  }
+
+  private static class KwArgStatementCallCollector extends PyElementVisitor {
+    private final List<Object> myRet;
+    private final PyParameter myKwArgs;
+
+    public KwArgStatementCallCollector(List<Object> ret, @NotNull PyParameter kwArgs) {
+      myRet = ret;
+      this.myKwArgs = kwArgs;
+    }
+
+    @Override
+    public void visitPyElement(PyElement node) {
+      node.acceptChildren(this);
+    }
+
+    @Override
+    public void visitPySubscriptionExpression(PySubscriptionExpression node) {
+      String operandName = node.getOperand().getName();
+      processGet(operandName, node.getIndexExpression());
+    }
+
+    @Override
+    public void visitPyCallExpression(PyCallExpression node) {
+      String callName = node.getCallee().getName();
+      if (callName.equals("pop") || callName.equals("get")) {
+        PyReferenceExpression child = PsiTreeUtil.getChildOfType(node.getCallee(), PyReferenceExpression.class);
+        if (child != null) {
+          String operandName = child.getName();
+          if (node.getArguments().length > 0) {
+            PyExpression argument = node.getArguments()[0];
+            processGet(operandName, argument);
+          }
+        }
+      }
+      super.visitPyCallExpression(node);
+    }
+
+    private void processGet(String operandName, PyExpression argument) {
+      if (myKwArgs.getName().equals(operandName) &&
+          argument instanceof PyStringLiteralExpression) {
+        String name = ((PyStringLiteralExpression)argument).getStringValue();
+        if (PyUtil.isPythonIdentifier(name)) {
+          myRet.add(PyUtil.createNamedParameterLookup(name));
+        }
      }
    }
  }
@@ -424,6 +498,7 @@ public class PyReferenceImpl implements PsiReferenceEx, PsiPolyVariantReference


  // our very own caching resolver
+
  private static class CachingResolver implements ResolveCache.PolyVariantResolver<PyReferenceImpl> {
    public static CachingResolver INSTANCE = new CachingResolver();
    private ThreadLocal<AtomicInteger> myNesting = new ThreadLocal<AtomicInteger>() {
@@ -0,0 +1,6 @@
+def foo(**kwargs):
+    x = kwargs.get('x_param', None)
+    run(x_param)
+
+def bar():
+    foo(x_param=)
@@ -0,0 +1,6 @@
+def foo(**kwargs):
+    x = kwargs.get('x_param', None)
+    run(x_param)
+
+def bar():
+    foo(x_<caret>)
@@ -0,0 +1,6 @@
+def foo(**kwargs):
+    x = kwargs['x_param']
+    run(x_param)
+
+def bar():
+    foo(x_param=)
@@ -0,0 +1,6 @@
+def foo(**kwargs):
+    x = kwargs['x_param']
+    run(x_param)
+
+def bar():
+    foo(x_<caret>)
@@ -83,6 +83,14 @@ public class PythonCompletionTest extends PyLightFixtureTestCase {
    doTest();
  }

+  public void testKwParamsInCodeUsage() throws Exception { //PY-1002
+    doTest();
+  }
+  
+  public void testKwParamsInCodeGetUsage() throws Exception { //PY-1002 
+    doTest();
+  }
+
  public void testImportModule() throws Exception {
    final String testName = "completion/" + getTestName(true);
    myFixture.configureByFiles(testName + ".py", "completion/someModule.py");