From 43cec6174f60cb5262c26d4c4bb33686a9dbd4ea Mon Sep 17 00:00:00 2001
From: Vince Broz <vince@broz.cc>
Date: Mon, 14 Oct 2024 12:52:12 -0400
Subject: [PATCH] Support writers in Rails models

For rails models with a 'foo' column, create pins for 'foo=' as
well as the existing 'foo' method.
---
 .github/workflows/ruby.yml             |  2 ++
 lib/solargraph/rails/annotate.rb       | 10 ++++++++++
 lib/solargraph/rails/schema.rb         | 23 +++++++++++++++++------
 lib/solargraph/rails/util.rb           | 19 ++++++++++++++++++-
 spec/helpers.rb                        | 10 +++++++++-
 spec/solargraph-rails/annotate_spec.rb |  2 ++
 spec/solargraph-rails/schema_spec.rb   |  3 ++-
 7 files changed, 60 insertions(+), 9 deletions(-)

diff --git a/.github/workflows/ruby.yml b/.github/workflows/ruby.yml
index cd894d7..9922a4c 100644
--- a/.github/workflows/ruby.yml
+++ b/.github/workflows/ruby.yml
@@ -20,9 +20,11 @@ jobs:
         ruby-version:
           - "3.1"
           - "3.2"
+          - "3.3"
         solargraph-version:
           - "0.48.0"
           - "0.49.0"
+          - "0.50.0"
       fail-fast: false
 
     steps:
diff --git a/lib/solargraph/rails/annotate.rb b/lib/solargraph/rails/annotate.rb
index 6cedc7d..d15e717 100644
--- a/lib/solargraph/rails/annotate.rb
+++ b/lib/solargraph/rails/annotate.rb
@@ -35,6 +35,16 @@ def process(source_map, ns)
               location:
                 Solargraph::Location.new(source_map.filename, snip.range)
             )
+
+          pins <<
+            Util.build_public_method(
+              ns,
+              "#{name}=",
+              types: [ruby_type],
+              params: { 'value' => [ruby_type] },
+              location:
+                Solargraph::Location.new(source_map.filename, snip.range)
+            )
         end
 
         pins
diff --git a/lib/solargraph/rails/schema.rb b/lib/solargraph/rails/schema.rb
index 93d6513..b4be7ce 100644
--- a/lib/solargraph/rails/schema.rb
+++ b/lib/solargraph/rails/schema.rb
@@ -43,15 +43,26 @@ def process(source_map, ns)
 
         return [] unless table
 
-        pins =
-          table.map do |column, data|
+        pins = table.flat_map do |column, data|
+          ruby_type = RUBY_TYPES.fetch(data.type.to_sym)
+          location = Util.build_location(data.ast, 'db/schema.rb')
+
+          [
             Util.build_public_method(
               ns,
               column,
-              types: [RUBY_TYPES.fetch(data.type.to_sym)],
-              location: Util.build_location(data.ast, 'db/schema.rb')
-            )
-          end
+              types: [ruby_type],
+              location: location
+            ),
+            Util.build_public_method(
+              ns,
+              "#{column}=",
+              types: [ruby_type],
+              params: { 'value' => [ruby_type] },
+              location: location
+            ),
+          ]
+        end
 
         if pins.any?
           Solargraph.logger.debug(
diff --git a/lib/solargraph/rails/util.rb b/lib/solargraph/rails/util.rb
index bc86745..2de643a 100644
--- a/lib/solargraph/rails/util.rb
+++ b/lib/solargraph/rails/util.rb
@@ -5,6 +5,7 @@ def self.build_public_method(
         ns,
         name,
         types: nil,
+        params: {},
         location: nil,
         attribute: false,
         scope: :instance
@@ -18,11 +19,27 @@ def self.build_public_method(
         }
 
         comments = []
+        params.each do |name, types|
+          comments << "@param [#{types.join(',')}] #{name}"
+        end
         comments << "@return [#{types.join(',')}]" if types
 
         opts[:comments] = comments.join("\n")
 
-        Solargraph::Pin::Method.new(**opts)
+        m = Solargraph::Pin::Method.new(**opts)
+        parameters = params.map do |name, type|
+            Solargraph::Pin::Parameter.new(
+              location: nil,
+              closure: m,
+              comments: '',
+              name: name,
+              presence: nil,
+              decl: :arg,
+              asgn_code: nil
+            )
+        end
+        m.parameters.concat(parameters)
+        m
       end
 
       def self.build_module_include(ns, module_name, location)
diff --git a/spec/helpers.rb b/spec/helpers.rb
index aff643c..885aac9 100644
--- a/spec/helpers.rb
+++ b/spec/helpers.rb
@@ -132,11 +132,19 @@ def use_workspace(folder, &block)
     map
   end
 
-  def assert_public_instance_method(map, query, return_type, &block)
+  def assert_public_instance_method(map, query, return_type, args: {}, &block)
     pin = find_pin(query, map)
     expect(pin).to_not be_nil
     expect(pin.scope).to eq(:instance)
     expect(pin.return_type.map(&:tag)).to eq(return_type)
+    args.each_pair do |name, type|
+      expect(parameter = pin.parameters.find { _1.name == name.to_s }).to_not be_nil
+      expect(parameter.return_type.tag).to eq(type)
+    end
+    pin.parameters.each do |param|
+      expect(args).to have_key(param.name.to_sym)
+      expect(param.return_type.tag).to eq(args[param.name.to_sym])
+    end
 
     yield pin if block_given?
   end
diff --git a/spec/solargraph-rails/annotate_spec.rb b/spec/solargraph-rails/annotate_spec.rb
index 84b8648..f2cc881 100644
--- a/spec/solargraph-rails/annotate_spec.rb
+++ b/spec/solargraph-rails/annotate_spec.rb
@@ -25,6 +25,8 @@ class MyModel < ApplicationRecord
     end
 
     assert_public_instance_method(api_map, 'MyModel#start_date', ['Date'])
+    assert_public_instance_method(api_map, 'MyModel#start_date=', ['Date'],
+                                  args: { value: 'Date' })
     assert_public_instance_method(
       api_map,
       'MyModel#living_expenses',
diff --git a/spec/solargraph-rails/schema_spec.rb b/spec/solargraph-rails/schema_spec.rb
index 7775586..b5f2d00 100644
--- a/spec/solargraph-rails/schema_spec.rb
+++ b/spec/solargraph-rails/schema_spec.rb
@@ -55,6 +55,8 @@ class Account < ActiveRecord::Base
     end
 
     assert_public_instance_method(map, "Account#balance", ["BigDecimal"])
+    assert_public_instance_method(map, "Account#balance=", ["BigDecimal"],
+                                  args: { value: 'BigDecimal' })
     assert_public_instance_method(map, "Account#some_int", ["Integer"])
     assert_public_instance_method(map, "Account#some_date", ["Date"])
     assert_public_instance_method(map, "Account#some_big_id", ["Integer"])
@@ -129,4 +131,3 @@ class Invoice < ActiveRecord::Base
     assert_public_instance_method(map, 'Invoice#amount', ['BigDecimal'])
   end
 end
-