microsoft/pyright
1
1
Fork 0
mirror of https://github.com/microsoft/pyright.git synced 2024-10-07 13:29:17 +03:00
Code Issues Projects Releases Wiki Activity

Added completion suggestion support for TypedDict keys and values. Thanks to Robert Cragie for this contribution!

Browse Source
This commit is contained in:
Eric Traut 2021-09-25 02:24:45 -07:00
parent 67a0bb459e
commit 464d21ab2e
4 changed files with 693 additions and 0 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

159
packages/pyright-internal/src/languageService/completionProvider.ts
View File

@ -90,6 +90,7 @@ import {
ArgumentCategory,
DecoratorNode,
DictionaryKeyEntryNode,
DictionaryNode,
ErrorExpressionCategory,
ErrorNode,
ExpressionNode,
@ -102,6 +103,7 @@ import {
ParameterNode,
ParseNode,
ParseNodeType,
SetNode,
StringNode,
} from '../parser/parseNodes';
import { ParseResults } from '../parser/parser';
@ -435,6 +437,13 @@ export class CompletionProvider {
return this._getMemberAccessCompletions(curNode.leftExpression, priorWord);
}
if (curNode.nodeType === ParseNodeType.Dictionary) {
const completionList = CompletionList.create();
if (this._addTypedDictKeys(curNode, /* stringNode */ undefined, priorText, postText, completionList)) {
return { completionList };
}
}
if (curNode.nodeType === ParseNodeType.Name) {
// This condition is little different than others since it does its own
// tree walk up to find context and let outer tree walk up to proceed if it can't find
@ -1435,6 +1444,42 @@ export class CompletionProvider {
});
}
private _getDictExpressionStringKeys(parseNode: ParseNode, excludeIds?: Set<number | undefined>) {
const node = getDictionaryLikeNode(parseNode);
if (!node) {
return [];
}
return node.entries.flatMap((entry) => {
if (entry.nodeType !== ParseNodeType.DictionaryKeyEntry || excludeIds?.has(entry.keyExpression.id)) {
return [];
}
if (entry.keyExpression.nodeType === ParseNodeType.StringList) {
return [entry.keyExpression.strings.map((s) => s.value).join('')];
}
return [];
});
function getDictionaryLikeNode(parseNode: ParseNode) {
// this method assumes the given parseNode is either a child of a dictionary or a dictionary itself
if (parseNode.nodeType === ParseNodeType.Dictionary) {
return parseNode;
}
let curNode: ParseNode | undefined = parseNode;
while (curNode && curNode.nodeType !== ParseNodeType.Dictionary && curNode.nodeType !== ParseNodeType.Set) {
curNode = curNode.parent;
if (!curNode) {
return;
}
}
return curNode;
}
}
private _getSubTypesWithLiteralValues(type: Type) {
const values: ClassType[] = [];
@ -1605,6 +1650,30 @@ export class CompletionProvider {
);
return { completionList };
}
if (parseNode.nodeType === ParseNodeType.String && parseNode.parent?.parent) {
const stringParent = parseNode.parent.parent;
// If the dictionary is not yet filled in, it will appear as though it's
// a set initially.
let dictOrSet: DictionaryNode | SetNode | undefined;
if (
stringParent.nodeType === ParseNodeType.DictionaryKeyEntry &&
stringParent.keyExpression === parseNode.parent &&
stringParent.parent?.nodeType === ParseNodeType.Dictionary
) {
dictOrSet = stringParent.parent;
} else if (stringParent?.nodeType === ParseNodeType.Set) {
dictOrSet = stringParent;
}
if (dictOrSet) {
if (this._addTypedDictKeys(dictOrSet, parseNode, priorText, postText, completionList)) {
return { completionList };
}
}
}
}
if (parentNode.nodeType !== ParseNodeType.Argument) {
@ -1670,6 +1739,96 @@ export class CompletionProvider {
return { completionList };
}
private _addTypedDictKeys(
dictionaryNode: DictionaryNode | SetNode,
stringNode: StringNode | undefined,
priorText: string,
postText: string,
completionList: CompletionList
) {
const expectedTypeResult = this._evaluator.getExpectedType(dictionaryNode);
if (!expectedTypeResult) {
return false;
}
// If the expected type result is associated with a node above the
// dictionaryNode in the parse tree, there are no typed dict keys to add.
if (ParseTreeUtils.getNodeDepth(expectedTypeResult.node) < ParseTreeUtils.getNodeDepth(dictionaryNode)) {
return false;
}
let typedDicts: ClassType[] = [];
doForEachSubtype(expectedTypeResult.type, (subtype) => {
if (isClassInstance(subtype) && ClassType.isTypedDictClass(subtype)) {
typedDicts.push(subtype);
}
});
if (typedDicts.length === 0) {
return false;
}
const keys = this._getDictExpressionStringKeys(
dictionaryNode,
stringNode ? new Set([stringNode.parent?.id]) : undefined
);
typedDicts = this._tryNarrowTypedDicts(typedDicts, keys);
const quoteValue = this._getQuoteValueFromPriorText(priorText);
const excludes = new Set(completionList.items.map((i) => i.label));
keys.forEach((key) => {
excludes.add(key);
});
typedDicts.forEach((typedDict) => {
getTypedDictMembersForClass(this._evaluator, typedDict, /* allowNarrowed */ true).forEach((_, key) => {
// Unions of TypedDicts may define the same key.
if (excludes.has(key)) {
return;
}
excludes.add(key);
this._addStringLiteralToCompletionList(
key,
quoteValue ? quoteValue.stringValue : undefined,
postText,
quoteValue
? quoteValue.quoteCharacter
: this._parseResults.tokenizerOutput.predominantSingleQuoteCharacter,
completionList
);
});
});
return true;
}
private _tryNarrowTypedDicts(types: ClassType[], keys: string[]): ClassType[] {
const newTypes = types.flatMap((type) => {
const entries = getTypedDictMembersForClass(this._evaluator, type, /* allowNarrowed */ true);
for (let index = 0; index < keys.length; index++) {
if (!entries.has(keys[index])) {
return [];
}
}
return [type];
});
if (newTypes.length === 0) {
// Couldn't narrow to any typed dicts. Just include all.
return types;
}
return newTypes;
}
// Given a string of text that precedes the current insertion point,
// determines which portion of it is the first part of a string literal
// (either starting with a single or double quote). Returns the quote

236
packages/pyright-internal/src/tests/fourslash/completions.call.typedDict.fourslash.ts Normal file
View File

@ -0,0 +1,236 @@
/// <reference path="fourslash.ts" />
// @filename: test.py
//// from typing import TypedDict, Optional, Union, List, Dict, Any
////
//// class Movie(TypedDict):
//// name: str
//// age: int
////
//// def thing(movie: Movie):
//// pass
////
//// thing({'[|/*marker1*/|]'})
//// thing({'name': '[|/*marker2*/|]'})
//// thing({'name': 'Robert','[|/*marker3*/|]'})
//// thing({'name': 'Robert', '[|/*marker4*/|]'})
//// thing('[|/*marker5*/|]')
//// thing({'na[|/*marker6*/|]'})
//// thing({[|/*marker7*/|]})
//// thing({'a', '[|/*marker8*/|]'})
////
//// class Episode(TypedDict):
//// title: str
//// score: int
////
//// def thing2(item: Union[Episode, Movie]):
//// pass
////
//// thing2({'[|/*marker9*/|]'})
//// thing2({'unknown': 'a', '[|/*marker10*/|]': ''})
//// thing2({'title': 'Episode 01', '[|/*marker11*/|]': ''})
////
//// class Wrapper(TypedDict):
//// age: int
//// wrapped: Union[bool, Movie]
//// data: Dict[str, Any]
////
//// def thing3(wrapper: Optional[Wrapper]):
//// pass
////
//// thing3({'data': {'[|/*marker12*/|]'}})
//// thing3({'wrapped': {'[|/*marker13*/|]'}})
//// thing3({'age': 1, 'wrapped': {'[|/*marker14*/|]'}})
//// thing3({'unknown': {'[|/*marker15*/|]'}})
//// thing3({'age': {'[|/*marker16*/|]'}})
{
const marker1Range = helper.expandPositionRange(helper.getPositionRange('marker1'), 1, 1);
const marker3Range = helper.expandPositionRange(helper.getPositionRange('marker3'), 1, 1);
const marker4Range = helper.expandPositionRange(helper.getPositionRange('marker4'), 1, 1);
const marker6Range = helper.expandPositionRange(helper.getPositionRange('marker6'), 3, 1);
const marker7Range = helper.getPositionRange('marker7');
const marker8Range = helper.expandPositionRange(helper.getPositionRange('marker8'), 1, 1);
const marker9Range = helper.expandPositionRange(helper.getPositionRange('marker9'), 1, 1);
const marker10Range = helper.expandPositionRange(helper.getPositionRange('marker10'), 1, 1);
const marker11Range = helper.expandPositionRange(helper.getPositionRange('marker11'), 1, 1);
const marker13Range = helper.expandPositionRange(helper.getPositionRange('marker13'), 1, 1);
const marker14Range = helper.expandPositionRange(helper.getPositionRange('marker14'), 1, 1);
// @ts-ignore
await helper.verifyCompletion('exact', 'markdown', {
marker1: {
completions: [
{
label: "'name'",
kind: Consts.CompletionItemKind.Constant,
textEdit: { range: marker1Range, newText: "'name'" },
},
{
label: "'age'",
kind: Consts.CompletionItemKind.Constant,
textEdit: { range: marker1Range, newText: "'age'" },
},
],
},
marker2: {
completions: [],
},
marker3: {
completions: [
{
label: "'age'",
kind: Consts.CompletionItemKind.Constant,
textEdit: { range: marker3Range, newText: "'age'" },
},
],
},
marker4: {
completions: [
{
label: "'age'",
kind: Consts.CompletionItemKind.Constant,
textEdit: { range: marker4Range, newText: "'age'" },
},
],
},
marker5: {
completions: [],
},
marker6: {
completions: [
{
label: "'name'",
kind: Consts.CompletionItemKind.Constant,
textEdit: { range: marker6Range, newText: "'name'" },
},
],
},
marker8: {
completions: [
{
label: "'age'",
kind: Consts.CompletionItemKind.Constant,
textEdit: { range: marker8Range, newText: "'age'" },
},
{
label: "'name'",
kind: Consts.CompletionItemKind.Constant,
textEdit: { range: marker8Range, newText: "'name'" },
},
],
},
marker9: {
completions: [
{
label: "'name'",
kind: Consts.CompletionItemKind.Constant,
textEdit: { range: marker9Range, newText: "'name'" },
},
{
label: "'age'",
kind: Consts.CompletionItemKind.Constant,
textEdit: { range: marker9Range, newText: "'age'" },
},
{
label: "'title'",
kind: Consts.CompletionItemKind.Constant,
textEdit: { range: marker9Range, newText: "'title'" },
},
{
label: "'score'",
kind: Consts.CompletionItemKind.Constant,
textEdit: { range: marker9Range, newText: "'score'" },
},
],
},
marker10: {
completions: [
{
label: "'name'",
kind: Consts.CompletionItemKind.Constant,
textEdit: { range: marker10Range, newText: "'name'" },
},
{
label: "'age'",
kind: Consts.CompletionItemKind.Constant,
textEdit: { range: marker10Range, newText: "'age'" },
},
{
label: "'title'",
kind: Consts.CompletionItemKind.Constant,
textEdit: { range: marker10Range, newText: "'title'" },
},
{
label: "'score'",
kind: Consts.CompletionItemKind.Constant,
textEdit: { range: marker10Range, newText: "'score'" },
},
],
},
marker11: {
completions: [
{
label: "'score'",
kind: Consts.CompletionItemKind.Constant,
textEdit: { range: marker11Range, newText: "'score'" },
},
],
},
marker12: {
completions: [],
},
marker13: {
completions: [
{
label: "'name'",
kind: Consts.CompletionItemKind.Constant,
textEdit: { range: marker13Range, newText: "'name'" },
},
{
label: "'age'",
kind: Consts.CompletionItemKind.Constant,
textEdit: { range: marker13Range, newText: "'age'" },
},
],
},
marker14: {
completions: [
{
label: "'name'",
kind: Consts.CompletionItemKind.Constant,
textEdit: { range: marker14Range, newText: "'name'" },
},
{
label: "'age'",
kind: Consts.CompletionItemKind.Constant,
textEdit: { range: marker14Range, newText: "'age'" },
},
],
},
marker15: {
completions: [],
},
marker16: {
completions: [],
},
});
// @ts-ignore
await helper.verifyCompletion('included', 'markdown', {
marker7: {
completions: [
{
label: "'name'",
kind: Consts.CompletionItemKind.Constant,
textEdit: { range: marker7Range, newText: "'name'" },
},
{
label: "'age'",
kind: Consts.CompletionItemKind.Constant,
textEdit: { range: marker7Range, newText: "'age'" },
},
],
},
});
}

163
packages/pyright-internal/src/tests/fourslash/completions.call.typedDict.list.fourslash.ts Normal file
View File

@ -0,0 +1,163 @@
/// <reference path="fourslash.ts" />
// @filename: test.py
//// from typing import TypedDict, Union, List
////
//// class Movie(TypedDict):
//// name: str
//// age: int
////
//// class MultipleInputs(TypedDict):
//// items: List[Movie]
//// union: Union[bool, List[Movie]]
//// unions: Union[Movie, Union[bool, List[Movie]]]
////
//// def thing(inputs: MultipleInputs):
//// pass
////
//// thing({'items': ['[|/*marker1*/|]']})
//// thing({'items': {'[|/*marker2*/|]'}})
//// thing({'items': [{'[|/*marker3*/|]'}]})
//// thing({'union': [{'[|/*marker4*/|]'}]})
//// thing({'unions': {'[|/*marker5*/|]'}})
//// thing({'unions': [{'[|/*marker6*/|]'}]})
////
//// def thing2(movies: List[Movie]):
//// pass
////
//// thing2([{'[|/*marker7*/|]'}])
//// thing2({'[|/*marker8*/|]'})
////
//// class Wrapper(TypedDict):
//// wrapped: MultipleInputs
////
//// def thing3(wrapper: Wrapper):
//// pass
////
//// thing3({'wrapped': {'items': [{'[|/*marker9*/|]'}]}})
//// thing3({'wrapped': {'items': {'[|/*marker10*/|]'}}})
//// thing3({'wrapped': {'items': [{'a': 'b'}, {'[|/*marker11*/|]'}]}})
{
const marker3Range = helper.expandPositionRange(helper.getPositionRange('marker3'), 1, 1);
const marker4Range = helper.expandPositionRange(helper.getPositionRange('marker4'), 1, 1);
const marker5Range = helper.expandPositionRange(helper.getPositionRange('marker5'), 1, 1);
const marker6Range = helper.expandPositionRange(helper.getPositionRange('marker6'), 1, 1);
const marker7Range = helper.expandPositionRange(helper.getPositionRange('marker7'), 1, 1);
const marker9Range = helper.expandPositionRange(helper.getPositionRange('marker9'), 1, 1);
const marker11Range = helper.expandPositionRange(helper.getPositionRange('marker11'), 1, 1);
// @ts-ignore
await helper.verifyCompletion('exact', 'markdown', {
marker1: {
completions: [],
},
marker2: {
completions: [],
},
marker3: {
completions: [
{
label: "'name'",
kind: Consts.CompletionItemKind.Constant,
textEdit: { range: marker3Range, newText: "'name'" },
},
{
label: "'age'",
kind: Consts.CompletionItemKind.Constant,
textEdit: { range: marker3Range, newText: "'age'" },
},
],
},
marker4: {
completions: [
{
label: "'name'",
kind: Consts.CompletionItemKind.Constant,
textEdit: { range: marker4Range, newText: "'name'" },
},
{
label: "'age'",
kind: Consts.CompletionItemKind.Constant,
textEdit: { range: marker4Range, newText: "'age'" },
},
],
},
marker5: {
completions: [
{
label: "'name'",
kind: Consts.CompletionItemKind.Constant,
textEdit: { range: marker5Range, newText: "'name'" },
},
{
label: "'age'",
kind: Consts.CompletionItemKind.Constant,
textEdit: { range: marker5Range, newText: "'age'" },
},
],
},
marker6: {
completions: [
{
label: "'name'",
kind: Consts.CompletionItemKind.Constant,
textEdit: { range: marker6Range, newText: "'name'" },
},
{
label: "'age'",
kind: Consts.CompletionItemKind.Constant,
textEdit: { range: marker6Range, newText: "'age'" },
},
],
},
marker7: {
completions: [
{
label: "'name'",
kind: Consts.CompletionItemKind.Constant,
textEdit: { range: marker7Range, newText: "'name'" },
},
{
label: "'age'",
kind: Consts.CompletionItemKind.Constant,
textEdit: { range: marker7Range, newText: "'age'" },
},
],
},
marker8: {
completions: [],
},
marker9: {
completions: [
{
label: "'name'",
kind: Consts.CompletionItemKind.Constant,
textEdit: { range: marker9Range, newText: "'name'" },
},
{
label: "'age'",
kind: Consts.CompletionItemKind.Constant,
textEdit: { range: marker9Range, newText: "'age'" },
},
],
},
marker10: {
completions: [],
},
marker11: {
completions: [
{
label: "'name'",
kind: Consts.CompletionItemKind.Constant,
textEdit: { range: marker11Range, newText: "'name'" },
},
{
label: "'age'",
kind: Consts.CompletionItemKind.Constant,
textEdit: { range: marker11Range, newText: "'age'" },
},
],
},
});
}

135
packages/pyright-internal/src/tests/fourslash/completions.call.typedDict.states.fourslash.ts Normal file
View File

@ -0,0 +1,135 @@
/// <reference path="fourslash.ts" />
// @filename: test.py
//// from typing import TypedDict
////
//// class Movie(TypedDict):
//// name: str
//// age: int
////
//// def thing(movie: Movie):
//// pass
////
//// thing(movie={'foo': 'a', '[|/*marker1*/|]'})
//// thing(movie={'foo': 'a', 'a[|/*marker2*/|]'})
//// thing(
//// movie={
//// 'name': 'Parasite',
//// '[|/*marker3*/|]
//// }
//// )
//// thing(
//// movie={
//// 'name': 'Parasite',
//// '[|/*marker4*/|]'
//// }
//// )
//// thing({
//// 'name': 'Parasite',
//// # hello world
//// '[|/*marker5*/|]'
//// })
//// thing({'foo': '[|/*marker6*/|]'})
{
// completions that rely on token parsing instead of node parsing
const marker1Range = helper.expandPositionRange(helper.getPositionRange('marker1'), 1, 1);
const marker2Range = helper.expandPositionRange(helper.getPositionRange('marker2'), 2, 1);
const marker3Range = helper.expandPositionRange(helper.getPositionRange('marker3'), 1, 0);
const marker4Range = helper.expandPositionRange(helper.getPositionRange('marker4'), 1, 1);
const marker5Range = helper.expandPositionRange(helper.getPositionRange('marker5'), 1, 1);
// @ts-ignore
await helper.verifyCompletion('exact', 'markdown', {
marker1: {
completions: [
{
label: "'name'",
kind: Consts.CompletionItemKind.Constant,
textEdit: { range: marker1Range, newText: "'name'" },
},
{
label: "'age'",
kind: Consts.CompletionItemKind.Constant,
textEdit: { range: marker1Range, newText: "'age'" },
},
],
},
marker2: {
completions: [
{
label: "'name'",
kind: Consts.CompletionItemKind.Constant,
textEdit: { range: marker2Range, newText: "'name'" },
},
{
label: "'age'",
kind: Consts.CompletionItemKind.Constant,
textEdit: { range: marker2Range, newText: "'age'" },
},
],
},
marker6: {
completions: [],
},
});
// @ts-ignore
await helper.verifyCompletion('included', 'markdown', {
marker3: {
completions: [
{
label: "'age'",
kind: Consts.CompletionItemKind.Constant,
textEdit: { range: marker3Range, newText: "'age'" },
},
],
},
marker4: {
completions: [
{
label: "'age'",
kind: Consts.CompletionItemKind.Constant,
textEdit: { range: marker4Range, newText: "'age'" },
},
],
},
marker5: {
completions: [
{
label: "'age'",
kind: Consts.CompletionItemKind.Constant,
textEdit: { range: marker5Range, newText: "'age'" },
},
],
},
});
// @ts-ignore
await helper.verifyCompletion('excluded', 'markdown', {
marker3: {
completions: [
{
label: "'name'",
kind: Consts.CompletionItemKind.Constant,
},
],
},
marker4: {
completions: [
{
label: "'name'",
kind: Consts.CompletionItemKind.Constant,
},
],
},
marker5: {
completions: [
{
label: "'name'",
kind: Consts.CompletionItemKind.Constant,
},
],
},
});
}