From 411951bad3d401802296444254957744c10b262c Mon Sep 17 00:00:00 2001
From: Dmitry Jemerov <yole@jetbrains.com>
Date: Tue, 14 Aug 2012 19:05:09 +0200
Subject: [PATCH] hook for providing type of 'with' block variables; provide
 type for contextlib.closing

---
 .../python/psi/impl/PyTypeProvider.java       |  3 ++
 .../python/psi/types/PyTypeProviderBase.java  |  6 +++
 .../stdlib/PyStdlibTypeProvider.java          | 12 +++++
 .../psi/impl/PyTargetExpressionImpl.java      | 48 +++++++++++--------
 4 files changed, 49 insertions(+), 20 deletions(-)

diff --git a/python/psi-api/src/com/jetbrains/python/psi/impl/PyTypeProvider.java b/python/psi-api/src/com/jetbrains/python/psi/impl/PyTypeProvider.java
index 722130b6f197..de6116ae3e20 100644
--- a/python/psi-api/src/com/jetbrains/python/psi/impl/PyTypeProvider.java
+++ b/python/psi-api/src/com/jetbrains/python/psi/impl/PyTypeProvider.java
@@ -28,4 +28,7 @@ public interface PyTypeProvider {
 
   @Nullable
   PyType getIterationType(PyClass iterable);
+
+  @Nullable
+  PyType getContextManagerVariableType(PyClass contextManager, PyExpression withExpression, TypeEvalContext context);
 }
diff --git a/python/psi-api/src/com/jetbrains/python/psi/types/PyTypeProviderBase.java b/python/psi-api/src/com/jetbrains/python/psi/types/PyTypeProviderBase.java
index f568f3cf2f7b..04a2ff40c895 100644
--- a/python/psi-api/src/com/jetbrains/python/psi/types/PyTypeProviderBase.java
+++ b/python/psi-api/src/com/jetbrains/python/psi/types/PyTypeProviderBase.java
@@ -96,6 +96,12 @@ public class PyTypeProviderBase implements PyTypeProvider {
     return null;
   }
 
+  @Nullable
+  @Override
+  public PyType getContextManagerVariableType(PyClass contextManager, PyExpression withExpression, TypeEvalContext context) {
+    return null;
+  }
+
   protected void registerSelfReturnType(String classQualifiedName, Collection<String> methods) {
     registerReturnType(classQualifiedName, methods, mySelfTypeCallback);
   }
diff --git a/python/src/com/jetbrains/python/codeInsight/stdlib/PyStdlibTypeProvider.java b/python/src/com/jetbrains/python/codeInsight/stdlib/PyStdlibTypeProvider.java
index 60659cd5f306..015accf5176c 100644
--- a/python/src/com/jetbrains/python/codeInsight/stdlib/PyStdlibTypeProvider.java
+++ b/python/src/com/jetbrains/python/codeInsight/stdlib/PyStdlibTypeProvider.java
@@ -124,6 +124,18 @@ public class PyStdlibTypeProvider extends PyTypeProviderBase {
     return null;
   }
 
+  @Nullable
+  @Override
+  public PyType getContextManagerVariableType(PyClass contextManager, PyExpression withExpression, TypeEvalContext context) {
+    if ("contextlib.closing".equals(contextManager.getQualifiedName()) && withExpression instanceof PyCallExpression) {
+      PyExpression closee = ((PyCallExpression)withExpression).getArgument(0, PyExpression.class);
+      if (closee != null) {
+        return closee.getType(context);
+      }
+    }
+    return null;
+  }
+
   @Nullable
   private PyType getOverloadedReturnTypeByQName(@NotNull Map<PyExpression, PyNamedParameter> arguments,
                                                 @NotNull String qname,
diff --git a/python/src/com/jetbrains/python/psi/impl/PyTargetExpressionImpl.java b/python/src/com/jetbrains/python/psi/impl/PyTargetExpressionImpl.java
index c1f647372133..cb39aea1dd47 100644
--- a/python/src/com/jetbrains/python/psi/impl/PyTargetExpressionImpl.java
+++ b/python/src/com/jetbrains/python/psi/impl/PyTargetExpressionImpl.java
@@ -172,26 +172,7 @@ public class PyTargetExpressionImpl extends PyPresentableElementImpl<PyTargetExp
         }
       }
       if (parent instanceof PyWithItem) {
-        final PyWithItem item = (PyWithItem)parent;
-        final PyExpression expression = item.getExpression();
-        if (expression != null) {
-          final PyType exprType = expression.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) {
-                final PyType enterType = enter.getReturnType(context, null);
-                if (enterType != null) {
-                  return enterType;
-                }
-                // Guess the return type of __enter__
-                return PyUnionType.createWeakType(exprType);
-              }
-            }
-          }
-        }
-        return null;
+        return getWithItemVariableType(context, (PyWithItem)parent);
       }
       PyType iterType = getTypeFromIteration(context);
       if (iterType != null) {
@@ -208,6 +189,33 @@ public class PyTargetExpressionImpl extends PyPresentableElementImpl<PyTargetExp
     }
   }
 
+  @Nullable
+  private static PyType getWithItemVariableType(TypeEvalContext context, PyWithItem item) {
+    final PyExpression expression = item.getExpression();
+    if (expression != null) {
+      final PyType exprType = expression.getType(context);
+      if (exprType instanceof PyClassType) {
+        final PyClass cls = ((PyClassType)exprType).getPyClass();
+        final PyFunction enter = cls.findMethodByName(PyNames.ENTER, true);
+        if (enter != null) {
+          final PyType enterType = enter.getReturnType(context, null);
+          if (enterType != null) {
+            return enterType;
+          }
+          for (PyTypeProvider provider: Extensions.getExtensions(PyTypeProvider.EP_NAME)) {
+            PyType typeFromProvider = provider.getContextManagerVariableType(cls, expression, context);
+            if (typeFromProvider != null) {
+              return typeFromProvider;
+            }
+          }
+          // Guess the return type of __enter__
+          return PyUnionType.createWeakType(exprType);
+        }
+      }
+    }
+    return null;
+  }
+
   @Nullable
   public static PyType getTypeFromDocString(PyTargetExpressionImpl targetExpression) {
     final String docString = PyPsiUtils.strValue(PyUtil.getAttributeDocString(targetExpression));