diff --git a/src/features/semanticTokens.ts b/src/features/semanticTokens.ts index 85239dd1..81a87534 100644 --- a/src/features/semanticTokens.ts +++ b/src/features/semanticTokens.ts @@ -1,4 +1,4 @@ -import { convertOffsetsToRange } from '@zzzen/pyright-internal/dist/common/positionUtils'; +import { convertOffsetsToRange, convertTextRangeToRange } from '@zzzen/pyright-internal/dist/common/positionUtils'; import { CancellationToken, DocumentSemanticTokensProvider, @@ -19,6 +19,7 @@ const tokenTypes: string[] = [ SemanticTokenTypes.enum, SemanticTokenTypes.enumMember, SemanticTokenTypes.function, + SemanticTokenTypes.keyword, SemanticTokenTypes.method, SemanticTokenTypes.namespace, SemanticTokenTypes.parameter, @@ -58,6 +59,14 @@ export class PythonSemanticTokensProvider implements DocumentSemanticTokensProvi if (token && token.isCancellationRequested) return null; const builder = new SemanticTokensBuilder(this.legend); + // @ts-ignore + for (const item of parsed.tokenizerOutput.tokens._items) { + if (item.type === 8 && item.keywordType) { + const range = convertTextRangeToRange(item, parsed.tokenizerOutput.lines); + builder.push(range.start.line, range.start.character, item.length, encodeTokenType(SemanticTokenTypes.keyword)); + } + } + const walker = new SemanticTokensWalker(); walker.walk(parsed.parseTree); diff --git a/src/parsers/semanticTokens.ts b/src/parsers/semanticTokens.ts index 8f2e18b6..48b537ec 100644 --- a/src/parsers/semanticTokens.ts +++ b/src/parsers/semanticTokens.ts @@ -138,13 +138,12 @@ export class SemanticTokensWalker extends ParseTreeWalker { modifiers.push(SemanticTokenModifiers.async); } const type = node.parent?.parent?.nodeType === 10 ? SemanticTokenTypes.method : SemanticTokenTypes.function; - this.addItem(node.name, type); + this.addItem(node.name, type, modifiers); for (const p of node.parameters) { if (!p.name) continue; - const modifiers = [SemanticTokenModifiers.declaration]; - this.addItem(p.name, SemanticTokenTypes.parameter, modifiers); + this.addItem(p.name, SemanticTokenTypes.parameter); if (p.typeAnnotation) { this.addItem(p.typeAnnotation, SemanticTokenTypes.typeParameter); }