Add hover support for :through associations

Author: joshuay03

lib/ruby_lsp/ruby_lsp_rails/hover.rb

@@ -29,6 +29,7 @@ def initialize(client, response_builder, node_context, global_state, dispatcher)
           :on_constant_path_node_enter,
           :on_constant_read_node_enter,
           :on_symbol_node_enter,
+          :on_string_node_enter,
         )
       end
 
@@ -56,6 +57,11 @@ def on_symbol_node_enter(node)
         handle_possible_dsl(node)
       end
 
+      #: (Prism::StringNode node) -> void
+      def on_string_node_enter(node)
+        handle_possible_dsl(node)
+      end
+
       private
 
       #: (String name) -> void
@@ -116,28 +122,67 @@ def format_default(default_value, type)
         end
       end
 
-      #: (Prism::SymbolNode node) -> void
+      #: ((Prism::SymbolNode | Prism::StringNode) node) -> void
       def handle_possible_dsl(node)
-        node = @node_context.call_node
-        return unless node
-        return unless self_receiver?(node)
-
-        message = node.message
+        call_node = @node_context.call_node
+        return unless call_node
+        return unless self_receiver?(call_node)
 
+        message = call_node.message
         return unless message
 
         if Support::Associations::ALL.include?(message)
-          handle_association(node)
+          handle_association(node, call_node)
+        end
+      end
+
+      #: ((Prism::SymbolNode | Prism::StringNode) node, Prism::CallNode call_node) -> void
+      def handle_association(node, call_node)
+        arguments = call_node.arguments&.arguments
+        return unless arguments
+
+        first_argument = arguments.first
+        return unless first_argument.is_a?(Prism::SymbolNode) || first_argument.is_a?(Prism::StringNode)
+
+        association_name = extract_string_from_node(first_argument)
+        return unless association_name
+
+        through_element = find_through_association_element(arguments)
+        clicked_symbol = extract_string_from_node(node)
+        return unless clicked_symbol
+
+        if through_element
+          through_association_name = extract_string_from_node(through_element.value)
+
+          if clicked_symbol == association_name
+            handle_association_name(association_name)
+          elsif through_association_name && clicked_symbol == through_association_name
+            handle_association_name(through_association_name)
+          end
+        else
+          handle_association_name(association_name)
         end
       end
 
-      #: (Prism::CallNode node) -> void
-      def handle_association(node)
-        first_argument = node.arguments&.arguments&.first
-        return unless first_argument.is_a?(Prism::SymbolNode)
+      #: (Array[Prism::Node]) -> Prism::AssocNode?
+      def find_through_association_element(arguments)
+        result = arguments
+          .filter_map { |arg| arg.elements if arg.is_a?(Prism::KeywordHashNode) }
+          .flatten
+          .find do |elem|
+            next false unless elem.is_a?(Prism::AssocNode)
+
+            key = elem.key
+            next false unless key.is_a?(Prism::SymbolNode)
 
-        association_name = first_argument.unescaped
+            key.value == "through"
+          end
 
+        result if result.is_a?(Prism::AssocNode)
+      end
+
+      #: (String association_name) -> void
+      def handle_association_name(association_name)
         result = @client.association_target(
           model_name: @nesting.join("::"),
           association_name: association_name,
@@ -164,6 +209,16 @@ def generate_hover(name)
           @response_builder.push(content, category: category)
         end
       end
+
+      #: (Prism::Node) -> String?
+      def extract_string_from_node(node)
+        case node
+        when Prism::SymbolNode
+          node.unescaped
+        when Prism::StringNode
+          node.content
+        end
+      end
     end
   end
 end

test/ruby_lsp_rails/hover_test.rb

@@ -321,6 +321,162 @@ class Bar < ApplicationRecord
         CONTENT
       end
 
+      test "returns main association on has_many :through association" do
+        expected_response = {
+          location: "#{dummy_root}/app/models/user.rb:2",
+          name: "User",
+        }
+        RunnerClient.any_instance.stubs(association_target: expected_response)
+
+        response = hover_on_source(<<~RUBY, { line: 2, character: 14 })
+          class Organization < ApplicationRecord
+            has_many :memberships
+            has_many :users, through: :memberships
+          end
+
+          class User < ApplicationRecord
+          end
+        RUBY
+
+        assert_equal(<<~CONTENT.chomp, response.contents.value)
+          ```ruby
+          User
+          ```
+
+          **Definitions**: [fake.rb](file:///fake.rb#L6,1-7,4)
+        CONTENT
+      end
+
+      test "returns through association on has_many :through association" do
+        expected_response = {
+          location: "#{dummy_root}/app/models/membership.rb:3",
+          name: "Membership",
+        }
+        RunnerClient.any_instance.stubs(association_target: expected_response)
+
+        response = hover_on_source(<<~RUBY, { line: 2, character: 31 })
+          class Organization < ApplicationRecord
+            has_many :memberships
+            has_many :users, through: :memberships
+          end
+
+          class Membership < ApplicationRecord
+          end
+        RUBY
+
+        assert_equal(<<~CONTENT.chomp, response.contents.value)
+          ```ruby
+          Membership
+          ```
+
+          **Definitions**: [fake.rb](file:///fake.rb#L6,1-7,4)
+        CONTENT
+      end
+
+      test "returns main association on has_one :through association" do
+        expected_response = {
+          location: "#{dummy_root}/app/models/flag.rb:2",
+          name: "Flag",
+        }
+        RunnerClient.any_instance.stubs(association_target: expected_response)
+
+        response = hover_on_source(<<~RUBY, { line: 2, character: 13 })
+          class User < ApplicationRecord
+            belongs_to :location, class_name: "Country"
+            has_one :country_flag, through: :location, source: :flag
+          end
+
+          class Flag < ApplicationRecord
+          end
+        RUBY
+
+        assert_equal(<<~CONTENT.chomp, response.contents.value)
+          ```ruby
+          Flag
+          ```
+
+          **Definitions**: [fake.rb](file:///fake.rb#L6,1-7,4)
+        CONTENT
+      end
+
+      test "returns through association on has_one :through association" do
+        expected_response = {
+          location: "#{dummy_root}/app/models/country.rb:2",
+          name: "Country",
+        }
+        RunnerClient.any_instance.stubs(association_target: expected_response)
+
+        response = hover_on_source(<<~RUBY, { line: 2, character: 37 })
+          class User < ApplicationRecord
+            belongs_to :location, class_name: "Country"
+            has_one :country_flag, through: :location, source: :flag
+          end
+
+          class Country < ApplicationRecord
+          end
+        RUBY
+
+        assert_equal(<<~CONTENT.chomp, response.contents.value)
+          ```ruby
+          Country
+          ```
+
+          **Definitions**: [fake.rb](file:///fake.rb#L6,1-7,4)
+        CONTENT
+      end
+
+      test "returns string main association on has_many :through association" do
+        expected_response = {
+          location: "#{dummy_root}/app/models/user.rb:2",
+          name: "User",
+        }
+        RunnerClient.any_instance.stubs(association_target: expected_response)
+
+        response = hover_on_source(<<~RUBY, { line: 2, character: 14 })
+          class Organization < ApplicationRecord
+            has_many :memberships
+            has_many "users", through: :memberships
+          end
+
+          class User < ApplicationRecord
+          end
+        RUBY
+
+        assert_equal(<<~CONTENT.chomp, response.contents.value)
+          ```ruby
+          User
+          ```
+
+          **Definitions**: [fake.rb](file:///fake.rb#L6,1-7,4)
+        CONTENT
+      end
+
+      test "returns string through association on has_many :through association" do
+        expected_response = {
+          location: "#{dummy_root}/app/models/membership.rb:3",
+          name: "Membership",
+        }
+        RunnerClient.any_instance.stubs(association_target: expected_response)
+
+        response = hover_on_source(<<~RUBY, { line: 2, character: 32 })
+          class Organization < ApplicationRecord
+            has_many "memberships"
+            has_many :users, through: "memberships"
+          end
+
+          class Membership < ApplicationRecord
+          end
+        RUBY
+
+        assert_equal(<<~CONTENT.chomp, response.contents.value)
+          ```ruby
+          Membership
+          ```
+
+          **Definitions**: [fake.rb](file:///fake.rb#L6,1-7,4)
+        CONTENT
+      end
+
       private
 
       def hover_on_source(source, position)