From 85257c075aef3f127600d54fbc48d2fb8409f748 Mon Sep 17 00:00:00 2001
From: Dmitry Jemerov <yole@jetbrains.com>
Date: Mon, 16 Jan 2012 18:21:30 +0100
Subject: [PATCH] include reads of decorators in control flow graph

---
 .../codeInsight/controlflow/PyControlFlowBuilder.java | 10 ++++++++++
 python/testData/codeInsight/controlflow/decorator.py  |  4 ++++
 python/testData/codeInsight/controlflow/decorator.txt |  5 +++++
 python/testData/codeInsight/controlflow/lambda.txt    |  7 ++++---
 .../jetbrains/python/PyControlFlowBuilderTest.java    | 11 ++++++++++-
 5 files changed, 33 insertions(+), 4 deletions(-)
 create mode 100644 python/testData/codeInsight/controlflow/decorator.py
 create mode 100644 python/testData/codeInsight/controlflow/decorator.txt

diff --git a/python/src/com/jetbrains/python/codeInsight/controlflow/PyControlFlowBuilder.java b/python/src/com/jetbrains/python/codeInsight/controlflow/PyControlFlowBuilder.java
index 40b1c7cbbe9d..ad2fe31e993c 100644
--- a/python/src/com/jetbrains/python/codeInsight/controlflow/PyControlFlowBuilder.java
+++ b/python/src/com/jetbrains/python/codeInsight/controlflow/PyControlFlowBuilder.java
@@ -39,12 +39,21 @@ public class PyControlFlowBuilder extends PyRecursiveElementVisitor {
     // Create node and stop here
     myBuilder.startNode(node);
     visitDefaultParameterValues(node.getParameterList());
+    visitDecorators(node.getDecoratorList());
 
     final ReadWriteInstruction instruction = ReadWriteInstruction.write(myBuilder, node, node.getName());
     myBuilder.addNode(instruction);
     myBuilder.checkPending(instruction);
   }
 
+  private void visitDecorators(PyDecoratorList list) {
+    if (list != null) {
+      for (PyDecorator decorator : list.getDecorators()) {
+        decorator.accept(this);
+      }
+    }
+  }
+
   private void visitDefaultParameterValues(PyParameterList parameterList) {
     ParamHelper.walkDownParamArray(parameterList.getParameters(), new ParamHelper.ParamVisitor() {
       @Override
@@ -64,6 +73,7 @@ public class PyControlFlowBuilder extends PyRecursiveElementVisitor {
     for (PsiElement element : node.getSuperClassExpressions()) {
       element.accept(this);
     }
+    visitDecorators(node.getDecoratorList());
     final ReadWriteInstruction instruction = ReadWriteInstruction.write(myBuilder, node, node.getName());
     myBuilder.addNode(instruction);
     myBuilder.checkPending(instruction);
diff --git a/python/testData/codeInsight/controlflow/decorator.py b/python/testData/codeInsight/controlflow/decorator.py
new file mode 100644
index 000000000000..eca16d28d0ca
--- /dev/null
+++ b/python/testData/codeInsight/controlflow/decorator.py
@@ -0,0 +1,4 @@
+class C:
+    @classmethod
+    def foo():
+        pass
diff --git a/python/testData/codeInsight/controlflow/decorator.txt b/python/testData/codeInsight/controlflow/decorator.txt
new file mode 100644
index 000000000000..55aa8e7e2997
--- /dev/null
+++ b/python/testData/codeInsight/controlflow/decorator.txt
@@ -0,0 +1,5 @@
+0(1) element: null
+1(2) element: PyFunction('foo')
+2(3) READ ACCESS: classmethod
+3(4) WRITE ACCESS: foo
+4() element: null
\ No newline at end of file
diff --git a/python/testData/codeInsight/controlflow/lambda.txt b/python/testData/codeInsight/controlflow/lambda.txt
index 0a48427a36b7..0a316efc37ad 100644
--- a/python/testData/codeInsight/controlflow/lambda.txt
+++ b/python/testData/codeInsight/controlflow/lambda.txt
@@ -2,9 +2,10 @@
 1(2) element: PyExpressionStatement
 2(3) element: PySetCompExpression
 3(4) element: PyCallExpression: range
-4(5,9) READ ACCESS: range
+4(5,10) READ ACCESS: range
 5(6) element: PyTargetExpression: i
 6(7) WRITE ACCESS: i
 7(8) element: PyParenthesizedExpression
-8(5,9) element: PyLambdaExpression
-9() element: null
\ No newline at end of file
+8(9) element: PyLambdaExpression
+9(5,10) READ ACCESS: i
+10() element: null
\ No newline at end of file
diff --git a/python/testSrc/com/jetbrains/python/PyControlFlowBuilderTest.java b/python/testSrc/com/jetbrains/python/PyControlFlowBuilderTest.java
index d004262879f5..23ed486402be 100644
--- a/python/testSrc/com/jetbrains/python/PyControlFlowBuilderTest.java
+++ b/python/testSrc/com/jetbrains/python/PyControlFlowBuilderTest.java
@@ -7,6 +7,7 @@ import com.intellij.openapi.vfs.LocalFileSystem;
 import com.intellij.openapi.vfs.VfsUtil;
 import com.intellij.openapi.vfs.VirtualFile;
 import com.jetbrains.python.codeInsight.controlflow.ControlFlowCache;
+import com.jetbrains.python.codeInsight.controlflow.ScopeOwner;
 import com.jetbrains.python.fixtures.LightMarkedTestCase;
 import com.jetbrains.python.psi.PyClass;
 import com.jetbrains.python.psi.PyFile;
@@ -164,6 +165,10 @@ public class PyControlFlowBuilderTest extends LightMarkedTestCase {
   public void testLambdaDefaultParameter() {
     doTest();
   }
+  
+  public void testDecorator() {
+    doTestFirstStatement();
+  }
 
   public void testQualifiedSelfReference() {
     final String testName = getTestName(false).toLowerCase();
@@ -192,10 +197,14 @@ public class PyControlFlowBuilderTest extends LightMarkedTestCase {
   }
 
   public void testFunction() {
+    doTestFirstStatement();
+  }
+
+  private void doTestFirstStatement() {
     final String testName = getTestName(false).toLowerCase();
     configureByFile(testName + ".py");
     final String fullPath = getTestDataPath() + testName + ".txt";
-    final ControlFlow flow = ControlFlowCache.getControlFlow((PyFunction)((PyFile)myFile).getStatements().get(0));
+    final ControlFlow flow = ControlFlowCache.getControlFlow((ScopeOwner)((PyFile)myFile).getStatements().get(0));
     check(fullPath, flow);
   }