Added Handler for Semantic Tokenization

src/PowerShellEditorServices/Services/TextDocument/Handlers/PsesSemanticTokens.cs

@@ -4,6 +4,7 @@
 //
 
 using System;
+using System.Collections.Generic;
 using System.Management.Automation.Language;
 using System.Threading;
 using System.Threading.Tasks;
@@ -18,13 +19,11 @@
 
 namespace Microsoft.PowerShell.EditorServices.Handlers
 {
-    //SemanticTokensHandler is labeled "Obsolete" because that is how Omnisharp marks proposed LSP features. Since we want this proposed feature, we disable this warning.
-#pragma warning disable 618
     internal class PsesSemanticTokens : SemanticTokensHandler
     {
         private readonly ILogger _logger;
         private readonly WorkspaceService _workspaceService;
-        static readonly SemanticTokensRegistrationOptions _registrationOptions = new SemanticTokensRegistrationOptions() {
+        private static readonly SemanticTokensRegistrationOptions s_registrationOptions = new SemanticTokensRegistrationOptions() {
             DocumentSelector = LspUtils.PowerShellDocumentSelector,
             Legend = new SemanticTokensLegend(),
             DocumentProvider = new Supports<SemanticTokensDocumentProviderOptions>(isSupported: true,
@@ -34,7 +33,7 @@ internal class PsesSemanticTokens : SemanticTokensHandler
             RangeProvider = true
         };
 
-        public PsesSemanticTokens(ILogger<PsesSemanticTokens> logger, WorkspaceService workspaceService) : base(_registrationOptions)
+        public PsesSemanticTokens(ILogger<PsesSemanticTokens> logger, WorkspaceService workspaceService) : base(s_registrationOptions)
         {
             _logger = logger;
             _workspaceService = workspaceService;
@@ -53,28 +52,44 @@ public PsesSemanticTokens(ILogger<PsesSemanticTokens> logger, WorkspaceService w
 
         private static void PushToken(Token token, SemanticTokensBuilder builder)
         {
+
+            List<SemanticToken> semanticToken = ConvertToSemanticTokens(token);
+            foreach(SemanticToken sToken in semanticToken)
+            {
+                builder.Push(sToken.Line, sToken.Index, length: sToken.Text.Length,
+                    sToken.Type, tokenModifiers: sToken.TokenModifiers);
+            }
+        }
+
+        internal static List<SemanticToken> ConvertToSemanticTokens(Token token){
             if(token is StringExpandableToken stringExpandableToken)
             {
                 // Try parsing tokens within the string
                 if (stringExpandableToken.NestedTokens != null)
                 {
+                    List<SemanticToken> tokens = new List<SemanticToken>();
                     foreach (Token t in stringExpandableToken.NestedTokens)
                     {
-                        PushToken(t, builder);
+                        tokens.AddRange(ConvertToSemanticTokens(t));
                     }
-                    return;
+                    return tokens;
                 }
             }
 
+            SemanticTokenType mappedType =  MapSemanticTokenType(token);
+            if(mappedType == null){
+                return new List<SemanticToken>();
+            }
+
             //Tokens line and col numbers indexed starting from 1, expecting indexing from 0
             int line = token.Extent.StartLineNumber - 1;
             int index = token.Extent.StartColumnNumber - 1;
-
-            builder.Push(line: line, @char: index, length: token.Text.Length,
-                tokenType: MapSemanticToken(token), tokenModifiers: Array.Empty<string>());
+            SemanticToken sToken = new SemanticToken(token.Text, mappedType,
+                line, index, Array.Empty<string>());
+            return new List<SemanticToken>{sToken};
         }
 
-        private static SemanticTokenType MapSemanticToken(Token token)
+        private static SemanticTokenType MapSemanticTokenType(Token token)
         {
             // First check token flags
             if ((token.TokenFlags & TokenFlags.Keyword) != 0)
@@ -130,8 +145,7 @@ private static SemanticTokenType MapSemanticToken(Token token)
                     return SemanticTokenType.Function;
             }
 
-            // Default semantic token
-            return SemanticTokenType.Documentation;
+            return null;
         }
 
         protected override Task<SemanticTokensDocument> GetSemanticTokensDocument(
@@ -141,5 +155,21 @@ private static SemanticTokenType MapSemanticToken(Token token)
             return Task.FromResult(new SemanticTokensDocument(GetRegistrationOptions().Legend));
         }
     }
-#pragma warning restore 618
+}
+
+class SemanticToken
+{
+    public SemanticToken(string text, SemanticTokenType type, int line, int index, IEnumerable<string> tokenModifiers)
+    {
+        Line = line;
+        Text = text;
+        Index = index;
+        Type = type;
+        TokenModifiers = tokenModifiers;
+    }
+    public string Text {get; set;}
+    public int Line {get; set;}
+    public int Index {get; set;}
+    public SemanticTokenType Type {get; set;}
+    public IEnumerable<string> TokenModifiers {get; set;}
 }

test/PowerShellEditorServices.Test.E2E/LanguageServerProtocolMessageTests.cs

@@ -1140,15 +1140,13 @@ await PsesLanguageClient
             Assert.Equal("Get-ChildItem", expandAliasResult.Text);
         }
 
-#pragma warning disable 618
-         [Fact]
+        [Fact]
         public async Task CanSendSemanticTokenRequest()
         {
-
-            //var tokens = semanticTokensBuilder.Commit().GetSemanticTokens();
             string scriptContent = "function";
             string scriptPath = NewTestFile(scriptContent);
-            var result =
+
+            SemanticTokens result =
                 await PsesLanguageClient
                     .SendRequest<SemanticTokensParams>(
                         "textDocument/semanticTokens",
@@ -1161,36 +1159,16 @@ await PsesLanguageClient
                         })
                     .Returning<SemanticTokens>(CancellationToken.None);
 
-            // Information about how this data is generated can be found at 0,0,11,14,0,0,12,0,0,0
+            // More information about how this data is generated can be found at
             // https://github.com/microsoft/vscode-extension-samples/blob/5ae1f7787122812dcc84e37427ca90af5ee09f14/semantic-tokens-sample/vscode.proposed.d.ts#L71
-
-
             var arr = new int[10]
                 {
                     //line, index, token length, token type, token modifiers
                     0, 0, scriptContent.Length, 2, 0, //token 1 is function
                     0, 9, 0, 0, 0 //token 2 is EOF
                 };
 
-
             Assert.Equal(result.Data.ToArray(), arr);
-
-            SemanticTokensRegistrationOptions options = new SemanticTokensRegistrationOptions() {
-                DocumentSelector = LspUtils.PowerShellDocumentSelector,
-                Legend = new SemanticTokensLegend(),
-                DocumentProvider = new Supports<SemanticTokensDocumentProviderOptions>(isSupported: true,
-                    new SemanticTokensDocumentProviderOptions {
-                        Edits = true
-                    }),
-                RangeProvider = true
-            };
-
-            SemanticTokensBuilder builder = new SemanticTokensBuilder(new SemanticTokensDocument(options), options.Legend);
-            builder.Push(0, 0, 8, SemanticTokenType.Keyword,  Array.Empty<string>());
-            SemanticTokens tokens = builder.Commit().GetSemanticTokens();
-
         }
-
-#pragma warning restore 618
     }
 }

test/PowerShellEditorServices.Test/Language/SemanticTokenTest.cs

@@ -0,0 +1,93 @@
+//
+// Copyright (c) Microsoft. All rights reserved.
+// Licensed under the MIT license. See LICENSE file in the project root for full license information.
+//
+
+using System;
+using System.Collections.Generic;
+using System.IO;
+using System.Management.Automation.Language;
+using System.Threading.Tasks;
+using Microsoft.PowerShell.EditorServices.Services.TextDocument;
+using Microsoft.PowerShell.EditorServices.Handlers;
+using OmniSharp.Extensions.LanguageServer.Protocol;
+using OmniSharp.Extensions.LanguageServer.Protocol.Document.Proposals;
+using OmniSharp.Extensions.LanguageServer.Protocol.Models.Proposals;
+using Xunit;
+
+namespace Microsoft.PowerShell.EditorServices.Test.Language
+{
+    public class SemanticTokenTest
+    {
+        [Fact]
+        public async Task MapsFunctionElements()
+        {
+            string text = @"
+function Get-Sum {
+    param( [int]$a, [int]$b )
+    return $a + $b
+}
+";
+            ScriptFile scriptFile = new ScriptFile(
+                // Use any absolute path. Even if it doesn't exist.
+                DocumentUri.FromFileSystemPath(Path.Combine(Path.GetTempPath(), "TestFile.ps1")),
+                text,
+                Version.Parse("5.0"));
+
+            foreach(Token t in scriptFile.ScriptTokens)
+            {
+                List<SemanticToken> mappedTokens = PsesSemanticTokens.ConvertToSemanticTokens(t);
+                switch(t.Text)
+                {
+                    case "function":
+                    case "param":
+                    case "return":
+                        Assert.Single(mappedTokens);
+                        Assert.Equal(SemanticTokenType.Keyword, mappedTokens[0].Type);
+                        break;
+                    case "Get-Sum":
+                        Assert.Single(mappedTokens);
+                        Assert.Equal(SemanticTokenType.Function, mappedTokens[0].Type);
+                        break;
+                    case "$a":
+                    case "$b":
+                        Assert.Single(mappedTokens);
+                        Assert.Equal(SemanticTokenType.Variable, mappedTokens[0].Type);
+                        break;
+                    case "[int]":
+                        Assert.Single(mappedTokens);
+                        Assert.Equal(SemanticTokenType.Type, mappedTokens[0].Type);
+                        break;
+                    case "+":
+                        Assert.Single(mappedTokens);
+                        Assert.Equal(SemanticTokenType.Operator, mappedTokens[0].Type);
+                        break;
+                }
+            }
+        }
+
+        [Fact]
+        public async Task TokenizesStringExpansion()
+        {
+            string text = "Write-Host \"$(Test-Property Get-Whatever) $(Get-Whatever)\"";
+            ScriptFile scriptFile = new ScriptFile(
+                // Use any absolute path. Even if it doesn't exist.
+                DocumentUri.FromFileSystemPath(Path.Combine(Path.GetTempPath(), "TestFile.ps1")),
+                text,
+                Version.Parse("5.0"));
+
+            Token commandToken = scriptFile.ScriptTokens[0];
+            List<SemanticToken> mappedTokens = PsesSemanticTokens.ConvertToSemanticTokens(commandToken);
+            Assert.Single(mappedTokens);
+            Assert.Equal(SemanticTokenType.Function, mappedTokens[0].Type);
+
+            Token stringExpandableToken = scriptFile.ScriptTokens[1];
+            mappedTokens = PsesSemanticTokens.ConvertToSemanticTokens(stringExpandableToken);
+            Assert.Collection(mappedTokens,
+                sToken => Assert.Equal(SemanticTokenType.Function, sToken.Type),
+                sToken => Assert.Equal(SemanticTokenType.Function, sToken.Type),
+                sToken => Assert.Equal(SemanticTokenType.Function, sToken.Type)
+            );
+        }
+    }
+}