handle qualified refs in control flow builder

python/src/META-INF/python-plugin-common.xml

@@ -269,6 +269,10 @@
       <add-to-group group-id="ToolsMenu" anchor="after" relative-to-action="Macros"/>
     </action>
 
+    <action id="Devmode.AnalyzeReturns" class="com.jetbrains.python.devmode.AnalyzeReturnsAction" text="Analyze return types" internal="true">
+      <add-to-group group-id="ProjectViewPopupMenu" anchor="last"/>
+    </action>
+
     <group id="PyTypeHierarchyPopupMenu">
       <reference ref="TypeHierarchyBase.BaseOnThisType"/>
       <reference ref="TypeHierarchy.Class"/>

python/src/com/jetbrains/python/codeInsight/controlflow/PyControlFlowBuilder.java

@@ -1,7 +1,7 @@
 package com.jetbrains.python.codeInsight.controlflow;
 
-import com.intellij.codeInsight.controlflow.ControlFlowBuilder;
 import com.intellij.codeInsight.controlflow.ControlFlow;
+import com.intellij.codeInsight.controlflow.ControlFlowBuilder;
 import com.intellij.codeInsight.controlflow.Instruction;
 import com.intellij.codeInsight.controlflow.impl.InstructionImpl;
 import com.intellij.openapi.util.Pair;
@@ -9,7 +9,6 @@ import com.intellij.openapi.util.Ref;
 import com.intellij.psi.PsiElement;
 import com.intellij.psi.util.PsiTreeUtil;
 import com.jetbrains.python.psi.*;
-import com.jetbrains.python.psi.impl.PyAssignmentStatementNavigator;
 import com.jetbrains.python.psi.impl.PyAugAssignmentStatementNavigator;
 import com.jetbrains.python.psi.impl.PyImportStatementNavigator;
 import org.jetbrains.annotations.NotNull;
@@ -87,7 +86,9 @@ public class PyControlFlowBuilder extends PyRecursiveElementVisitor {
     final PyExpression qualifier = node.getQualifier();
     if (qualifier != null){
       qualifier.accept(this);
-      return;
+      if (!isSelf(qualifier)) {
+        return;
+      }
     }
     if (PyImportStatementNavigator.getImportStatementByElement(node) != null){
       return;
@@ -101,6 +102,19 @@ public class PyControlFlowBuilder extends PyRecursiveElementVisitor {
     myBuilder.checkPending(readWriteInstruction);
   }
 
+  private static boolean isSelf(PsiElement qualifier) {
+    PyFunction func = PsiTreeUtil.getParentOfType(qualifier, PyFunction.class);
+    if (func == null || PsiTreeUtil.getParentOfType(func, PyClass.class) == null) {
+      return false;
+    }
+    final PyParameter[] params = func.getParameterList().getParameters();
+    if (params.length == 0) {
+      return false;
+    }
+    final PyNamedParameter named = params[0].getAsNamed();
+    return named != null && named.getName().equals(qualifier.getText());
+  }
+
   @Override
   public void visitPyAssignmentStatement(final PyAssignmentStatement node) {
     myBuilder.startNode(node);
@@ -127,13 +141,14 @@ public class PyControlFlowBuilder extends PyRecursiveElementVisitor {
   public void visitPyTargetExpression(final PyTargetExpression node) {
     final PsiElement[] children = node.getChildren();
     // Case of non qualified reference
-    if (children.length == 0){
+    if (children.length == 0 || (children.length == 1 && isSelf(children [0]))){
       final ReadWriteInstruction.ACCESS access = node.getParent() instanceof PySliceExpression 
                                                  ? ReadWriteInstruction.ACCESS.READ : ReadWriteInstruction.ACCESS.WRITE;
       final ReadWriteInstruction instruction = new ReadWriteInstruction(myBuilder, node, node.getName(), access);
       myBuilder.addNode(instruction);
       myBuilder.checkPending(instruction);
-    } else {
+    }
+    else {
       for (PsiElement child : children) {
         child.accept(this);
       }

python/testData/codeInsight/controlflow/self.py

@@ -0,0 +1,5 @@
+class A:
+    def __init__(self):
+        self.foo = 1
+        print self.foo
+        

python/testData/codeInsight/controlflow/self.txt

@@ -0,0 +1,8 @@
+0(1) element: null
+1(2) WRITE ACCESS: self
+2(3) element: PyAssignmentStatement
+3(4) WRITE ACCESS: foo
+4(5) element: PyPrintStatement
+5(6) READ ACCESS: self
+6(7) READ ACCESS: foo
+7() element: null

python/testData/codeInsight/controlflow/while.py

@@ -1,4 +1,4 @@
-while condition
+while condition:
     print("Alloha!")
 else:
     print("Bye")

python/testSrc/com/jetbrains/python/PyControlFlowBuilderTest.java

@@ -1,14 +1,15 @@
 package com.jetbrains.python;
 
+import com.intellij.codeInsight.controlflow.ControlFlow;
 import com.intellij.codeInsight.controlflow.Instruction;
 import com.intellij.openapi.util.text.StringUtil;
 import com.intellij.openapi.vfs.LocalFileSystem;
 import com.intellij.openapi.vfs.VfsUtil;
 import com.intellij.openapi.vfs.VirtualFile;
 import com.jetbrains.python.fixtures.LightMarkedTestCase;
+import com.jetbrains.python.psi.PyClass;
 import com.jetbrains.python.psi.PyFile;
 import com.jetbrains.python.psi.PyFunction;
-import com.intellij.codeInsight.controlflow.ControlFlow;
 
 import java.io.File;
 import java.io.IOException;
@@ -110,6 +111,15 @@ public class PyControlFlowBuilderTest extends LightMarkedTestCase {
     doTest();
   }
 
+  public void testSelf() throws Exception {
+    final String testName = getTestName(false).toLowerCase();
+    configureByFile(testName + ".py");
+    final String fullPath = getTestDataPath() + testName + ".txt";
+    final PyClass pyClass = ((PyFile) myFile).getTopLevelClasses().get(0);
+    final ControlFlow flow = pyClass.getMethods() [0].getControlFlow();
+    check(fullPath, flow);
+  }
+
   public void testTryBreak() throws Exception {
     final String testName = getTestName(false).toLowerCase();
     configureByFile(testName + ".py");

python/testSrc/com/jetbrains/python/PythonAllTestsSuite.java

@@ -1,5 +1,6 @@
 package com.jetbrains.python;
 
+import com.jetbrains.python.refactoring.PyCodeFragmentTest;
 import com.jetbrains.python.refactoring.PyExtractMethodTest;
 import com.jetbrains.python.refactoring.PyInlineLocalTest;
 import com.jetbrains.python.refactoring.PyRenameTest;
@@ -50,7 +51,9 @@ public class PythonAllTestsSuite {
     PyInlineLocalTest.class,
     PyAutoUnindentTest.class,
     PyFindUsagesTest.class,
-    PyTypeTest.class
+    PyTypeTest.class,
+    PyControlFlowBuilderTest.class,
+    PyCodeFragmentTest.class
   };
 
   public static TestSuite suite() {