From 5f50fe098c0491bcdb622dfaba9f3d5a42061b2f Mon Sep 17 00:00:00 2001
From: Alexandre Terrasa <alexandre.terrasa@shopify.com>
Date: Fri, 20 Dec 2024 11:09:55 -0500
Subject: [PATCH] Fix RBS translation of attr_writer signatures

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

diff --git a/lib/rbi/rbs_printer.rb b/lib/rbi/rbs_printer.rb
index dea265e1..4954ea66 100644
--- a/lib/rbi/rbs_printer.rb
+++ b/lib/rbi/rbs_printer.rb
@@ -273,19 +273,24 @@ def visit_attr(node)
 
     sig { params(node: RBI::Attr, sig: Sig).void }
     def print_attr_sig(node, sig)
+      ret_type = sig.return_type
+
       type = case node
-      when AttrAccessor, AttrReader
-        parse_type(sig.return_type)
+      when AttrReader, AttrAccessor
+        parse_type(ret_type).rbs_string
       else
-        first_arg = sig.params.first
-        if first_arg
-          parse_type(first_arg.type)
+        # For attr_writer, Sorbet will prioritize the return type over the argument type in case of mismatch
+        arg_type = sig.params.first
+        if arg_type && (ret_type.is_a?(Type::Void) || ret_type == "void")
+          # If we have an argument type and the return type is void, we prioritize the argument type
+          parse_type(arg_type.type).rbs_string
         else
-          Type.untyped
+          # Otherwise, we prioritize the return type
+          parse_type(ret_type).rbs_string
         end
       end
 
-      print(type.rbs_string)
+      print(type)
     end
 
     sig { override.params(node: Method).void }
@@ -416,6 +421,24 @@ def print_method_sig(node, sig)
       print(" # #{loc}") if loc && print_locs
     end
 
+    sig { params(node: Sig).void }
+    def visit_sig(node)
+      if node.params
+        print("(")
+        node.params.each do |param|
+          visit(param)
+        end
+        print(") ")
+      end
+
+      print("-> #{parse_type(node.return_type).rbs_string}")
+    end
+
+    sig { params(node: SigParam).void }
+    def visit_sig_param(node)
+      print(parse_type(node.type).rbs_string)
+    end
+
     sig { override.params(node: ReqParam).void }
     def visit_req_param(node)
       print("untyped #{node.name}")
diff --git a/test/rbi/rbs_printer_test.rb b/test/rbi/rbs_printer_test.rb
index 7ca06723..389751f3 100644
--- a/test/rbi/rbs_printer_test.rb
+++ b/test/rbi/rbs_printer_test.rb
@@ -233,7 +233,7 @@ def test_print_attributes_with_signatures
         sig { returns(Integer) }
         attr_reader :bar
 
-        sig { params(baz: String).returns(String) }
+        sig { params(baz: String).returns(T.nilable(String)) }
         attr_writer :baz
 
         sig { params(qux: String).void }
@@ -249,9 +249,9 @@ def test_print_attributes_with_signatures
       assert_equal(<<~RBI, rbi.rbs_string)
         attr_accessor foo: Integer
         attr_reader bar: String
-        attr_writer baz: String
+        attr_writer baz: String?
         attr_writer qux: String
-        attr_writer quux: untyped
+        attr_writer quux: Integer
       RBI
     end