Merge pull request #368 from Shopify/at-rbs-attr

Do not rely on return type for `attr_writer` RBS signatures
Shopify · Oct 3, 2024 · 80784c2 · 80784c2
2 parents 174dcc9 + 2d2ff32
commit 80784c2
Show file tree

Hide file tree

Showing 2 changed files with 29 additions and 5 deletions.
diff --git 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
@@ -230,15 +230,28 @@ 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
       # if we have multiple signatures we just use the return type of the first one
       # 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