Move Sorbet level and locate_first_within_range to RubyDocument (#2435)

vinistock · web-flow · commit 6c4702c17b90 · 2024-08-13T15:03:47.000-04:00
* Move sorbet_level to RubyDocument

* Move locate_first_within_range to RubyDocument
diff --git a/lib/ruby_lsp/document.rb b/lib/ruby_lsp/document.rb
@@ -10,16 +10,6 @@ class LanguageId < T::Enum
       end
     end
 
-    class SorbetLevel < T::Enum
-      enums do
-        None = new("none")
-        Ignore = new("ignore")
-        False = new("false")
-        True = new("true")
-        Strict = new("strict")
-      end
-    end
-
     extend T::Sig
     extend T::Helpers
 
@@ -223,60 +213,6 @@ def locate(node, char_position, node_types: [])
       NodeContext.new(closest, parent, nesting_nodes, call_node)
     end
 
-    sig do
-      params(
-        range: T::Hash[Symbol, T.untyped],
-        node_types: T::Array[T.class_of(Prism::Node)],
-      ).returns(T.nilable(Prism::Node))
-    end
-    def locate_first_within_range(range, node_types: [])
-      scanner = create_scanner
-      start_position = scanner.find_char_position(range[:start])
-      end_position = scanner.find_char_position(range[:end])
-      desired_range = (start_position...end_position)
-      queue = T.let(@parse_result.value.child_nodes.compact, T::Array[T.nilable(Prism::Node)])
-
-      until queue.empty?
-        candidate = queue.shift
-
-        # Skip nil child nodes
-        next if candidate.nil?
-
-        # Add the next child_nodes to the queue to be processed. The order here is important! We want to move in the
-        # same order as the visiting mechanism, which means searching the child nodes before moving on to the next
-        # sibling
-        T.unsafe(queue).unshift(*candidate.child_nodes)
-
-        # Skip if the current node doesn't cover the desired position
-        loc = candidate.location
-
-        if desired_range.cover?(loc.start_offset...loc.end_offset) &&
-            (node_types.empty? || node_types.any? { |type| candidate.class == type })
-          return candidate
-        end
-      end
-    end
-
-    sig { returns(SorbetLevel) }
-    def sorbet_level
-      sigil = parse_result.magic_comments.find do |comment|
-        comment.key == "typed"
-      end&.value
-
-      case sigil
-      when "ignore"
-        SorbetLevel::Ignore
-      when "false"
-        SorbetLevel::False
-      when "true"
-        SorbetLevel::True
-      when "strict", "strong"
-        SorbetLevel::Strict
-      else
-        SorbetLevel::None
-      end
-    end
-
     class Scanner
       extend T::Sig
 
diff --git a/lib/ruby_lsp/listeners/completion.rb b/lib/ruby_lsp/listeners/completion.rb
@@ -56,7 +56,7 @@ class Completion
           response_builder: ResponseBuilders::CollectionResponseBuilder[Interface::CompletionItem],
           global_state: GlobalState,
           node_context: NodeContext,
-          sorbet_level: Document::SorbetLevel,
+          sorbet_level: RubyDocument::SorbetLevel,
           dispatcher: Prism::Dispatcher,
           uri: URI::Generic,
           trigger_character: T.nilable(String),
@@ -99,7 +99,7 @@ def initialize( # rubocop:disable Metrics/ParameterLists
       def on_constant_read_node_enter(node)
         # The only scenario where Sorbet doesn't provide constant completion is on ignored files. Even if the file has
         # no sigil, Sorbet will still provide completion for constants
-        return if @sorbet_level != Document::SorbetLevel::Ignore
+        return if @sorbet_level != RubyDocument::SorbetLevel::Ignore
 
         name = constant_name(node)
         return if name.nil?
@@ -122,7 +122,7 @@ def on_constant_read_node_enter(node)
       def on_constant_path_node_enter(node)
         # The only scenario where Sorbet doesn't provide constant completion is on ignored files. Even if the file has
         # no sigil, Sorbet will still provide completion for constants
-        return if @sorbet_level != Document::SorbetLevel::Ignore
+        return if @sorbet_level != RubyDocument::SorbetLevel::Ignore
 
         name = constant_name(node)
         return if name.nil?
@@ -134,7 +134,7 @@ def on_constant_path_node_enter(node)
       def on_call_node_enter(node)
         # The only scenario where Sorbet doesn't provide constant completion is on ignored files. Even if the file has
         # no sigil, Sorbet will still provide completion for constants
-        if @sorbet_level == Document::SorbetLevel::Ignore
+        if @sorbet_level == RubyDocument::SorbetLevel::Ignore
           receiver = node.receiver
 
           # When writing `Foo::`, the AST assigns a method call node (because you can use that syntax to invoke
@@ -257,7 +257,7 @@ def constant_path_completion(name, range)
       def handle_instance_variable_completion(name, location)
         # Sorbet enforces that all instance variables be declared on typed strict or higher, which means it will be able
         # to provide all features for them
-        return if @sorbet_level == Document::SorbetLevel::Strict
+        return if @sorbet_level == RubyDocument::SorbetLevel::Strict
 
         type = @type_inferrer.infer_receiver_type(@node_context)
         return unless type
diff --git a/lib/ruby_lsp/listeners/definition.rb b/lib/ruby_lsp/listeners/definition.rb
@@ -20,7 +20,7 @@ class Definition
           uri: URI::Generic,
           node_context: NodeContext,
           dispatcher: Prism::Dispatcher,
-          sorbet_level: Document::SorbetLevel,
+          sorbet_level: RubyDocument::SorbetLevel,
         ).void
       end
       def initialize(response_builder, global_state, language_id, uri, node_context, dispatcher, sorbet_level) # rubocop:disable Metrics/ParameterLists
@@ -181,7 +181,7 @@ def handle_super_node_definition
       def handle_instance_variable_definition(name)
         # Sorbet enforces that all instance variables be declared on typed strict or higher, which means it will be able
         # to provide all features for them
-        return if @sorbet_level == Document::SorbetLevel::Strict
+        return if @sorbet_level == RubyDocument::SorbetLevel::Strict
 
         type = @type_inferrer.infer_receiver_type(@node_context)
         return unless type
@@ -289,7 +289,7 @@ def find_in_index(value)
           # additional behavior on top of jumping to RBIs. The only sigil where Sorbet cannot handle constants is typed
           # ignore
           file_path = entry.file_path
-          next if @sorbet_level != Document::SorbetLevel::Ignore && not_in_dependencies?(file_path)
+          next if @sorbet_level != RubyDocument::SorbetLevel::Ignore && not_in_dependencies?(file_path)
 
           @response_builder << Interface::LocationLink.new(
             target_uri: URI::Generic.from_path(path: file_path).to_s,
diff --git a/lib/ruby_lsp/listeners/hover.rb b/lib/ruby_lsp/listeners/hover.rb
@@ -42,7 +42,7 @@ class Hover
           uri: URI::Generic,
           node_context: NodeContext,
           dispatcher: Prism::Dispatcher,
-          sorbet_level: Document::SorbetLevel,
+          sorbet_level: RubyDocument::SorbetLevel,
         ).void
       end
       def initialize(response_builder, global_state, uri, node_context, dispatcher, sorbet_level) # rubocop:disable Metrics/ParameterLists
@@ -73,7 +73,7 @@ def initialize(response_builder, global_state, uri, node_context, dispatcher, so
 
       sig { params(node: Prism::ConstantReadNode).void }
       def on_constant_read_node_enter(node)
-        return if @sorbet_level != Document::SorbetLevel::Ignore
+        return if @sorbet_level != RubyDocument::SorbetLevel::Ignore
 
         name = constant_name(node)
         return if name.nil?
@@ -83,14 +83,14 @@ def on_constant_read_node_enter(node)
 
       sig { params(node: Prism::ConstantWriteNode).void }
       def on_constant_write_node_enter(node)
-        return if @sorbet_level != Document::SorbetLevel::Ignore
+        return if @sorbet_level != RubyDocument::SorbetLevel::Ignore
 
         generate_hover(node.name.to_s, node.name_loc)
       end
 
       sig { params(node: Prism::ConstantPathNode).void }
       def on_constant_path_node_enter(node)
-        return if @sorbet_level != Document::SorbetLevel::Ignore
+        return if @sorbet_level != RubyDocument::SorbetLevel::Ignore
 
         name = constant_name(node)
         return if name.nil?
@@ -193,7 +193,7 @@ def handle_method_hover(message, inherited_only: false)
       def handle_instance_variable_hover(name)
         # Sorbet enforces that all instance variables be declared on typed strict or higher, which means it will be able
         # to provide all features for them
-        return if @sorbet_level == Document::SorbetLevel::Strict
+        return if @sorbet_level == RubyDocument::SorbetLevel::Strict
 
         type = @type_inferrer.infer_receiver_type(@node_context)
         return unless type
diff --git a/lib/ruby_lsp/listeners/signature_help.rb b/lib/ruby_lsp/listeners/signature_help.rb
@@ -13,7 +13,7 @@ class SignatureHelp
           global_state: GlobalState,
           node_context: NodeContext,
           dispatcher: Prism::Dispatcher,
-          sorbet_level: Document::SorbetLevel,
+          sorbet_level: RubyDocument::SorbetLevel,
         ).void
       end
       def initialize(response_builder, global_state, node_context, dispatcher, sorbet_level)
diff --git a/lib/ruby_lsp/requests/code_action_resolve.rb b/lib/ruby_lsp/requests/code_action_resolve.rb
@@ -38,7 +38,7 @@ class Error < ::T::Enum
         end
       end
 
-      sig { params(document: Document, code_action: T::Hash[Symbol, T.untyped]).void }
+      sig { params(document: RubyDocument, code_action: T::Hash[Symbol, T.untyped]).void }
       def initialize(document, code_action)
         super()
         @document = document
diff --git a/lib/ruby_lsp/requests/completion.rb b/lib/ruby_lsp/requests/completion.rb
@@ -49,7 +49,7 @@ def provider
           document: Document,
           global_state: GlobalState,
           params: T::Hash[Symbol, T.untyped],
-          sorbet_level: Document::SorbetLevel,
+          sorbet_level: RubyDocument::SorbetLevel,
           dispatcher: Prism::Dispatcher,
         ).void
       end
diff --git a/lib/ruby_lsp/requests/definition.rb b/lib/ruby_lsp/requests/definition.rb
@@ -42,7 +42,7 @@ class Definition < Request
           global_state: GlobalState,
           position: T::Hash[Symbol, T.untyped],
           dispatcher: Prism::Dispatcher,
-          sorbet_level: Document::SorbetLevel,
+          sorbet_level: RubyDocument::SorbetLevel,
         ).void
       end
       def initialize(document, global_state, position, dispatcher, sorbet_level)
diff --git a/lib/ruby_lsp/requests/hover.rb b/lib/ruby_lsp/requests/hover.rb
@@ -36,7 +36,7 @@ def provider
           global_state: GlobalState,
           position: T::Hash[Symbol, T.untyped],
           dispatcher: Prism::Dispatcher,
-          sorbet_level: Document::SorbetLevel,
+          sorbet_level: RubyDocument::SorbetLevel,
         ).void
       end
       def initialize(document, global_state, position, dispatcher, sorbet_level)
diff --git a/lib/ruby_lsp/requests/signature_help.rb b/lib/ruby_lsp/requests/signature_help.rb
@@ -46,7 +46,7 @@ def provider
           position: T::Hash[Symbol, T.untyped],
           context: T.nilable(T::Hash[Symbol, T.untyped]),
           dispatcher: Prism::Dispatcher,
-          sorbet_level: Document::SorbetLevel,
+          sorbet_level: RubyDocument::SorbetLevel,
         ).void
       end
       def initialize(document, global_state, position, context, dispatcher, sorbet_level) # rubocop:disable Metrics/ParameterLists
diff --git a/lib/ruby_lsp/requests/support/common.rb b/lib/ruby_lsp/requests/support/common.rb
@@ -209,9 +209,9 @@ def kind_for_entry(entry)
           end
         end
 
-        sig { params(sorbet_level: Document::SorbetLevel).returns(T::Boolean) }
+        sig { params(sorbet_level: RubyDocument::SorbetLevel).returns(T::Boolean) }
         def sorbet_level_true_or_higher?(sorbet_level)
-          sorbet_level == Document::SorbetLevel::True || sorbet_level == Document::SorbetLevel::Strict
+          sorbet_level == RubyDocument::SorbetLevel::True || sorbet_level == RubyDocument::SorbetLevel::Strict
         end
       end
     end
diff --git a/lib/ruby_lsp/ruby_document.rb b/lib/ruby_lsp/ruby_document.rb
@@ -3,6 +3,16 @@
 
 module RubyLsp
   class RubyDocument < Document
+    class SorbetLevel < T::Enum
+      enums do
+        None = new("none")
+        Ignore = new("ignore")
+        False = new("false")
+        True = new("true")
+        Strict = new("strict")
+      end
+    end
+
     sig { override.returns(Prism::ParseResult) }
     def parse
       return @parse_result unless @needs_parsing
@@ -20,5 +30,59 @@ def syntax_error?
     def language_id
       LanguageId::Ruby
     end
+
+    sig { returns(SorbetLevel) }
+    def sorbet_level
+      sigil = parse_result.magic_comments.find do |comment|
+        comment.key == "typed"
+      end&.value
+
+      case sigil
+      when "ignore"
+        SorbetLevel::Ignore
+      when "false"
+        SorbetLevel::False
+      when "true"
+        SorbetLevel::True
+      when "strict", "strong"
+        SorbetLevel::Strict
+      else
+        SorbetLevel::None
+      end
+    end
+
+    sig do
+      params(
+        range: T::Hash[Symbol, T.untyped],
+        node_types: T::Array[T.class_of(Prism::Node)],
+      ).returns(T.nilable(Prism::Node))
+    end
+    def locate_first_within_range(range, node_types: [])
+      scanner = create_scanner
+      start_position = scanner.find_char_position(range[:start])
+      end_position = scanner.find_char_position(range[:end])
+      desired_range = (start_position...end_position)
+      queue = T.let(@parse_result.value.child_nodes.compact, T::Array[T.nilable(Prism::Node)])
+
+      until queue.empty?
+        candidate = queue.shift
+
+        # Skip nil child nodes
+        next if candidate.nil?
+
+        # Add the next child_nodes to the queue to be processed. The order here is important! We want to move in the
+        # same order as the visiting mechanism, which means searching the child nodes before moving on to the next
+        # sibling
+        T.unsafe(queue).unshift(*candidate.child_nodes)
+
+        # Skip if the current node doesn't cover the desired position
+        loc = candidate.location
+
+        if desired_range.cover?(loc.start_offset...loc.end_offset) &&
+            (node_types.empty? || node_types.any? { |type| candidate.class == type })
+          return candidate
+        end
+      end
+    end
   end
 end
diff --git a/lib/ruby_lsp/server.rb b/lib/ruby_lsp/server.rb
@@ -495,9 +495,10 @@ def text_document_hover(message)
       )
     end
 
-    sig { params(document: Document).returns(Document::SorbetLevel) }
+    sig { params(document: Document).returns(RubyDocument::SorbetLevel) }
     def sorbet_level(document)
-      return Document::SorbetLevel::Ignore unless @global_state.has_type_checker
+      return RubyDocument::SorbetLevel::Ignore unless @global_state.has_type_checker
+      return RubyDocument::SorbetLevel::Ignore unless document.is_a?(RubyDocument)
 
       document.sorbet_level
     end
@@ -535,6 +536,12 @@ def code_action_resolve(message)
       params = message[:params]
       uri = URI(params.dig(:data, :uri))
       document = @store.get(uri)
+
+      unless document.is_a?(RubyDocument)
+        send_message(Notification.window_show_error("Code actions are currently only available for Ruby documents"))
+        raise Requests::CodeActionResolve::CodeActionError
+      end
+
       result = Requests::CodeActionResolve.new(document, params).perform
 
       case result
diff --git a/test/ruby_document_test.rb b/test/ruby_document_test.rb
