Enhanced callee resolution for decorators. Tests for decorators.

python/src/com/jetbrains/python/psi/PyDecorator.java

@@ -7,10 +7,14 @@ import org.jetbrains.annotations.Nullable;
 /**
  * Covers a decorator call, e.g. <tt>@staticmethod</tt>.
  * Decorators happen contextually above the function definition, but are stored inside it for convenience.
+ * <b>Note:</b>
+ * In <code>@foo</code> form, <code>PyCallExpression</code>'s methods are related to invocation of <code>foo</code>
+ * as decorator. In <code>@foo(...)</code> form, these very methods are related to the call that returns the decorator
+ * to be applied. In either case, they are related to an invocation of <code>foo</code>.
  * User: dcheryasov
  * Date: Sep 26, 2008
  */
-public interface PyDecorator extends /*PyElement*/ PyCallExpression, StubBasedPsiElement<PyDecoratorStub> {
+public interface PyDecorator extends PyCallExpression, StubBasedPsiElement<PyDecoratorStub> {
   /**
    * @return the function being decorated, or null.
    */
@@ -18,9 +22,14 @@ public interface PyDecorator extends /*PyElement*/ PyCallExpression, StubBasedPs
   PyFunction getTarget();
 
   /**
-   * True if the annotating function is a builtin. Uses a stub, does not incur parsing, useful togeter with getName().
+   * True if the annotating function is a builtin, useful togeter with getName(). Implementation uses stub info.
    * @see com.jetbrains.python.psi.PyElement#getName()
    */
   boolean isBuiltin();
 
+  /**
+   * @return true if invocation has a form of <code>@foo(...)</code>.
+   */
+  boolean hasArgumentList();
+
 }

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

@@ -71,6 +71,10 @@ public class PyDecoratorImpl extends PyPresentableElementImpl<PyDecoratorStub> i
     }
   }
 
+  public boolean hasArgumentList() {
+    return getNode().findChildByType(PyElementTypes.ARGUMENT_LIST) != null;
+  }
+
   public PyExpression getCallee() {
     try {
       return (PyExpression)getFirstChild().getNextSibling(); // skip the @ before call
@@ -93,7 +97,12 @@ public class PyDecoratorImpl extends PyPresentableElementImpl<PyDecoratorStub> i
   }
 
   public PyMarkedFunction resolveCallee() {
-    return PyCallExpressionHelper.resolveCallee(this);
+    PyMarkedFunction callee = PyCallExpressionHelper.resolveCallee(this);
+    if (callee == null) return null;
+    if (! hasArgumentList()) {
+      callee.getFlags().add(Flag.IMPLICIT_FIRST_ARG); // NOTE: assumes mutability
+    }
+    return callee;
   }
 
   @Override

python/testData/psi/DecoratedFunction.py

@@ -0,0 +1,5 @@
+@staticmethod
+@xmlize(node="foo")
+def foo():
+  pass
+

python/testData/psi/DecoratedFunction.txt

@@ -0,0 +1,32 @@
+PyFile:DecoratedFunction.py
+  PyFunction
+    PyDecoratorList
+      PyDecorator: @staticmethod
+        PsiElement(Py:AT)('@')
+        PyReferenceExpression: staticmethod
+          PsiElement(Py:IDENTIFIER)('staticmethod')
+      PsiWhiteSpace('\n')
+      PyDecorator: @xmlize
+        PsiElement(Py:AT)('@')
+        PyReferenceExpression: xmlize
+          PsiElement(Py:IDENTIFIER)('xmlize')
+        PyArgumentList
+          PsiElement(Py:LPAR)('(')
+          PyKeywordArgumentImpl: node
+            PsiElement(Py:IDENTIFIER)('node')
+            PsiElement(Py:EQ)('=')
+            PyStringLiteralExpression: foo
+              PsiElement(Py:STRING_LITERAL)('"foo"')
+          PsiElement(Py:RPAR)(')')
+    PsiWhiteSpace('\n')
+    PsiElement(Py:DEF_KEYWORD)('def')
+    PsiWhiteSpace(' ')
+    PsiElement(Py:IDENTIFIER)('foo')
+    PyParameterList
+      PsiElement(Py:LPAR)('(')
+      PsiElement(Py:RPAR)(')')
+    PsiElement(Py:COLON)(':')
+    PsiWhiteSpace('\n  ')
+    PyStatementList
+      PyPassStatement
+        PsiElement(Py:PASS_KEYWORD)('pass')

python/testData/resolve/callee/DecoCall.py

@@ -0,0 +1,6 @@
+def deco(fun):
+  return fun
+
+@<ref>deco
+def foo():
+  pass

python/testData/resolve/callee/DecoParamCall.py

@@ -0,0 +1,11 @@
+def deco(prefix):
+  def fun(f):
+    print f
+    def dfun():
+      return [prefix, f()]
+    return dfun
+  return fun
+
+@<ref>deco(1)
+def foo():
+  pass

python/testData/resolve/decorators/DecoCall.py

@@ -0,0 +1,8 @@
+def deco(fun):
+  return fun
+
+@<ref>deco
+def foo():
+  pass
+
+# same as in callee test

python/testData/resolve/decorators/DecoParamCall.py

@@ -0,0 +1,13 @@
+def deco(prefix):
+  def fun(f):
+    print f
+    def dfun():
+      return [prefix, f()]
+    return dfun
+  return fun
+
+@<ref>deco(1)
+def foo():
+  pass
+
+# same as in callee test

python/testData/stubs/StubStructure.py

@@ -1,9 +1,12 @@
+def deco(fun):
+  return fun # valid
+
 class FooClass:
   staticField = 1
   def __init__(self):
     self.instanceField = 2
 
-  @staticmethod
+  @deco
   def fooFunction(fooParam1, fooParam2) :
     pass
 

python/testSrc/com/jetbrains/python/PyDecoratorTest.java

@@ -0,0 +1,54 @@
+package com.jetbrains.python;
+
+import com.intellij.openapi.application.PathManager;
+import com.intellij.psi.PsiElement;
+import com.intellij.psi.PsiReference;
+import com.intellij.testFramework.ResolveTestCase;
+import com.jetbrains.python.psi.PyArgumentList;
+import com.jetbrains.python.psi.PyDecorator;
+import com.jetbrains.python.psi.PyFunction;
+import com.jetbrains.python.psi.PyExpression;
+
+/**
+ * Decorator-specific tests.
+ * User: dcheryasov
+ * Date: Dec 28, 2008 3:50:23 AM
+ */
+public class PyDecoratorTest extends ResolveTestCase {
+  private PsiElement find() throws Exception {
+    PsiReference ref = configureByFile(getTestName(false) + ".py");
+    return ref.getElement();
+  }
+
+  public void testDecoCall() throws Exception {
+    PsiElement targetElement = find().getParent();
+    assertTrue(targetElement instanceof PyDecorator);
+    PyDecorator deco = (PyDecorator)targetElement;
+    PyFunction decofun = deco.getTarget();
+    assertNotNull(decofun);
+    assertEquals("foo", decofun.getName());
+    assertFalse(deco.isBuiltin());
+    assertFalse(deco.hasArgumentList());
+  }
+
+  public void testDecoParamCall() throws Exception {
+    PsiElement targetElement = find().getParent();
+    assertTrue(targetElement instanceof PyDecorator);
+    PyDecorator deco = (PyDecorator)targetElement;
+    PyFunction decofun = deco.getTarget();
+    assertNotNull(decofun);
+    assertEquals("foo", decofun.getName());
+    assertFalse(deco.isBuiltin());
+    assertTrue(deco.hasArgumentList());
+    PyArgumentList arglist = deco.getArgumentList();
+    assertNotNull(arglist);
+    PyExpression[] args = arglist.getArguments();
+    assertEquals("argument count", 1, args.length);
+    assertEquals("argument value", "1", args[0].getText());
+  }
+
+  @Override
+  protected String getTestDataPath() {
+    return PathManager.getHomePath() + "/plugins/python/testData/resolve/decorators/";
+  }
+}

python/testSrc/com/jetbrains/python/PyResolveCalleeTest.java

@@ -33,6 +33,18 @@ public class PyResolveCalleeTest extends ResolveTestCase {
     assertTrue(resolved.getFlags().equals(EnumSet.noneOf(PyCallExpression.Flag.class)));
   }
 
+  public void testDecoCall() throws Exception {
+    PyCallExpression.PyMarkedFunction resolved = resolveCallee();
+    assertNotNull(resolved.getFunction());
+    assertTrue(resolved.getFlags().equals(EnumSet.of(PyCallExpression.Flag.IMPLICIT_FIRST_ARG)));
+  }
+
+  public void testDecoParamCall() throws Exception {
+    PyCallExpression.PyMarkedFunction resolved = resolveCallee();
+    assertNotNull(resolved.getFunction());
+    assertTrue(resolved.getFlags().equals(EnumSet.noneOf(PyCallExpression.Flag.class)));
+  }
+
   @Override
   protected String getTestDataPath() {
     return PathManager.getHomePath() + "/plugins/python/testData/resolve/callee/";

python/testSrc/com/jetbrains/python/PyStubsTest.java

@@ -75,23 +75,20 @@ public class PyStubsTest extends CodeInsightTestCase {
     assertNotParsed(file);
     PyDecorator deco = da[0];
     assertNotNull(deco);
-    /*
-    assertEquals(deco.getTarget(), decorated);
+    assertEquals("deco", deco.getName());
     assertNotParsed(file);
-    assertEquals(deco.getName(), "staticmethod");
-    */
-    assertEquals(deco.getName(), "staticmethod");
-    assertNotParsed(file);
-    //assertNotNull(deco.getCallee());
+    assertFalse(deco.isBuiltin());
 
     final PyTargetExpression[] instanceAttrs = pyClass.getInstanceAttributes();
     assertEquals(1, instanceAttrs.length);
     assertEquals("instanceField", instanceAttrs [0].getName());
 
     final List<PyFunction> functions = file.getTopLevelFunctions();
-    assertEquals(1, functions.size());
+    assertEquals(2, functions.size()); // "deco" and "topLevelFunction"
     PyFunction func = functions.get(0);
+    assertEquals("deco", func.getName());
 
+    func = functions.get(1);
     assertEquals("topLevelFunction", func.getName());
 
     final List<PyTargetExpression> exprs = file.getTopLevelAttributes();

python/testSrc/com/jetbrains/python/PythonParsingTest.java

@@ -77,4 +77,7 @@ public class PythonParsingTest extends ParsingTestCase {
     doTest(true);
   }
 
+  public void testDecoratedFunction() throws Exception {
+    doTest(true);
+  }
 }