Skip to content

Commit

Permalink
fix inlay hint insertions for typing/typing_extensions depending …
Browse files Browse the repository at this point in the history 
…on the python version
  • Loading branch information
@DetachHead
DetachHead committed Jan 26, 2025
1 parent 565d0e9 commit 723e70a
Show file tree
Hide file tree
Showing 2 changed files with 46 additions and 33 deletions.
24 changes: 14 additions & 10 deletions packages/pyright-internal/src/analyzer/typeInlayHintsWalker.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -147,8 +147,10 @@ export class TypeInlayHintsWalker extends ParseTreeWalker {
!isUnderscoreOnlyName(node.d.value) &&
!this._variablesThatShouldntHaveInlayHints.has(node)
) {
const type = this._program.evaluator?.getType(node);
const evaluator = this._program.evaluator;
const type = evaluator?.getType(node);
if (
evaluator &&
type &&
!isAny(type) &&
!(isClass(type) && isLiteralType(type)) &&
Expand All @@ -164,10 +166,10 @@ export class TypeInlayHintsWalker extends ParseTreeWalker {
getTypeAliasInfo(type)?.shared.name === node.d.value
) {
inlayHintValue = 'TypeAlias';
importTracker = new ImportTracker(this._fileUri);
importTracker.add('typing', inlayHintValue);
importTracker = new ImportTracker(this._fileUri, (name) => evaluator.getTypingType(node, name));
importTracker.addTypingImport(inlayHintValue);
} else {
const result = this._printType(type);
const result = this._printType(node, type);
inlayHintValue = result.value;
importTracker = result.imports;
}
Expand Down Expand Up @@ -196,7 +198,7 @@ export class TypeInlayHintsWalker extends ParseTreeWalker {
if (functionType !== undefined && !functionType.shared.declaredReturnType) {
const inferredReturnType = evaluator?.getInferredReturnType(functionType);
if (inferredReturnType) {
const { imports, value } = this._printType(inferredReturnType);
const { imports, value } = this._printType(node, inferredReturnType);
this.featureItems.push({
inlayHintType: 'functionReturn',
position: node.d.suite.start,
Expand All @@ -222,7 +224,7 @@ export class TypeInlayHintsWalker extends ParseTreeWalker {
) {
const valueType = evaluator.getType(node.d.valueExpr);
if (valueType) {
const { value, imports } = this._printType(valueType);
const { value, imports } = this._printType(node, valueType);
this.featureItems.push({
inlayHintType: 'generic',
position: this._endOfNode(node),
Expand Down Expand Up @@ -268,8 +270,8 @@ export class TypeInlayHintsWalker extends ParseTreeWalker {
) {
const printedTypeArgs = returnType.priv.typeArgs.flatMap((typeArg) =>
isClass(typeArg) && typeArg.priv.tupleTypeArgs && typeArg.priv.isUnpacked
? typeArg.priv.tupleTypeArgs.map((tupleTypeArg) => this._printType(tupleTypeArg.type))
: this._printType(typeArg)
? typeArg.priv.tupleTypeArgs.map((tupleTypeArg) => this._printType(node, tupleTypeArg.type))
: this._printType(node, typeArg)
);
const imports = printedTypeArgs.map((inlayHintInfo) => inlayHintInfo.imports);
const values = printedTypeArgs.map((inlayHintInfo) => inlayHintInfo.value);
Expand Down Expand Up @@ -352,8 +354,10 @@ export class TypeInlayHintsWalker extends ParseTreeWalker {

private _endOfNode = (node: ParseNode) => node.start + node.length;

private _printType = (type: Type): { value: string; imports: ImportTracker } => {
const importTracker = new ImportTracker(this._fileUri);
private _printType = (node: ParseNode, type: Type): { value: string; imports: ImportTracker } => {
const importTracker = new ImportTracker(this._fileUri, (name) =>
this._program.evaluator!.getTypingType(node, name)
);
return {
value: this._program.evaluator!.printType(type, { enforcePythonSyntax: true, importTracker }),
imports: importTracker,
Expand Down
55 changes: 32 additions & 23 deletions packages/pyright-internal/src/analyzer/typePrinter.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -155,7 +155,7 @@ export function printLiteralValueTruncated(type: ClassType, importTracker: Impor

assert(type.shared.name === 'str');
const name = 'LiteralString';
importTracker?.add('typing', name);
importTracker?.addTypingImport(name);
return name;
}

Expand Down Expand Up @@ -219,7 +219,7 @@ export class ImportTracker {
private readonly _importFroms = new Map<string, Set<string>>();
readonly result: ImportTrackerResults = { imports: this._imports, importFroms: this._importFroms };

constructor(private _fileUri: Uri) {}
constructor(private _fileUri: Uri, private _getTypingType: (name: string) => Type | undefined) {}

/**
* @param module the name of the module being imported. if it's possible for the module to be the same as the current module, you should
Expand All @@ -246,6 +246,15 @@ export class ImportTracker {
this._imports.add(module);
}
};

/**
* use this instead of {@link add} when importing something from `typing`, since it may need to import from `typing_extensions` instead
* depending on the python version
*/
addTypingImport = (name: string) => {
const module = this._getTypingType(name);
this.add(module?.category === TypeCategory.Class ? module.shared.moduleName : 'typing', name);
};
}

function printTypeInternal(
Expand All @@ -259,7 +268,7 @@ function printTypeInternal(
): string {
if (recursionCount > maxTypeRecursionCount) {
if (printTypeFlags & PrintTypeFlags.PythonSyntax) {
importTracker?.add('typing', 'Any');
importTracker?.addTypingImport('Any');
return 'Any';
}
return '<Recursive>';
Expand Down Expand Up @@ -459,23 +468,23 @@ function printTypeInternal(
switch (type.category) {
case TypeCategory.Unbound: {
if (printTypeFlags & PrintTypeFlags.PythonSyntax) {
importTracker?.add('typing', 'Any');
importTracker?.addTypingImport('Any');
return 'Any';
}
return 'Unbound';
}

case TypeCategory.Unknown: {
if (printTypeFlags & (PrintTypeFlags.PythonSyntax | PrintTypeFlags.PrintUnknownWithAny)) {
importTracker?.add('typing', 'Any');
importTracker?.addTypingImport('Any');
return 'Any';
}
return 'Unknown';
}

case TypeCategory.Module: {
if (printTypeFlags & PrintTypeFlags.PythonSyntax) {
importTracker?.add('typing', 'Any');
importTracker?.addTypingImport('Any');
return 'Any';
}
return `Module("${type.priv.moduleName}")`;
Expand All @@ -487,7 +496,7 @@ function printTypeInternal(
if (isLiteralValueTruncated(type) && (printTypeFlags & PrintTypeFlags.PythonSyntax) !== 0) {
return printLiteralValueTruncated(type, importTracker);
} else {
importTracker?.add('typing', 'Literal');
importTracker?.addTypingImport('Literal');
return `Literal[${printLiteralValue(type, "'", importTracker)}]`;
}
}
Expand All @@ -508,7 +517,7 @@ function printTypeInternal(
if (isLiteralValueTruncated(type) && (printTypeFlags & PrintTypeFlags.PythonSyntax) !== 0) {
typeToWrap = printLiteralValueTruncated(type, importTracker);
} else {
importTracker?.add('typing', 'Literal');
importTracker?.addTypingImport('Literal');
typeToWrap = `Literal[${printLiteralValue(type, "'", importTracker)}]`;
}

Expand Down Expand Up @@ -582,8 +591,8 @@ function printTypeInternal(
);

if ((printTypeFlags & PrintTypeFlags.PythonSyntax) !== 0) {
importTracker?.add('typing', 'Callable');
importTracker?.add('typing', 'Any');
importTracker?.addTypingImport('Callable');
importTracker?.addTypingImport('Any');
return 'Callable[..., Any]';
}

Expand Down Expand Up @@ -675,7 +684,7 @@ function printTypeInternal(
) {
boundTypeString = `Self@${boundTypeString}`;
} else {
importTracker?.add('typing', 'Self');
importTracker?.addTypingImport('Self');
boundTypeString = `Self`;
}
}
Expand All @@ -687,7 +696,7 @@ function printTypeInternal(
return boundTypeString;
}
if ((printTypeFlags & (PrintTypeFlags.PrintUnknownWithAny | PrintTypeFlags.PythonSyntax)) !== 0) {
importTracker?.add('typing', 'Self');
importTracker?.addTypingImport('Self');
return 'Any';
} else {
return 'Unknown';
Expand Down Expand Up @@ -718,7 +727,7 @@ function printTypeInternal(
}

if (isTypeVarTuple(type) && type.priv.isInUnion) {
importTracker?.add('typing', 'Union');
importTracker?.addTypingImport('Union');
typeVarName = `Union[${typeVarName}]`;
}

Expand All @@ -738,7 +747,7 @@ function printTypeInternal(

case TypeCategory.Never: {
const result = type.priv.isNoReturn ? 'NoReturn' : 'Never';
importTracker?.add('typing', result);
importTracker?.addTypingImport(result);
return result;
}

Expand All @@ -747,7 +756,7 @@ function printTypeInternal(
if (anyType.priv.isEllipsis) {
return '...';
}
importTracker?.add('typing', 'Any');
importTracker?.addTypingImport('Any');
return 'Any';
}
}
Expand Down Expand Up @@ -850,7 +859,7 @@ function printUnionType(
return unionString;
}

importTracker?.add('typing', 'Optional');
importTracker?.addTypingImport('Optional');
return 'Optional[' + optionalType + ']';
}

Expand Down Expand Up @@ -892,14 +901,14 @@ function printUnionType(
if (literalObjectStrings.size > 0) {
const literalStrings: string[] = [];
literalObjectStrings.forEach((s) => literalStrings.push(s));
importTracker?.add('typing', 'Literal');
importTracker?.addTypingImport('Literal');
dedupedSubtypeStrings.push(`Literal[${literalStrings.join(', ')}]`);
}

if (literalClassStrings.size > 0) {
const literalStrings: string[] = [];
literalClassStrings.forEach((s) => literalStrings.push(s));
importTracker?.add('typing', 'Literal');
importTracker?.addTypingImport('Literal');
dedupedSubtypeStrings.push(`type[Literal[${literalStrings.join(', ')}]]`);
}

Expand All @@ -914,7 +923,7 @@ function printUnionType(
}
return unionString;
}
importTracker?.add('typing', 'Union');
importTracker?.addTypingImport('Union');
return `Union[${dedupedSubtypeStrings.join(', ')}]`;
}

Expand Down Expand Up @@ -958,7 +967,7 @@ function printFunctionType(
importTracker
);
} else {
importTracker?.add('typing', 'Any');
importTracker?.addTypingImport('Any');
}

let result: string;
Expand All @@ -981,15 +990,15 @@ function printFunctionType(
)
);
} else {
importTracker?.add('typing', 'Any');
importTracker?.addTypingImport('Any');
paramTypes.push('Any');
}
}
});

if (paramSpec) {
if (paramTypes.length > 0) {
importTracker?.add('typing', 'Concatenate');
importTracker?.addTypingImport('Concatenate');
result = `Concatenate[${paramTypes.join(', ')}, ${paramSpec.shared.name}]`;
} else {
result = paramSpec.shared.name;
Expand All @@ -1002,7 +1011,7 @@ function printFunctionType(
// a "catch all" Callable.
result = '...';
}
importTracker?.add('typing', 'Callable');
importTracker?.addTypingImport('Callable');
return FunctionType.isParamSpecValue(type) ? result : `Callable[${result}, ${returnTypeString}]`;
} else {
const parts = printFunctionPartsInternal(
Expand Down

0 comments on commit 723e70a

Please sign in to comment.