in packages/pyright-internal/src/analyzer/typeEvaluator.ts [14561:15000]
function getTypeOfFunction(node: FunctionNode): FunctionTypeResult | undefined {
const fileInfo = AnalyzerNodeInfo.getFileInfo(node);
// Is this type already cached?
const cachedFunctionType = readTypeCache(node.name, EvaluatorFlags.None) as FunctionType;
if (cachedFunctionType) {
if (!isFunction(cachedFunctionType)) {
// This can happen in certain rare circumstances where the
// function declaration falls within an unreachable code block.
return undefined;
}
return {
functionType: cachedFunctionType,
decoratedType: readTypeCache(node, EvaluatorFlags.None) || UnknownType.create(),
};
}
let functionDecl: FunctionDeclaration | undefined;
const decl = AnalyzerNodeInfo.getDeclaration(node);
if (decl) {
functionDecl = decl as FunctionDeclaration;
}
// There was no cached type, so create a new one.
// Retrieve the containing class node if the function is a method.
const containingClassNode = ParseTreeUtils.getEnclosingClass(node, /* stopAtFunction */ true);
let containingClassType: ClassType | undefined;
if (containingClassNode) {
const classInfo = getTypeOfClass(containingClassNode);
if (!classInfo) {
return undefined;
}
containingClassType = classInfo.classType;
}
let functionFlags = getFunctionFlagsFromDecorators(node, !!containingClassNode);
if (functionDecl?.isGenerator) {
functionFlags |= FunctionTypeFlags.Generator;
}
// Special-case magic method __class_getitem__, which is implicitly a class method.
if (containingClassNode && node.name.value === '__class_getitem__') {
functionFlags |= FunctionTypeFlags.ClassMethod;
}
if (fileInfo.isStubFile) {
functionFlags |= FunctionTypeFlags.StubDefinition;
} else if (fileInfo.isInPyTypedPackage && evaluatorOptions.disableInferenceForPyTypedSources) {
functionFlags |= FunctionTypeFlags.PyTypedDefinition;
}
if (node.isAsync) {
functionFlags |= FunctionTypeFlags.Async;
}
const functionType = FunctionType.createInstance(
node.name.value,
getFunctionFullName(node, fileInfo.moduleName, node.name.value),
fileInfo.moduleName,
functionFlags,
ParseTreeUtils.getDocString(node.suite.statements)
);
functionType.details.typeVarScopeId = getScopeIdForNode(node);
if (fileInfo.isBuiltInStubFile || fileInfo.isTypingStubFile || fileInfo.isTypingExtensionsStubFile) {
// Stash away the name of the function since we need to handle
// 'namedtuple', 'abstractmethod', 'dataclass' and 'NewType'
// specially.
functionType.details.builtInName = node.name.value;
}
functionType.details.declaration = functionDecl;
// Allow recursion by registering the partially-constructed
// function type.
const scope = ScopeUtils.getScopeForNode(node);
const functionSymbol = scope?.lookUpSymbolRecursive(node.name.value);
if (functionDecl && functionSymbol) {
setSymbolResolutionPartialType(functionSymbol.symbol, functionDecl, functionType);
}
writeTypeCache(node, functionType, /* flags */ undefined, /* isIncomplete */ false);
writeTypeCache(node.name, functionType, /* flags */ undefined, /* isIncomplete */ false);
// Is this an "__init__" method within a pseudo-generic class? If so,
// we'll add generic types to the constructor's parameters.
const addGenericParamTypes =
containingClassType &&
ClassType.isPseudoGenericClass(containingClassType) &&
node.name.value === '__init__';
const paramTypes: Type[] = [];
let typeParamIndex = 0;
// Determine if the first parameter should be skipped for comment-based
// function annotations.
let firstCommentAnnotationIndex = 0;
if (containingClassType && (functionType.details.flags & FunctionTypeFlags.StaticMethod) === 0) {
firstCommentAnnotationIndex = 1;
}
// If there is a function annotation comment, validate that it has the correct
// number of parameter annotations.
if (node.functionAnnotationComment && !node.functionAnnotationComment.isParamListEllipsis) {
const expected = node.parameters.length - firstCommentAnnotationIndex;
const received = node.functionAnnotationComment.paramTypeAnnotations.length;
// For methods with "self" or "cls" parameters, the annotation list
// can either include or exclude the annotation for the first parameter.
if (firstCommentAnnotationIndex > 0 && received === node.parameters.length) {
firstCommentAnnotationIndex = 0;
} else if (received !== expected) {
addError(
Localizer.Diagnostic.annotatedParamCountMismatch().format({
expected,
received,
}),
node.functionAnnotationComment
);
}
}
const markParamAccessed = (param: ParameterNode) => {
if (param.name) {
const symbolWithScope = lookUpSymbolRecursive(param.name, param.name.value, /* honorCodeFlow */ false);
if (symbolWithScope) {
setSymbolAccessed(fileInfo, symbolWithScope.symbol, param.name);
}
}
};
let paramsArePositionOnly = true;
node.parameters.forEach((param, index) => {
let paramType: Type | undefined;
let annotatedType: Type | undefined;
let isNoneWithoutOptional = false;
let paramTypeNode: ExpressionNode | undefined;
if (param.name) {
if (
index === 0 &&
containingClassType &&
(FunctionType.isClassMethod(functionType) ||
FunctionType.isInstanceMethod(functionType) ||
FunctionType.isConstructorMethod(functionType))
) {
// Mark "self/cls" as accessed.
markParamAccessed(param);
} else if (FunctionType.isAbstractMethod(functionType)) {
// Mark all parameters in abstract methods as accessed.
markParamAccessed(param);
} else if (containingClassType && ClassType.isProtocolClass(containingClassType)) {
// Mark all parameters in protocol methods as accessed.
markParamAccessed(param);
}
}
if (param.typeAnnotation) {
paramTypeNode = param.typeAnnotation;
} else if (param.typeAnnotationComment) {
paramTypeNode = param.typeAnnotationComment;
} else if (node.functionAnnotationComment && !node.functionAnnotationComment.isParamListEllipsis) {
const adjustedIndex = index - firstCommentAnnotationIndex;
if (adjustedIndex >= 0 && adjustedIndex < node.functionAnnotationComment.paramTypeAnnotations.length) {
paramTypeNode = node.functionAnnotationComment.paramTypeAnnotations[adjustedIndex];
}
}
if (paramTypeNode) {
annotatedType = getTypeOfAnnotation(paramTypeNode, {
associateTypeVarsWithScope: true,
allowTypeVarTuple: param.category === ParameterCategory.VarArgList,
disallowRecursiveTypeAlias: true,
});
if (isVariadicTypeVar(annotatedType) && !annotatedType.isVariadicUnpacked) {
addError(
Localizer.Diagnostic.unpackedTypeVarTupleExpected().format({
name1: annotatedType.details.name,
name2: annotatedType.details.name,
}),
paramTypeNode
);
annotatedType = UnknownType.create();
}
}
if (!annotatedType && addGenericParamTypes) {
if (index > 0 && param.category === ParameterCategory.Simple && param.name) {
annotatedType = containingClassType!.details.typeParameters[typeParamIndex];
typeParamIndex++;
}
}
if (annotatedType) {
const adjustedAnnotatedType = adjustParameterAnnotatedType(param, annotatedType);
if (adjustedAnnotatedType !== annotatedType) {
annotatedType = adjustedAnnotatedType;
isNoneWithoutOptional = true;
}
}
let defaultValueType: Type | undefined;
if (param.defaultValue) {
defaultValueType = getTypeOfExpression(
param.defaultValue,
annotatedType,
EvaluatorFlags.ConvertEllipsisToAny
).type;
}
if (annotatedType) {
// If there was both a type annotation and a default value, verify
// that the default value matches the annotation.
if (param.defaultValue && defaultValueType) {
const diagAddendum = new DiagnosticAddendum();
const typeVarMap = new TypeVarMap(functionType.details.typeVarScopeId);
if (containingClassType && containingClassType.details.typeVarScopeId !== undefined) {
if (node.name.value === '__init__' || node.name.value === '__new__') {
typeVarMap.addSolveForScope(containingClassType.details.typeVarScopeId);
}
}
if (!canAssignType(annotatedType, defaultValueType, diagAddendum, typeVarMap)) {
const diag = addDiagnostic(
fileInfo.diagnosticRuleSet.reportGeneralTypeIssues,
DiagnosticRule.reportGeneralTypeIssues,
Localizer.Diagnostic.paramAssignmentMismatch().format({
sourceType: printType(defaultValueType),
paramType: printType(annotatedType),
}) + diagAddendum.getString(),
param.defaultValue
);
if (isNoneWithoutOptional && paramTypeNode) {
const addOptionalAction: AddMissingOptionalToParamAction = {
action: Commands.addMissingOptionalToParam,
offsetOfTypeNode: paramTypeNode.start + 1,
};
if (diag) {
diag.addAction(addOptionalAction);
}
}
}
}
paramType = annotatedType;
}
const isPositionOnlyParam =
param.category === ParameterCategory.Simple && param.name && isPrivateName(param.name.value);
const isPositionOnlySeparator = param.category === ParameterCategory.Simple && !param.name;
if (index > 0 && paramsArePositionOnly && !isPositionOnlyParam && !isPositionOnlySeparator) {
// Insert an implicit "position-only parameter" separator.
FunctionType.addParameter(functionType, {
category: ParameterCategory.Simple,
type: UnknownType.create(),
});
}
if (!isPositionOnlyParam || isPositionOnlySeparator) {
paramsArePositionOnly = false;
}
const functionParam: FunctionParameter = {
category: param.category,
name: param.name ? param.name.value : undefined,
hasDefault: !!param.defaultValue,
defaultValueExpression: param.defaultValue,
defaultType: defaultValueType,
type: paramType || UnknownType.create(),
typeAnnotation: paramTypeNode,
hasDeclaredType: !!paramTypeNode,
};
FunctionType.addParameter(functionType, functionParam);
if (param.name) {
const variadicParamType = transformVariadicParamType(node, param.category, functionParam.type);
paramTypes.push(variadicParamType);
} else {
paramTypes.push(functionParam.type);
}
});
if (paramsArePositionOnly && functionType.details.parameters.length > 0) {
// Insert an implicit "position-only parameter" separator.
FunctionType.addParameter(functionType, {
category: ParameterCategory.Simple,
type: UnknownType.create(),
});
}
if (containingClassType) {
// If the first parameter doesn't have an explicit type annotation,
// provide a type if it's an instance, class or constructor method.
if (functionType.details.parameters.length > 0) {
const typeAnnotation = getTypeAnnotationForParameter(node, 0);
if (!typeAnnotation) {
const inferredParamType = inferFirstParamType(functionType.details.flags, containingClassType);
if (inferredParamType) {
functionType.details.parameters[0].type = inferredParamType;
if (!isAnyOrUnknown(inferredParamType)) {
functionType.details.parameters[0].isTypeInferred = true;
}
paramTypes[0] = inferredParamType;
}
}
}
}
// Update the types for the nodes associated with the parameters.
paramTypes.forEach((paramType, index) => {
const paramNameNode = node.parameters[index].name;
if (paramNameNode) {
if (isUnknown(paramType)) {
functionType.details.flags |= FunctionTypeFlags.UnannotatedParams;
}
writeTypeCache(paramNameNode, paramType, EvaluatorFlags.None, /* isIncomplete */ false);
}
});
// If the function ends in P.args and P.kwargs parameters, make it exempt from
// args/kwargs compatibility checks. This is important for protocol comparisons.
if (paramTypes.length >= 2) {
const paramType1 = paramTypes[paramTypes.length - 2];
const paramType2 = paramTypes[paramTypes.length - 1];
if (
isParamSpec(paramType1) &&
paramType1.paramSpecAccess === 'args' &&
isParamSpec(paramType2) &&
paramType2.paramSpecAccess === 'kwargs'
) {
functionType.details.flags |= FunctionTypeFlags.SkipArgsKwargsCompatibilityCheck;
}
}
// If there was a defined return type, analyze that first so when we
// walk the contents of the function, return statements can be
// validated against this type.
if (node.returnTypeAnnotation) {
// Temporarily set the return type to unknown in case of recursion.
functionType.details.declaredReturnType = UnknownType.create();
const returnType = getTypeOfAnnotation(node.returnTypeAnnotation, {
associateTypeVarsWithScope: true,
disallowRecursiveTypeAlias: true,
});
functionType.details.declaredReturnType = returnType;
} else if (node.functionAnnotationComment) {
// Temporarily set the return type to unknown in case of recursion.
functionType.details.declaredReturnType = UnknownType.create();
const returnType = getTypeOfAnnotation(node.functionAnnotationComment.returnTypeAnnotation, {
associateTypeVarsWithScope: true,
disallowRecursiveTypeAlias: true,
});
functionType.details.declaredReturnType = returnType;
} else {
// If there was no return type annotation and this is a type stub,
// we have no opportunity to infer the return type, so we'll indicate
// that it's unknown.
if (fileInfo.isStubFile) {
// Special-case the __init__ method, which is commonly left without
// an annotated return type, but we can assume it returns None.
if (node.name.value === '__init__') {
functionType.details.declaredReturnType = NoneType.createInstance();
} else {
functionType.details.declaredReturnType = UnknownType.create();
}
}
}
// If the return type is explicitly annotated as a generator, mark the
// function as a generator even though it may not contain a "yield" statement.
// This is important for generator functions declared in stub files, abstract
// methods or protocol definitions.
if (fileInfo.isStubFile || ParseTreeUtils.isSuiteEmpty(node.suite)) {
if (
functionType.details.declaredReturnType &&
isClassInstance(functionType.details.declaredReturnType) &&
ClassType.isBuiltIn(functionType.details.declaredReturnType, [
'Generator',
'AsyncGenerator',
'AwaitableGenerator',
])
) {
functionType.details.flags |= FunctionTypeFlags.Generator;
}
}
// If it's an async function, wrap the return type in an Awaitable or Generator.
const preDecoratedType = node.isAsync ? createAsyncFunction(node, functionType) : functionType;
// Apply all of the decorators in reverse order.
let decoratedType: Type = preDecoratedType;
let foundUnknown = false;
for (let i = node.decorators.length - 1; i >= 0; i--) {
const decorator = node.decorators[i];
const newDecoratedType = applyFunctionDecorator(decoratedType, functionType, decorator, node);
if (containsUnknown(newDecoratedType)) {
// Report this error only on the first unknown type.
if (!foundUnknown) {
addDiagnostic(
fileInfo.diagnosticRuleSet.reportUntypedFunctionDecorator,
DiagnosticRule.reportUntypedFunctionDecorator,
Localizer.Diagnostic.functionDecoratorTypeUnknown(),
node.decorators[i].expression
);
foundUnknown = true;
}
} else {
// Apply the decorator only if the type is known.
decoratedType = newDecoratedType;
}
}
// See if there are any overloads provided by previous function declarations.
if (isFunction(decoratedType)) {
if (FunctionType.isOverloaded(decoratedType)) {
// Mark all the parameters as accessed.
node.parameters.forEach((param) => {
markParamAccessed(param);
});
}
decoratedType = addOverloadsToFunctionType(node, decoratedType);
}
writeTypeCache(node.name, functionType, EvaluatorFlags.None, /* isIncomplete */ false);
writeTypeCache(node, decoratedType, EvaluatorFlags.None, /* isIncomplete */ false);
return { functionType, decoratedType };
}