From ee7d00561290f0739c137c0e955488fb6fcdc74a Mon Sep 17 00:00:00 2001
From: Andrey Vlasovskikh <andrey.vlasovskikh@jetbrains.com>
Date: Wed, 27 Jul 2011 22:37:16 +0400
Subject: [PATCH] Type inference for with statements according to context
 management procotol (PY-4198)

---
 python/src/com/jetbrains/python/PyNames.java  |  1 +
 .../com/jetbrains/python/psi/PyWithItem.java  |  5 ++++-
 .../psi/impl/PyTargetExpressionImpl.java      | 21 ++++++++++++++++---
 .../python/psi/impl/PyWithItemImpl.java       | 20 ++++++++++++++++--
 .../python/psi/impl/PyWithStatementImpl.java  |  2 +-
 python/testData/completion/withType.after.py  | 13 ++++++++++++
 python/testData/completion/withType.py        | 13 ++++++++++++
 .../python/PythonCompletionTest.java          |  7 ++++++-
 8 files changed, 74 insertions(+), 8 deletions(-)
 create mode 100644 python/testData/completion/withType.after.py
 create mode 100644 python/testData/completion/withType.py

diff --git a/python/src/com/jetbrains/python/PyNames.java b/python/src/com/jetbrains/python/PyNames.java
index 1d240aa0a7ef..d13e6f90bc46 100644
--- a/python/src/com/jetbrains/python/PyNames.java
+++ b/python/src/com/jetbrains/python/PyNames.java
@@ -78,6 +78,7 @@ public class PyNames {
   public static final String GETITEM = "__getitem__";
 
   public static final String NAME = "__name__";
+  public static final String ENTER = "__enter__";
 
   /**
    * Contains all known predefined names of "__foo__" form.
diff --git a/python/src/com/jetbrains/python/psi/PyWithItem.java b/python/src/com/jetbrains/python/psi/PyWithItem.java
index 2b6998943f0e..c9586fe95c8f 100644
--- a/python/src/com/jetbrains/python/psi/PyWithItem.java
+++ b/python/src/com/jetbrains/python/psi/PyWithItem.java
@@ -9,5 +9,8 @@ public interface PyWithItem extends PyElement {
   PyWithItem[] EMPTY_ARRAY = new PyWithItem[0];
 
   @Nullable
-  PyExpression getTargetExpression();
+  PyExpression getExpression();
+
+  @Nullable
+  PyExpression getTarget();
 }
diff --git a/python/src/com/jetbrains/python/psi/impl/PyTargetExpressionImpl.java b/python/src/com/jetbrains/python/psi/impl/PyTargetExpressionImpl.java
index f7cbf45107b7..ce3d0df03294 100644
--- a/python/src/com/jetbrains/python/psi/impl/PyTargetExpressionImpl.java
+++ b/python/src/com/jetbrains/python/psi/impl/PyTargetExpressionImpl.java
@@ -136,8 +136,9 @@ public class PyTargetExpressionImpl extends PyPresentableElementImpl<PyTargetExp
       if (type != null) {
         return type;
       }
-      if (getParent() instanceof PyAssignmentStatement) {
-        final PyAssignmentStatement assignmentStatement = (PyAssignmentStatement)getParent();
+      final PsiElement parent = getParent();
+      if (parent instanceof PyAssignmentStatement) {
+        final PyAssignmentStatement assignmentStatement = (PyAssignmentStatement)parent;
         final PyExpression assignedValue = assignmentStatement.getAssignedValue();
         if (assignedValue != null) {
           if (assignedValue instanceof PyReferenceExpressionImpl) {
@@ -166,12 +167,26 @@ public class PyTargetExpressionImpl extends PyPresentableElementImpl<PyTargetExp
           return context.getType(assignedValue);
         }
       }
-      if (getParent() instanceof PyTupleExpression) {
+      if (parent instanceof PyTupleExpression) {
         final PyType typeFromTupleAssignment = getTypeFromTupleAssignment(context);
         if (typeFromTupleAssignment != null) {
           return typeFromTupleAssignment;
         }
       }
+      if (parent instanceof PyWithItem) {
+        final PyWithItem item = (PyWithItem)parent;
+        final PyType exprType = item.getExpression().getType(context);
+        if (exprType instanceof PyClassType) {
+          final PyClass cls = ((PyClassType)exprType).getPyClass();
+          if (cls != null) {
+            final PyFunction enter = cls.findMethodByName(PyNames.ENTER, true);
+            if (enter != null) {
+              return enter.getReturnType(context, null);
+            }
+          }
+        }
+        return null;
+      }
       PyType iterType = getTypeFromIteration(context);
       if (iterType != null) {
         return iterType;
diff --git a/python/src/com/jetbrains/python/psi/impl/PyWithItemImpl.java b/python/src/com/jetbrains/python/psi/impl/PyWithItemImpl.java
index 8e1e187e2834..57a2fbac5ac5 100644
--- a/python/src/com/jetbrains/python/psi/impl/PyWithItemImpl.java
+++ b/python/src/com/jetbrains/python/psi/impl/PyWithItemImpl.java
@@ -1,6 +1,7 @@
 package com.jetbrains.python.psi.impl;
 
 import com.intellij.lang.ASTNode;
+import com.intellij.psi.PsiElement;
 import com.jetbrains.python.PyTokenTypes;
 import com.jetbrains.python.psi.PyExpression;
 import com.jetbrains.python.psi.PyWithItem;
@@ -14,8 +15,23 @@ public class PyWithItemImpl extends PyElementImpl implements PyWithItem {
     super(astNode);
   }
 
-  @Nullable
-  public PyExpression getTargetExpression() {
+  @Override
+  public PyExpression getExpression() {
+    ASTNode[] children = getNode().getChildren(null);
+    for (ASTNode child: children) {
+      final PsiElement e = child.getPsi();
+      if (e instanceof PyExpression) {
+        return (PyExpression)e;
+      }
+      else if (child.getElementType() == PyTokenTypes.AS_KEYWORD) {
+        break;
+      }
+    }
+    return null;
+  }
+
+  @Override
+  public PyExpression getTarget() {
     ASTNode[] children = getNode().getChildren(null);
     boolean foundAs = false;
     for (ASTNode child : children) {
diff --git a/python/src/com/jetbrains/python/psi/impl/PyWithStatementImpl.java b/python/src/com/jetbrains/python/psi/impl/PyWithStatementImpl.java
index 89e61743b613..cada8da1b2f6 100644
--- a/python/src/com/jetbrains/python/psi/impl/PyWithStatementImpl.java
+++ b/python/src/com/jetbrains/python/psi/impl/PyWithStatementImpl.java
@@ -31,7 +31,7 @@ public class PyWithStatementImpl extends PyElementImpl implements PyWithStatemen
     List<PyElement> result = new ArrayList<PyElement>();
     if (items != null) {
       for (PyWithItem item : items) {
-        PyExpression targetExpression = item.getTargetExpression();
+        PyExpression targetExpression = item.getTarget();
         result.addAll(PyUtil.flattenedParensAndTuples(targetExpression));
       }
     }
diff --git a/python/testData/completion/withType.after.py b/python/testData/completion/withType.after.py
new file mode 100644
index 000000000000..9c8bab5a705c
--- /dev/null
+++ b/python/testData/completion/withType.after.py
@@ -0,0 +1,13 @@
+class Eggs(object):
+    def __enter__(self):
+        return u'foo'
+
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        pass
+
+class Spam(Eggs):
+    pass
+
+def f():
+    with Spam() as spam:
+        spam.encode()
diff --git a/python/testData/completion/withType.py b/python/testData/completion/withType.py
new file mode 100644
index 000000000000..0734fea58af9
--- /dev/null
+++ b/python/testData/completion/withType.py
@@ -0,0 +1,13 @@
+class Eggs(object):
+    def __enter__(self):
+        return u'foo'
+
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        pass
+
+class Spam(Eggs):
+    pass
+
+def f():
+    with Spam() as spam:
+        spam.enc<caret>
diff --git a/python/testSrc/com/jetbrains/python/PythonCompletionTest.java b/python/testSrc/com/jetbrains/python/PythonCompletionTest.java
index 5d295e23d52e..cb324f96d7da 100644
--- a/python/testSrc/com/jetbrains/python/PythonCompletionTest.java
+++ b/python/testSrc/com/jetbrains/python/PythonCompletionTest.java
@@ -177,7 +177,12 @@ public class PythonCompletionTest extends PyLightFixtureTestCase {
     doTest();
   }
 
-  public void testReturnType() {  
+  public void testReturnType() {
+    doTest();
+  }
+
+  public void testWithType() { // PY-4198
+    setLanguageLevel(LanguageLevel.PYTHON26);
     doTest();
   }