Initial support added for textDocument/documentHighlight

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

@@ -107,6 +107,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

@@ -92,6 +92,7 @@ public Server(IServiceManager services) {
                 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,71 @@
+// 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 enum DocumentHighlightSearchOptions {
+        All,
+        ExcludeLibraries
+    }
+
+    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) {
+                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 name = definingMember.GetName();
+
+                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;
+            }
+
+            return Array.Empty<DocumentHighlight>();
+        }
+    }
+}