Initial support added for textDocument/documentHighlight (#1767)

* Initial support added for textDocument/documentHighlight

Change-Id: Ib13d86bc96a3702b0e0d79b27b7791898388e104

* Added two tests for DocumentHighlightSource

Change-Id: I411c1d4daac84a6a12a95b11fa5781eaf42eeae9

* Code refactoring according to comments from reviewer

Change-Id: If45beeafc9a40af5ac12b79d124c8021e9e492a9

Co-authored-by: Anders Stenman <[email protected]>
Co-authored-by: Jake Bailey <[email protected]>
(cherry picked from commit 1a9d6e5)

Author: 3 people

src/LanguageServer/Impl/Implementation/Server.Editor.cs

@@ -108,6 +108,12 @@ public Task<Reference[]> FindReferences(ReferencesParams @params, CancellationTo
             return new ReferenceSource(Services).FindAllReferencesAsync(uri, @params.position, ReferenceSearchOptions.All, cancellationToken);
         }
 
+        public Task<DocumentHighlight[]> DocumentHighlight(ReferencesParams @params, CancellationToken cancellationToken) {
+            var uri = @params.textDocument.uri;
+            _log?.Log(TraceEventType.Verbose, $"Document highlight in {uri} at {@params.position}");
+            return new DocumentHighlightSource(Services).DocumentHighlightAsync(uri, @params.position, cancellationToken);
+        }
+
         public Task<WorkspaceEdit> Rename(RenameParams @params, CancellationToken cancellationToken) {
             var uri = @params.textDocument.uri;
             _log?.Log(TraceEventType.Verbose, $"Rename in {uri} at {@params.position}");

src/LanguageServer/Impl/Implementation/Server.cs

@@ -98,6 +98,7 @@ private InitializeResult GetInitializeResult() {
                     referencesProvider = true,
                     workspaceSymbolProvider = true,
                     documentSymbolProvider = true,
+                    documentHighlightProvider = true,
                     renameProvider = true,
                     declarationProvider = true,
                     documentOnTypeFormattingProvider = new DocumentOnTypeFormattingOptions {

src/LanguageServer/Impl/LanguageServer.cs

@@ -211,11 +211,14 @@ public async Task<Reference[]> FindReferences(JToken token, CancellationToken ca
             }
         }
 
-        //[JsonRpcMethod("textDocument/documentHighlight")]
-        //public async Task<DocumentHighlight[]> DocumentHighlight(JToken token, CancellationToken cancellationToken) {
-        //    await _prioritizer.DefaultPriorityAsync(cancellationToken);
-        //    return await _server.DocumentHighlight(ToObject<TextDocumentPositionParams>(token), cancellationToken);
-        //}
+        [JsonRpcMethod("textDocument/documentHighlight")]
+        public async Task<DocumentHighlight[]> DocumentHighlight(JToken token, CancellationToken cancellationToken) {
+            using (_requestTimer.Time("textDocument/documentHighlight")) {
+                await _prioritizer.DefaultPriorityAsync(cancellationToken);
+                Debug.Assert(_initialized);
+                return await _server.DocumentHighlight(ToObject<ReferencesParams>(token), cancellationToken);
+            }
+        }
 
         [JsonRpcMethod("textDocument/documentSymbol")]
         public async Task<DocumentSymbol[]> DocumentSymbol(JToken token, CancellationToken cancellationToken) {

src/LanguageServer/Impl/Sources/DocumentHighlightSource.cs

@@ -0,0 +1,64 @@
+// Copyright(c) Microsoft Corporation
+// All rights reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the License); you may not use
+// this file except in compliance with the License. You may obtain a copy of the
+// License at http://www.apache.org/licenses/LICENSE-2.0
+//
+// THIS CODE IS PROVIDED ON AN  *AS IS* BASIS, WITHOUT WARRANTIES OR CONDITIONS
+// OF ANY KIND, EITHER EXPRESS OR IMPLIED, INCLUDING WITHOUT LIMITATION ANY
+// IMPLIED WARRANTIES OR CONDITIONS OF TITLE, FITNESS FOR A PARTICULAR PURPOSE,
+// MERCHANTABILITY OR NON-INFRINGEMENT.
+//
+// See the Apache Version 2.0 License for specific language governing
+// permissions and limitations under the License.
+
+using System;
+using System.Collections.Generic;
+using System.IO;
+using System.Linq;
+using System.Threading;
+using System.Threading.Tasks;
+using Microsoft.Python.Analysis;
+using Microsoft.Python.Analysis.Documents;
+using Microsoft.Python.Analysis.Modules;
+using Microsoft.Python.Analysis.Types;
+using Microsoft.Python.Core;
+using Microsoft.Python.Core.IO;
+using Microsoft.Python.Core.Text;
+using Microsoft.Python.LanguageServer.Documents;
+using Microsoft.Python.LanguageServer.Protocol;
+
+namespace Microsoft.Python.LanguageServer.Sources {
+    internal sealed class DocumentHighlightSource {
+        private const int DocumentHighlightAnalysisTimeout = 10000;
+        private readonly IServiceContainer _services;
+
+        public DocumentHighlightSource(IServiceContainer services) {
+            _services = services;
+        }
+
+        public async Task<DocumentHighlight[]> DocumentHighlightAsync(Uri uri, SourceLocation location, CancellationToken cancellationToken = default) {
+            if (uri == null) {
+                return Array.Empty<DocumentHighlight>();
+            }
+
+            var analysis = await Document.GetAnalysisAsync(uri, _services, DocumentHighlightAnalysisTimeout, cancellationToken);
+            var definitionSource = new DefinitionSource(_services);
+
+            var definition = definitionSource.FindDefinition(analysis, location, out var definingMember);
+            if (definition == null || definingMember == null) {
+                return Array.Empty<DocumentHighlight>();
+            }
+
+            var rootDefinition = definingMember.GetRootDefinition();
+
+            var result = rootDefinition.References
+                .Where(r => r.DocumentUri.Equals(uri))
+                .Select((r, i) => new DocumentHighlight { kind = (i == 0) ? DocumentHighlightKind.Write : DocumentHighlightKind.Read, range = r.Span })
+                .ToArray();
+
+            return result;
+        }
+    }
+}

src/LanguageServer/Test/DocumentHighlightTests.cs

@@ -0,0 +1,85 @@
+// Copyright(c) Microsoft Corporation
+// All rights reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the License); you may not use
+// this file except in compliance with the License. You may obtain a copy of the
+// License at http://www.apache.org/licenses/LICENSE-2.0
+//
+// THIS CODE IS PROVIDED ON AN  *AS IS* BASIS, WITHOUT WARRANTIES OR CONDITIONS
+// OF ANY KIND, EITHER EXPRESS OR IMPLIED, INCLUDING WITHOUT LIMITATION ANY
+// IMPLIED WARRANTIES OR CONDITIONS OF TITLE, FITNESS FOR A PARTICULAR PURPOSE,
+// MERCHANTABILITY OR NON-INFRINGEMENT.
+//
+// See the Apache Version 2.0 License for specific language governing
+// permissions and limitations under the License.
+
+using System;
+using System.IO;
+using System.Threading.Tasks;
+using FluentAssertions;
+using Microsoft.Python.Analysis.Analyzer;
+using Microsoft.Python.Analysis.Documents;
+using Microsoft.Python.Core.Text;
+using Microsoft.Python.LanguageServer.Protocol;
+using Microsoft.Python.LanguageServer.Sources;
+using Microsoft.Python.LanguageServer.Tests.FluentAssertions;
+using Microsoft.Python.Parsing.Tests;
+using Microsoft.VisualStudio.TestTools.UnitTesting;
+using TestUtilities;
+
+namespace Microsoft.Python.LanguageServer.Tests {
+    [TestClass]
+    public class DocumentHighlightTests : LanguageServerTestBase {
+        public TestContext TestContext { get; set; }
+
+        [TestInitialize]
+        public void TestInitialize()
+            => TestEnvironmentImpl.TestInitialize($"{TestContext.FullyQualifiedTestClassName}.{TestContext.TestName}");
+
+        [TestCleanup]
+        public void Cleanup() => TestEnvironmentImpl.TestCleanup();
+
+
+        [TestMethod, Priority(0)]
+        public async Task HighlightBasic() {
+            const string code = @"
+x = 1
+
+def func(x):
+    return x
+
+y = func(x)
+x = 2
+";
+            var analysis = await GetAnalysisAsync(code);
+            var dhs = new DocumentHighlightSource(Services);
+
+            // Test global scope
+            var highlights1 = await dhs.DocumentHighlightAsync(analysis.Document.Uri, new SourceLocation(8, 1));
+
+            highlights1.Should().HaveCount(3);
+            highlights1[0].range.Should().Be(1, 0, 1, 1);
+            highlights1[0].kind.Should().Be(DocumentHighlightKind.Write);
+            highlights1[1].range.Should().Be(6, 9, 6, 10);
+            highlights1[1].kind.Should().Be(DocumentHighlightKind.Read);
+            highlights1[2].range.Should().Be(7, 0, 7, 1);
+
+            // Test local scope in func()
+            var highlights2 = await dhs.DocumentHighlightAsync(analysis.Document.Uri, new SourceLocation(4, 10));
+
+            highlights2.Should().HaveCount(2);
+            highlights2[0].range.Should().Be(3, 9, 3, 10);
+            highlights2[0].kind.Should().Be(DocumentHighlightKind.Write);
+            highlights2[1].range.Should().Be(4, 11, 4, 12);
+            highlights2[1].kind.Should().Be(DocumentHighlightKind.Read);
+        }
+
+        [TestMethod, Priority(0)]
+        public async Task HighlightEmptyDocument() {
+            await GetAnalysisAsync(string.Empty);
+            var dhs = new DocumentHighlightSource(Services);
+            var references = await dhs.DocumentHighlightAsync(null, new SourceLocation(1, 1));
+            references.Should().BeEmpty();
+        }
+    }
+}