From ab034cf35a5eb0b776248e9e21ba0518ca9e4867 Mon Sep 17 00:00:00 2001
From: Andrey Vlasovskikh <andrey.vlasovskikh@jetbrains.com>
Date: Mon, 5 Aug 2013 16:04:27 +0400
Subject: [PATCH] Override function signatures in user skeletons

---
 .../com/jetbrains/python/psi/impl/PyTypeProvider.java  |  3 +++
 .../jetbrains/python/psi/types/PyTypeProviderBase.java |  6 ++++++
 .../userSkeletons/PyUserSkeletonsTypeProvider.java     | 10 ++++++++++
 .../com/jetbrains/python/psi/impl/PyFunctionImpl.java  |  6 ++++++
 .../python/psi/impl/PyLambdaExpressionImpl.java        |  7 +++++++
 5 files changed, 32 insertions(+)

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 38dfcac02dae..cc3452d52dc6 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
@@ -31,4 +31,7 @@ public interface PyTypeProvider {
 
   @Nullable
   PyType getContextManagerVariableType(PyClass contextManager, PyExpression withExpression, TypeEvalContext context);
+
+  @Nullable
+  PyType getCallableType(@NotNull Callable callable, @NotNull 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 81988147bd01..aa23c396bab6 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
@@ -102,6 +102,12 @@ public class PyTypeProviderBase implements PyTypeProvider {
     return null;
   }
 
+  @Nullable
+  @Override
+  public PyType getCallableType(@NotNull Callable callable, @NotNull 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/userSkeletons/PyUserSkeletonsTypeProvider.java b/python/src/com/jetbrains/python/codeInsight/userSkeletons/PyUserSkeletonsTypeProvider.java
index f93b3b22978c..0a7d129b0018 100644
--- a/python/src/com/jetbrains/python/codeInsight/userSkeletons/PyUserSkeletonsTypeProvider.java
+++ b/python/src/com/jetbrains/python/codeInsight/userSkeletons/PyUserSkeletonsTypeProvider.java
@@ -46,4 +46,14 @@ public class PyUserSkeletonsTypeProvider extends PyTypeProviderBase {
     }
     return null;
   }
+
+  @Nullable
+  @Override
+  public PyType getCallableType(@NotNull Callable callable, @NotNull TypeEvalContext context) {
+    final Callable callableSkeleton = PyUserSkeletonsUtil.getUserSkeleton(callable);
+    if (callableSkeleton != null) {
+      return context.getType(callableSkeleton);
+    }
+    return null;
+  }
 }
diff --git a/python/src/com/jetbrains/python/psi/impl/PyFunctionImpl.java b/python/src/com/jetbrains/python/psi/impl/PyFunctionImpl.java
index 81392e382669..310641880a04 100644
--- a/python/src/com/jetbrains/python/psi/impl/PyFunctionImpl.java
+++ b/python/src/com/jetbrains/python/psi/impl/PyFunctionImpl.java
@@ -342,6 +342,12 @@ public class PyFunctionImpl extends PyPresentableElementImpl<PyFunctionStub> imp
 
   @Override
   public PyType getType(@NotNull TypeEvalContext context, @NotNull TypeEvalContext.Key key) {
+    for (PyTypeProvider provider : Extensions.getExtensions(PyTypeProvider.EP_NAME)) {
+      final PyType type = provider.getCallableType(this, context);
+      if (type != null) {
+        return type;
+      }
+    }
     final PyFunctionType type = new PyFunctionType(this);
     if (getDecoratorList() != null) {
       return PyUnionType.createWeakType(type);
diff --git a/python/src/com/jetbrains/python/psi/impl/PyLambdaExpressionImpl.java b/python/src/com/jetbrains/python/psi/impl/PyLambdaExpressionImpl.java
index fc31ce2e21a3..9bf0de63425e 100644
--- a/python/src/com/jetbrains/python/psi/impl/PyLambdaExpressionImpl.java
+++ b/python/src/com/jetbrains/python/psi/impl/PyLambdaExpressionImpl.java
@@ -1,6 +1,7 @@
 package com.jetbrains.python.psi.impl;
 
 import com.intellij.lang.ASTNode;
+import com.intellij.openapi.extensions.Extensions;
 import com.intellij.psi.util.PsiTreeUtil;
 import com.jetbrains.python.PyElementTypes;
 import com.jetbrains.python.codeInsight.controlflow.ControlFlowCache;
@@ -24,6 +25,12 @@ public class PyLambdaExpressionImpl extends PyElementImpl implements PyLambdaExp
   }
 
   public PyType getType(@NotNull TypeEvalContext context, @NotNull TypeEvalContext.Key key) {
+    for (PyTypeProvider provider : Extensions.getExtensions(PyTypeProvider.EP_NAME)) {
+      final PyType type = provider.getCallableType(this, context);
+      if (type != null) {
+        return type;
+      }
+    }
     return new PyFunctionType(this);
   }