in src/Analysis/Ast/Impl/Types/PythonFunctionOverload.cs [35:223]
public delegate IMember ReturnValueProvider(
IPythonModule declaringModule,
IPythonFunctionOverload overload,
IArgumentSet args,
IndexSpan indexSpan);
internal sealed class PythonFunctionOverload : LocatedMember, IPythonFunctionOverload {
private readonly string _returnDocumentation;
// Allow dynamic function specialization, such as defining return types for builtin
// functions that are impossible to scrape and that are missing from stubs.
// FormalGenericParameters: declaring module, overload for the return value, list of arguments.
private ReturnValueProvider _returnValueProvider;
// Return value can be an instance or a type info. Consider type(C()) returning
// type info of C vs. return C() that returns an instance of C.
private bool _fromAnnotation;
public PythonFunctionOverload(IPythonClassMember cm, FunctionDefinition fd, Location location, string returnDocumentation)
: this(cm, location) {
Documentation = fd.GetDocumentation();
cm.DeclaringModule.AddAstNode(this, fd);
_returnDocumentation = returnDocumentation;
}
public PythonFunctionOverload(IPythonClassMember cm, Location location) : base(location) {
ClassMember = cm ?? throw new ArgumentNullException(nameof(cm));
Name = cm.Name ?? throw new ArgumentNullException(nameof(cm.Name));
}
#region ILocatedMember
public override PythonMemberType MemberType => PythonMemberType.Function;
#endregion
internal void SetParameters(IReadOnlyList<IParameterInfo> parameters) => Parameters = parameters;
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();
if (!_fromAnnotation && !currentType.Equals(valueType)) {
var type = PythonUnionType.Combine(currentType, valueType);
// Track instance vs type info.
StaticReturnValue = value is IPythonInstance ? type.CreateInstance(ArgumentSet.WithoutContext) : (IMember)type;
}
}
internal void SetReturnValue(IMember value, bool fromAnnotation) {
StaticReturnValue = value;
_fromAnnotation = fromAnnotation;
}
internal void SetReturnValueProvider(ReturnValueProvider provider) => _returnValueProvider = provider;
internal void SetDocumentation(string documentation) => Documentation = documentation;
#region IPythonFunctionOverload
public FunctionDefinition FunctionDefinition => ClassMember?.DeclaringModule?.GetAstNode<FunctionDefinition>(this);
public IPythonClassMember ClassMember { get; }
public string Name { get; }
public string Documentation { get; private set; }
public string GetReturnDocumentation(IPythonType self = null) {
if (self != null) {
var returnType = GetSpecificReturnType(self as IPythonClassType, null);
if (!returnType.IsUnknown()) {
return returnType.GetPythonType().Name;
}
}
// Use annotation value if it is there
if(_fromAnnotation && !StaticReturnValue.IsUnknown()) {
return StaticReturnValue.GetPythonType().Name;
}
return _returnDocumentation;
}
public IReadOnlyList<IParameterInfo> Parameters { get; private set; } = Array.Empty<IParameterInfo>();
public IMember StaticReturnValue { get; private set; }
public IMember Call(IArgumentSet args, IPythonType self) {
if (!_fromAnnotation) {
// First try supplied specialization callback.
var rt = _returnValueProvider?.Invoke(args.Eval.Module, this, args, default);
if (!rt.IsUnknown()) {
return rt;
}
}
return GetSpecificReturnType(self as IPythonClassType, args);
}
#endregion
private IMember GetSpecificReturnType(IPythonClassType selfClassType, IArgumentSet args) {
var returnValueType = StaticReturnValue.GetPythonType();
switch (returnValueType) {
case PythonClassType cls when cls.IsGeneric:
return CreateSpecificReturnFromClassType(selfClassType, cls, args); // -> A[_T1, _T2, ...]
case IGenericType gt when gt.IsGeneric && args != null: // -> CLASS[T] on standalone function (i.e. -> List[T]).
var typeArgs = ExpressionEval.GetTypeArgumentsFromParameters(this, args);
if (typeArgs != null) {
return gt.CreateSpecificType(new ArgumentSet(typeArgs, args.Expression, args.Eval));
}
break;
case IGenericTypeParameter gtd1 when selfClassType != null:
return CreateSpecificReturnFromTypeVar(selfClassType, args, gtd1); // -> _T
case IGenericTypeParameter gtd2 when args != null: // -> T on standalone function.
return args.Arguments.FirstOrDefault(a => gtd2.Equals(a.Type))?.Value as IMember;
}
return StaticReturnValue;
}
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 newArgs = new ArgumentSet(typeArgs, args?.Expression, args?.Eval);
var specificReturnValue = returnClassType.CreateSpecificType(newArgs);
return specificReturnValue.CreateInstance(newArgs);
}
return null;
}
private IMember CreateSpecificReturnFromTypeVar(IPythonClassType selfClassType, IArgumentSet args, IGenericTypeParameter returnType) {
if (selfClassType.GetSpecificType(returnType.Name, out var specificType)) {
return specificType.CreateInstance(args);
}
// Find first base class type in which function was declared
var baseType = selfClassType.Mro
.OfType<IPythonClassType>()
.Skip(1)
.FirstOrDefault(b => b.GetMember(ClassMember.Name) != null && b.GenericParameters.ContainsKey(returnType.Name));
// Try and infer return value from base class
if (baseType != null && baseType.GetSpecificType(returnType.Name, out specificType)) {
return specificType.CreateInstance(args);
}
// Try getting type from passed in arguments
if (args?.Arguments.FirstOrDefault(a => returnType.Equals(a.Type))?.Value is IMember typeFromArgs) {
return typeFromArgs;
}
// Try getting the type from the type parameter bound
if (!returnType.Bound.IsUnknown()) {
return returnType.Bound.CreateInstance(args);
}
// 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 IGenericTypeParameter 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;
}
}