Handle generic standalone functions (#1172)

* Handle generic standalone functions (#1164)

* Initial

* Cleanup

* Usings

* Fix return type

* Restore lost code

* Merge mess

Author: Mikhail Arkhipov

src/Analysis/Ast/Impl/Analyzer/Evaluation/ExpressionEval.Generics.cs

@@ -194,6 +194,8 @@ private static IReadOnlyList<IPythonType> GetSpecificTypeFromArgumentValue(objec
                     var itemType = iter.GetIterator().Next.GetPythonType();
                     if (!itemType.IsUnknown()) {
                         specificTypes.Add(itemType);
+                    } else if(argumentValue is IPythonInstance inst) {
+                        specificTypes.Add(inst.GetPythonType());
                     }
                     break;
                 case IMember m:

src/Analysis/Ast/Impl/Analyzer/ModuleWalker.cs

@@ -281,8 +281,12 @@ private void MergeStub() {
             }
         }
 
-        private static bool IsStubBetterType(IPythonType sourceType, IPythonType stubType)
+        private static bool IsStubBetterType(IPythonType sourceType, IPythonType stubType) {
             // If stub says 'Any' but we have better type, keep the current type.
-            => sourceType.IsUnknown() || !(stubType.DeclaringModule is TypingModule) || stubType.Name != "Any";
+            if (stubType.IsUnknown()) {
+                return false;
+            }
+            return sourceType.IsUnknown() || !(stubType.DeclaringModule is TypingModule) || stubType.Name != "Any";
+        }
     }
 }

src/Analysis/Ast/Impl/Analyzer/Symbols/FunctionEvaluator.cs

@@ -117,7 +117,7 @@ private IPythonType TryDetermineReturnValue() {
                 // If instance could not be created, such as when return type is List[T] and
                 // type of T is not yet known, just use the type.
                  var instance = t.IsUnknown() ? annotationType : t;
-                _overload.SetReturnValue(instance, true); _overload.SetReturnValue(instance, true);
+                _overload.SetReturnValue(instance, true);
             } else {
                 // Check if function is a generator
                 var suite = FunctionDefinition.Body as SuiteStatement;
@@ -158,10 +158,9 @@ private void DeclareParameters(bool declareVariables) {
             // Declare parameters in scope
             IMember defaultValue = null;
             for (var i = skip; i < FunctionDefinition.Parameters.Length; i++) {
-                var isGeneric = false;
                 var p = FunctionDefinition.Parameters[i];
                 if (!string.IsNullOrEmpty(p.Name)) {
-                    var paramType = Eval.GetTypeFromAnnotation(p.Annotation);
+                    var paramType = Eval.GetTypeFromAnnotation(p.Annotation, out var isGeneric);
                     if (paramType.IsUnknown() && p.DefaultValue != null) {
                         defaultValue = Eval.GetValueFromExpression(p.DefaultValue);
                         // If parameter has default value, look for the annotation locally first
@@ -175,7 +174,7 @@ private void DeclareParameters(bool declareVariables) {
                     }
                     // If all else fails, look up globally.
                     paramType = paramType ?? Eval.GetTypeFromAnnotation(p.Annotation, out isGeneric) ?? Eval.UnknownType;
-                    var pi = new ParameterInfo(Ast, p, paramType, defaultValue, isGeneric);
+                    var pi = new ParameterInfo(Ast, p, paramType, defaultValue, isGeneric | paramType.IsGeneric());
                     if (declareVariables) {
                         DeclareParameter(p, pi);
                     }

src/Analysis/Ast/Impl/Types/PythonFunctionOverload.cs

@@ -15,6 +15,7 @@
 
 using System;
 using System.Collections.Generic;
+using System.Diagnostics;
 using System.Linq;
 using Microsoft.Python.Analysis.Analyzer.Evaluation;
 using Microsoft.Python.Analysis.Documents;
@@ -75,10 +76,12 @@ internal void AddReturnValue(IMember value) {
             if (value.IsUnknown()) {
                 return; // Don't add useless values.
             }
+
             if (StaticReturnValue.IsUnknown()) {
                 SetReturnValue(value, false);
                 return;
             }
+
             // If return value is set from annotation, it should not be changing.
             var currentType = StaticReturnValue.GetPythonType();
             var valueType = value.GetPythonType();
@@ -98,6 +101,7 @@ internal void SetReturnValueProvider(ReturnValueProvider provider)
             => _returnValueProvider = provider;
 
         #region IPythonFunctionOverload
+
         public FunctionDefinition FunctionDefinition { get; }
         public IPythonClassMember ClassMember { get; }
         public string Name { get; }
@@ -113,38 +117,11 @@ public string Documentation {
         }
 
         public string GetReturnDocumentation(IPythonType self = null) {
-            if (self == null) {
-                return _returnDocumentation;
-            }
-            var returnType = StaticReturnValue.GetPythonType();
-            switch (returnType) {
-                case PythonClassType cls when cls.IsGeneric(): {
-                        // -> A[_T1, _T2, ...]
-                        // Match arguments 
-                        var typeArgs = cls.GenericParameters.Keys
-                            .Select(n => cls.GenericParameters.TryGetValue(n, out var t) ? t : null)
-                            .ExcludeDefault()
-                            .ToArray();
-                        var specificReturnValue = cls.CreateSpecificType(new ArgumentSet(typeArgs));
-                        return specificReturnValue.Name;
-                    }
-                case IGenericTypeDefinition gtp1 when self is IPythonClassType cls: {
-                        // -> _T
-                        if (cls.GenericParameters.TryGetValue(gtp1.Name, out var specificType)) {
-                            return specificType.Name;
-                        }
-                        // Try returning the constraint
-                        // TODO: improve this, the heuristic is pretty basic and tailored to simple func(_T) -> _T
-                        var name = StaticReturnValue.GetPythonType()?.Name;
-                        var typeDefVar = DeclaringModule.Analysis.GlobalScope.Variables[name];
-                        if (typeDefVar?.Value is IGenericTypeDefinition gtp2) {
-                            var t = gtp2.Constraints.FirstOrDefault();
-                            if (t != null) {
-                                return t.Name;
-                            }
-                        }
-                        break;
-                    }
+            if (self != null) {
+                var returnType = GetSpecificReturnType(self as IPythonClassType, null);
+                if (!returnType.IsUnknown()) {
+                    return returnType.GetPythonType().Name;
+                }
             }
             return _returnDocumentation;
         }
@@ -161,59 +138,72 @@ public IMember Call(IArgumentSet args, IPythonType self, Node callLocation = nul
                 }
             }
 
-            // If function returns generic, determine actual type based on the passed in specific type (self).
-            // If there is no self and no declaring type, the function is standalone.
-            if (self == null && StaticReturnValue.IsGeneric() && Parameters.Any(p => p.IsGeneric)) {
-                return null; // Evaluate standalone generic with arguments instead.
-            }
-            if (!(self is IPythonClassType selfClassType)) {
-                return StaticReturnValue;
-            }
+            return GetSpecificReturnType(self as IPythonClassType, args);
+        }
+
+        #endregion
 
-            var returnType = StaticReturnValue.GetPythonType();
-            switch (returnType) {
+        private IMember GetSpecificReturnType(IPythonClassType selfClassType, IArgumentSet args) {
+            var returnValueType = StaticReturnValue.GetPythonType();
+            switch (returnValueType) {
                 case PythonClassType cls when cls.IsGeneric():
-                    // -> A[_T1, _T2, ...]
-                    // Match arguments
-                    IReadOnlyList<IPythonType> typeArgs = null;
-                    var classGenericParameters = selfClassType.GenericParameters.Keys.ToArray();
-                    if (classGenericParameters.Length > 0) {
-                        // Declaring class is specific and provides definitions of generic parameters
-                        typeArgs = classGenericParameters
-                            .Select(n => selfClassType.GenericParameters.TryGetValue(n, out var t) ? t : null)
-                            .ExcludeDefault()
-                            .ToArray();
-                    } else {
-                        typeArgs = ExpressionEval.GetTypeArgumentsFromParameters(this, args);
-                    }
-
-                    if (typeArgs != null) {
-                        var specificReturnValue = cls.CreateSpecificType(new ArgumentSet(typeArgs));
-                        return new PythonInstance(specificReturnValue);
-                    }
-                    break;
-
-                case IGenericTypeDefinition gtp1: {
-                        // -> _T
-                        if (selfClassType.GenericParameters.TryGetValue(gtp1.Name, out var specificType)) {
-                            return new PythonInstance(specificType);
-                        }
-                        // Try returning the constraint
-                        // TODO: improve this, the heuristic is pretty basic and tailored to simple func(_T) -> _T
-                        var name = StaticReturnValue.GetPythonType()?.Name;
-                        var typeDefVar = DeclaringModule.Analysis.GlobalScope.Variables[name];
-                        if (typeDefVar?.Value is IGenericTypeDefinition gtp2) {
-                            // See if the instance (self) type satisfies one of the constraints.
-                            return selfClassType.Mro.Any(b => gtp2.Constraints.Any(c => c.Equals(b)))
-                                ? selfClassType
-                                : gtp2.Constraints.FirstOrDefault();
-                        }
-
-                        break;
-                    }
+                    return CreateSpecificReturnFromClassType(selfClassType, cls, args); // -> A[_T1, _T2, ...]
+
+                case IGenericTypeDefinition gtd1 when selfClassType != null:
+                    return CreateSpecificReturnFromTypeVar(selfClassType, gtd1); // -> _T
+
+                case IGenericTypeDefinition gtd2 when args != null: // -> T on standalone function.
+                    return args.Arguments.FirstOrDefault(a => gtd2.Equals(a.Type))?.Value as IMember;
+
+                case IGenericType gt when args != null: // -> CLASS[T] on standalone function (i.e. -> List[T]).
+                    var typeArgs = ExpressionEval.GetTypeArgumentsFromParameters(this, args);
+                    Debug.Assert(typeArgs != null);
+                    return gt.CreateSpecificType(typeArgs);
             }
+
             return StaticReturnValue;
         }
-        #endregion
+
+        private IMember CreateSpecificReturnFromClassType(IPythonClassType selfClassType, PythonClassType returnClassType, IArgumentSet args) {
+            // -> A[_T1, _T2, ...]
+            // Match arguments
+            IReadOnlyList<IPythonType> typeArgs = null;
+            var classGenericParameters = selfClassType?.GenericParameters.Keys.ToArray() ?? Array.Empty<string>();
+            if (classGenericParameters.Length > 0 && selfClassType != null) {
+                // Declaring class is specific and provides definitions of generic parameters
+                typeArgs = classGenericParameters
+                    .Select(n => selfClassType.GenericParameters.TryGetValue(n, out var t) ? t : null)
+                    .ExcludeDefault()
+                    .ToArray();
+            } else if (args != null) {
+                typeArgs = ExpressionEval.GetTypeArgumentsFromParameters(this, args);
+            }
+
+            if (typeArgs != null) {
+                var specificReturnValue = returnClassType.CreateSpecificType(new ArgumentSet(typeArgs));
+                return new PythonInstance(specificReturnValue);
+            }
+
+            return null;
+        }
+
+        private IMember CreateSpecificReturnFromTypeVar(IPythonClassType selfClassType, IGenericTypeDefinition returnType) {
+            if (selfClassType.GenericParameters.TryGetValue(returnType.Name, out var specificType)) {
+                return new PythonInstance(specificType);
+            }
+
+            // Try returning the constraint
+            // TODO: improve this, the heuristic is pretty basic and tailored to simple func(_T) -> _T
+            var name = StaticReturnValue.GetPythonType()?.Name;
+            var typeDefVar = DeclaringModule.Analysis.GlobalScope.Variables[name];
+            if (typeDefVar?.Value is IGenericTypeDefinition gtp2) {
+                // See if the instance (self) type satisfies one of the constraints.
+                return selfClassType.Mro.Any(b => gtp2.Constraints.Any(c => c.Equals(b)))
+                    ? selfClassType
+                    : gtp2.Constraints.FirstOrDefault();
+            }
+
+            return null;
+        }
     }
 }