Move sorbet_level to RubyDocument

Author: vinistock

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
 
@@ -257,26 +247,6 @@ def locate_first_within_range(range, node_types: [])
       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
 

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

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,

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

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)

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

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)

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)

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

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

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,25 @@ 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
   end
 end

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

test/ruby_document_test.rb

@@ -601,32 +601,32 @@ def test_no_sigil
       version: 1,
       uri: URI("file:///foo/bar.rb"),
     )
-    assert_equal(RubyLsp::Document::SorbetLevel::None, document.sorbet_level)
+    assert_equal(RubyLsp::RubyDocument::SorbetLevel::None, document.sorbet_level)
   end
 
   def test_sigil_ignore
     document = RubyLsp::RubyDocument.new(source: +"# typed: ignore", version: 1, uri: URI("file:///foo/bar.rb"))
-    assert_equal(RubyLsp::Document::SorbetLevel::Ignore, document.sorbet_level)
+    assert_equal(RubyLsp::RubyDocument::SorbetLevel::Ignore, document.sorbet_level)
   end
 
   def test_sigil_false
     document = RubyLsp::RubyDocument.new(source: +"# typed: false", version: 1, uri: URI("file:///foo/bar.rb"))
-    assert_equal(RubyLsp::Document::SorbetLevel::False, document.sorbet_level)
+    assert_equal(RubyLsp::RubyDocument::SorbetLevel::False, document.sorbet_level)
   end
 
   def test_sigil_true
     document = RubyLsp::RubyDocument.new(source: +"# typed: true", version: 1, uri: URI("file:///foo/bar.rb"))
-    assert_equal(RubyLsp::Document::SorbetLevel::True, document.sorbet_level)
+    assert_equal(RubyLsp::RubyDocument::SorbetLevel::True, document.sorbet_level)
   end
 
   def test_sigil_strict
     document = RubyLsp::RubyDocument.new(source: +"# typed: strict", version: 1, uri: URI("file:///foo/bar.rb"))
-    assert_equal(RubyLsp::Document::SorbetLevel::Strict, document.sorbet_level)
+    assert_equal(RubyLsp::RubyDocument::SorbetLevel::Strict, document.sorbet_level)
   end
 
   def test_sigil_strong
     document = RubyLsp::RubyDocument.new(source: +"# typed: strong", version: 1, uri: URI("file:///foo/bar.rb"))
-    assert_equal(RubyLsp::Document::SorbetLevel::Strict, document.sorbet_level)
+    assert_equal(RubyLsp::RubyDocument::SorbetLevel::Strict, document.sorbet_level)
   end
 
   def test_sorbet_sigil_only_in_magic_comment
@@ -651,7 +651,7 @@ def baz
         CODE
       end
     RUBY
-    assert_equal(RubyLsp::Document::SorbetLevel::False, document.sorbet_level)
+    assert_equal(RubyLsp::RubyDocument::SorbetLevel::False, document.sorbet_level)
   end
 
   def test_locating_compact_namespace_declaration