Make document symbol extensible (#949)

* Make DocumentSymbol extensible

* Add test_extension helper for extension tests

lib/ruby_lsp/executor.rb

@@ -104,6 +104,7 @@ def run(request)
         emitter.visit(document.tree) if document.parsed?
 
         code_lens.merge_external_listeners_responses!
+        document_symbol.merge_external_listeners_responses!
 
         # Store all responses retrieve in this round of visits in the cache and then return the response for the request
         # we actually received

lib/ruby_lsp/extension.rb

@@ -124,5 +124,14 @@ def create_code_lens_listener(uri, emitter, message_queue); end
       ).returns(T.nilable(Listener[T.nilable(Interface::Hover)]))
     end
     def create_hover_listener(emitter, message_queue); end
+
+    # Creates a new DocumentSymbol listener. This method is invoked on every DocumentSymbol request
+    sig do
+      overridable.params(
+        emitter: EventEmitter,
+        message_queue: Thread::Queue,
+      ).returns(T.nilable(Listener[T.nilable(Interface::DocumentSymbol)]))
+    end
+    def create_document_symbol_listener(emitter, message_queue); end
   end
 end

lib/ruby_lsp/requests/document_symbol.rb

@@ -60,6 +60,10 @@ def initialize(emitter, message_queue)
           T::Array[T.any(SymbolHierarchyRoot, Interface::DocumentSymbol)],
         )
 
+        @external_listeners.concat(
+          Extension.extensions.filter_map { |ext| ext.create_document_symbol_listener(emitter, message_queue) },
+        )
+
         emitter.register(
           self,
           :on_class,
@@ -75,6 +79,13 @@ def initialize(emitter, message_queue)
         )
       end
 
+      # Merges responses from other listeners
+      sig { override.params(other: Listener[ResponseType]).returns(T.self_type) }
+      def merge_response!(other)
+        @response.concat(other.response)
+        self
+      end
+
       sig { params(node: SyntaxTree::ClassDeclaration).void }
       def on_class(node)
         @stack << create_document_symbol(

test/expectations/expectations_test_runner.rb

@@ -105,6 +105,26 @@ def ruby_requirement_magic_comment_version(fixture_path)
 
   private
 
+  def test_extension(extension_creation_method, source:, &block)
+    RubyLsp::DependencyDetector.const_set(:HAS_TYPECHECKER, false)
+    message_queue = Thread::Queue.new
+
+    send(extension_creation_method)
+
+    store = RubyLsp::Store.new
+    uri = URI::Generic.from_path(path: "/fake.rb")
+    store.set(uri: uri, source: source, version: 1)
+
+    executor = RubyLsp::Executor.new(store, message_queue)
+    executor.instance_variable_get(:@index).index_single(uri.to_standardized_path, source)
+
+    yield(executor)
+  ensure
+    RubyLsp::Extension.extensions.clear
+    RubyLsp::DependencyDetector.const_set(:HAS_TYPECHECKER, true)
+    T.must(message_queue).close
+  end
+
   def diff(expected, actual)
     res = super
     return unless res

test/requests/code_lens_expectations_test.rb

@@ -98,27 +98,22 @@ def test_bar; end
   end
 
   def test_code_lens_extensions
-    message_queue = Thread::Queue.new
-    create_code_lens_extension
-
-    store = RubyLsp::Store.new
-    store.set(uri: URI("file:///fake.rb"), source: <<~RUBY, version: 1)
+    source = <<~RUBY
       class Test < Minitest::Test; end
     RUBY
 
-    response = RubyLsp::Executor.new(store, message_queue).execute({
-      method: "textDocument/codeLens",
-      params: { textDocument: { uri: "file:///fake.rb" }, position: { line: 1, character: 2 } },
-    }).response
-
-    assert_equal(response.size, 4)
-    assert_match("Run", response[0].command.title)
-    assert_match("Run In Terminal", response[1].command.title)
-    assert_match("Debug", response[2].command.title)
-    assert_match("Run Test", response[3].command.title)
-  ensure
-    RubyLsp::Extension.extensions.clear
-    T.must(message_queue).close
+    test_extension(:create_code_lens_extension, source: source) do |executor|
+      response = executor.execute({
+        method: "textDocument/codeLens",
+        params: { textDocument: { uri: "file:///fake.rb" }, position: { line: 1, character: 2 } },
+      }).response
+
+      assert_equal(response.size, 4)
+      assert_match("Run", response[0].command.title)
+      assert_match("Run In Terminal", response[1].command.title)
+      assert_match("Debug", response[2].command.title)
+      assert_match("Run Test", response[3].command.title)
+    end
   end
 
   private

test/requests/document_symbol_expectations_test.rb

@@ -1,9 +1,67 @@
-# typed: strict
+# typed: true
 # frozen_string_literal: true
 
 require "test_helper"
 require "expectations/expectations_test_runner"
 
 class DocumentSymbolExpectationsTest < ExpectationsTestRunner
   expectations_tests RubyLsp::Requests::DocumentSymbol, "document_symbol"
+
+  def test_document_symbol_extensions
+    source = <<~RUBY
+      test "foo" do
+      end
+    RUBY
+
+    test_extension(:create_document_symbol_extension, source: source) do |executor|
+      response = executor.execute({
+        method: "textDocument/documentSymbol",
+        params: { textDocument: { uri: "file:///fake.rb" }, position: { line: 0, character: 1 } },
+      }).response
+
+      assert_equal("foo", response.first.name)
+      assert_equal(LanguageServer::Protocol::Constant::SymbolKind::METHOD, response.first.kind)
+    end
+  end
+
+  private
+
+  def create_document_symbol_extension
+    Class.new(RubyLsp::Extension) do
+      def activate; end
+
+      def name
+        "Document SymbolsExtension"
+      end
+
+      def create_document_symbol_listener(emitter, message_queue)
+        klass = Class.new(RubyLsp::Listener) do
+          attr_reader :response
+
+          def initialize(emitter, message_queue)
+            super
+            emitter.register(self, :on_command)
+          end
+
+          def on_command(node)
+            T.bind(self, RubyLsp::Listener[T.untyped])
+            message_value = node.message.value
+            return unless message_value == "test" && node.arguments.parts.any?
+
+            first_argument = node.arguments.parts.first
+            test_name = first_argument.parts.map(&:value).join
+
+            @response = [RubyLsp::Interface::DocumentSymbol.new(
+              name: test_name,
+              kind: LanguageServer::Protocol::Constant::SymbolKind::METHOD,
+              selection_range: range_from_syntax_tree_node(node),
+              range: range_from_syntax_tree_node(node),
+            )]
+          end
+        end
+
+        klass.new(emitter, message_queue)
+      end
+    end
+  end
 end

test/requests/hover_expectations_test.rb

@@ -33,34 +33,22 @@ def run_expectations(source)
   end
 
   def test_hover_extensions
-    RubyLsp::DependencyDetector.const_set(:HAS_TYPECHECKER, false)
-    message_queue = Thread::Queue.new
-    create_hover_extension
-
-    store = RubyLsp::Store.new
     source = <<~RUBY
       # Hello
       class Post
       end
 
       Post
     RUBY
-    uri = URI::Generic.from_path(path: "/fake.rb")
-    store.set(uri: uri, source: source, version: 1)
 
-    executor = RubyLsp::Executor.new(store, message_queue)
-    executor.instance_variable_get(:@index).index_single(uri.to_standardized_path, source)
-
-    response = executor.execute({
-      method: "textDocument/hover",
-      params: { textDocument: { uri: "file:///fake.rb" }, position: { line: 4, character: 0 } },
-    }).response
-
-    assert_match("Hello\n\nHello from middleware: Post", response.contents.value)
-  ensure
-    RubyLsp::Extension.extensions.clear
-    RubyLsp::DependencyDetector.const_set(:HAS_TYPECHECKER, true)
-    T.must(message_queue).close
+    test_extension(:create_hover_extension, source: source) do |executor|
+      response = executor.execute({
+        method: "textDocument/hover",
+        params: { textDocument: { uri: "file:///fake.rb" }, position: { line: 4, character: 0 } },
+      }).response
+
+      assert_match("Hello\n\nHello from middleware: Post", response.contents.value)
+    end
   end
 
   private