Fix RBS translation of attr_writer signatures

Signed-off-by: Alexandre Terrasa <alexandre.terrasa@shopify.com>
Shopify · Jan 15, 2025 · 5f50fe0 · 5f50fe0
1 parent e9e5148
commit 5f50fe0
Show file tree

Hide file tree

Showing 2 changed files with 33 additions and 10 deletions.
diff --git 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
@@ -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