From 2d2ff32ff933ec50268590f0c8310959d47ad3b2 Mon Sep 17 00:00:00 2001
From: Alexandre Terrasa <alexandre.terrasa@shopify.com>
Date: Thu, 3 Oct 2024 14:41:20 -0400
Subject: [PATCH] Do not rely on return type for `attr_writer` RBS signatures

Signed-off-by: Alexandre Terrasa <alexandre.terrasa@shopify.com>
---
 lib/rbi/rbs_printer.rb       | 17 ++++++++++++++---
 test/rbi/rbs_printer_test.rb | 17 +++++++++++++++--
 2 files changed, 29 insertions(+), 5 deletions(-)

diff --git a/lib/rbi/rbs_printer.rb b/lib/rbi/rbs_printer.rb
index ff4290c..60e818b 100644
--- a/lib/rbi/rbs_printer.rb
+++ b/lib/rbi/rbs_printer.rb
@@ -262,8 +262,8 @@ def visit_attr(node)
         print(" #{name}")
         first_sig, *_rest = node.sigs # discard remaining signatures
         if first_sig
-          type = parse_type(first_sig.return_type)
-          print(": #{type.rbs_string}")
+          print(": ")
+          print_attr_sig(node, first_sig)
         else
           print(": untyped")
         end
@@ -273,7 +273,18 @@ def visit_attr(node)
 
     sig { params(node: RBI::Attr, sig: Sig).void }
     def print_attr_sig(node, sig)
-      type = parse_type(sig.return_type)
+      type = case node
+      when AttrAccessor, AttrReader
+        parse_type(sig.return_type)
+      else
+        first_arg = sig.params.first
+        if first_arg
+          parse_type(first_arg.type)
+        else
+          Type.untyped
+        end
+      end
+
       print(type.rbs_string)
     end
 
diff --git a/test/rbi/rbs_printer_test.rb b/test/rbi/rbs_printer_test.rb
index c03ec3b..8cec0bb 100644
--- a/test/rbi/rbs_printer_test.rb
+++ b/test/rbi/rbs_printer_test.rb
@@ -230,7 +230,17 @@ def test_print_attributes_with_signatures
         attr_accessor :foo
 
         sig { returns(String) }
-        attr_accessor :bar
+        sig { returns(Integer) }
+        attr_reader :bar
+
+        sig { params(baz: String).returns(String) }
+        attr_writer :baz
+
+        sig { params(qux: String).void }
+        attr_writer :qux
+
+        sig { returns(Integer) }
+        attr_writer :quux
       RBI
 
       # With RBS, attributes can only have one and only one return type
@@ -238,7 +248,10 @@ def test_print_attributes_with_signatures
       # and ignore the rest
       assert_equal(<<~RBI, rbi.rbs_string)
         attr_accessor foo: Integer
-        attr_accessor bar: String
+        attr_reader bar: String
+        attr_writer baz: String
+        attr_writer qux: String
+        attr_writer quux: untyped
       RBI
     end