Override function signatures in user skeletons

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);
 }

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);
   }

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;
+  }
 }

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);

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);
   }