From d321f8630e9d2b0ed117d32706dfd7e980eb06eb Mon Sep 17 00:00:00 2001
From: Alexander Momchilov <alexander.momchilov@shopify.com>
Date: Thu, 6 Jun 2024 12:00:01 -0400
Subject: [PATCH] Rewrite `attr_*` into `def` methods

---
 lib/rbi.rb                                    |   1 +
 lib/rbi/model.rb                              |  56 ++++
 .../replace_attributes_with_methods.rb        | 146 +++++++++
 .../replace_attributes_with_methods_test.rb   | 292 ++++++++++++++++++
 4 files changed, 495 insertions(+)
 create mode 100644 lib/rbi/rewriters/replace_attributes_with_methods.rb
 create mode 100644 test/rbi/rewriters/replace_attributes_with_methods_test.rb

diff --git a/lib/rbi.rb b/lib/rbi.rb
index e2ad99fd..bcf22837 100644
--- a/lib/rbi.rb
+++ b/lib/rbi.rb
@@ -16,6 +16,7 @@
 require "rbi/rewriters/nest_non_public_methods"
 require "rbi/rewriters/group_nodes"
 require "rbi/rewriters/remove_known_definitions"
+require "rbi/rewriters/replace_attributes_with_methods"
 require "rbi/rewriters/sort_nodes"
 require "rbi/parser"
 require "rbi/printer"
diff --git a/lib/rbi/model.rb b/lib/rbi/model.rb
index ea579190..91672606 100644
--- a/lib/rbi/model.rb
+++ b/lib/rbi/model.rb
@@ -42,6 +42,16 @@ def replace(node)
       self.parent_tree = nil
     end
 
+    sig { params(nodes: T::Enumerable[Node]).void }
+    def replace_with_multiple(nodes)
+      tree = parent_tree
+      raise unless tree
+
+      # Does this work?
+      nodes.each { |node| tree << node }
+      detach
+    end
+
     sig { returns(T.nilable(Scope)) }
     def parent_scope
       parent = T.let(parent_tree, T.nilable(Tree))
@@ -1153,6 +1163,43 @@ def initialize(
       block&.call(self)
     end
 
+    sig do
+      params(
+        params: T::Array[SigParam],
+        return_type: T.nilable(String),
+        is_abstract: T::Boolean,
+        is_override: T::Boolean,
+        is_overridable: T::Boolean,
+        is_final: T::Boolean,
+        type_params: T::Array[String],
+        checked: T.nilable(Symbol),
+        loc: T.nilable(Loc),
+      ).returns(Sig)
+    end
+    def new_with(
+      params: @params,
+      return_type: @return_type,
+      is_abstract: @is_abstract,
+      is_override: @is_override,
+      is_overridable: @is_overridable,
+      is_final: @is_final,
+      type_params: @type_params.dup,
+      checked: @checked,
+      loc: @loc.dup
+    )
+      Sig.new(
+        params: params,
+        return_type: return_type,
+        is_abstract: is_abstract,
+        is_override: is_override,
+        is_overridable: is_overridable,
+        is_final: is_final,
+        type_params: type_params,
+        checked: checked,
+        loc: loc,
+      )
+    end
+
     sig { params(param: SigParam).void }
     def <<(param)
       @params << param
@@ -1171,6 +1218,15 @@ def ==(other)
         is_override == other.is_override && is_overridable == other.is_overridable && is_final == other.is_final &&
         type_params == other.type_params && checked == other.checked
     end
+
+    sig { override.returns(String) }
+    def inspect
+      io = StringIO.new
+
+      Printer.new(out: io, indent: 0).visit(self)
+
+      io.string.chomp
+    end
   end
 
   class SigParam < NodeWithComments
diff --git a/lib/rbi/rewriters/replace_attributes_with_methods.rb b/lib/rbi/rewriters/replace_attributes_with_methods.rb
new file mode 100644
index 00000000..b875346e
--- /dev/null
+++ b/lib/rbi/rewriters/replace_attributes_with_methods.rb
@@ -0,0 +1,146 @@
+# typed: strict
+# frozen_string_literal: true
+
+module RBI
+  module TempHelpers
+    extend T::Sig
+  end
+
+  module Rewriters
+    class ReplaceAttributesWithMethods < Visitor
+      extend T::Sig
+      include TempHelpers
+
+      sig { override.params(node: T.nilable(Node)).void }
+      def visit(node)
+        return unless node
+
+        case node
+        when Tree
+          node.nodes.dup.each do |child|
+            visit(child)
+            next unless (attr = child).is_a?(Attr)
+
+            new_methods = convert_to_methods(attr)
+
+            child.replace_with_multiple(new_methods)
+          end
+        end
+      end
+
+      private
+
+      sig { params(attr: Attr).returns(T::Array[Method]) }
+      def convert_to_methods(attr)
+        sig, attribute_type = parse_sig_of(attr)
+
+        case attr
+        when AttrReader then convert_attr_reader_to_methods(attr, sig, attribute_type)
+        when AttrWriter then convert_attr_writer_to_methods(attr, sig, attribute_type)
+        when AttrAccessor then convert_attr_accessor_to_methods(attr, sig, attribute_type)
+        else raise NotImplementedError, "Unknown attribute type: #{attr.class}"
+        end
+      end
+
+      sig { params(attr: AttrReader, sig: T.nilable(Sig), attribute_type: T.nilable(String)).returns(T::Array[Method]) }
+      def convert_attr_reader_to_methods(attr, sig, attribute_type)
+        attr.names.map do |name|
+          create_getter_method(name.to_s, sig, attr.visibility, attr.loc, attr.comments)
+        end
+      end
+
+      sig { params(attr: AttrWriter, sig: T.nilable(Sig), attribute_type: T.nilable(String)).returns(T::Array[Method]) }
+      def convert_attr_writer_to_methods(attr, sig, attribute_type)
+        attr.names.map do |name|
+          create_setter_method(name.to_s, sig, attribute_type, attr.visibility, attr.loc, attr.comments)
+        end
+      end
+
+      sig do
+        params(attr: AttrAccessor, sig: T.nilable(Sig), attribute_type: T.nilable(String)).returns(T::Array[Method])
+      end
+      def convert_attr_accessor_to_methods(attr, sig, attribute_type)
+        readers = attr.names.flat_map do |name|
+          create_getter_method(name.to_s, sig, attr.visibility, attr.loc, attr.comments)
+        end
+
+        writers = attr.names.map do |name|
+          create_setter_method(name.to_s, sig, attribute_type, attr.visibility, attr.loc, attr.comments)
+        end
+
+        readers + writers
+      end
+
+      sig { params(attr: Attr).returns([T.nilable(Sig), T.nilable(String)]) }
+      def parse_sig_of(attr)
+        raise "Attributes cannot have more than 1 sig" if 1 < attr.sigs.count
+
+        sig = attr.sigs.first
+        return [nil, nil] unless sig
+
+        attribute_type = case attr
+        when AttrReader, AttrAccessor then sig.return_type
+        when AttrWriter then sig.params.first&.type
+        end
+
+        [sig, attribute_type]
+      end
+
+      sig do
+        params(
+          name: String,
+          sig: T.nilable(Sig),
+          visibility: Visibility,
+          loc: T.nilable(Loc),
+          comments: T::Array[Comment],
+        ).returns(Method)
+      end
+      def create_getter_method(name, sig, visibility, loc, comments)
+        Method.new(
+          name,
+          params: [],
+          visibility: visibility,
+          sigs: sig ? [sig] : [],
+          loc: loc,
+          comments: comments,
+        )
+      end
+
+      sig do
+        params(
+          name: String,
+          sig: T.nilable(Sig),
+          attribute_type: T.nilable(String),
+          visibility: Visibility,
+          loc: T.nilable(Loc),
+          comments: T::Array[Comment],
+        ).returns(Method)
+      end
+      def create_setter_method(name, sig, attribute_type, visibility, loc, comments) # rubocop:disable Metrics/ParameterLists
+        sig = if sig # Modify the original sig to correct the name, and remove the return type
+          params = attribute_type ? [SigParam.new(name, attribute_type)] : []
+          sig.new_with(params: params, return_type: "void")
+        end
+
+        Method.new(
+          "#{name}=",
+          params: [ReqParam.new(name)],
+          visibility: visibility,
+          sigs: sig ? [sig] : sig,
+          loc: loc,
+          comments: comments,
+        )
+      end
+    end
+  end
+
+  class Tree
+    extend T::Sig
+
+    sig { void }
+    def replace_attributes_with_methods!
+      visitor = Rewriters::ReplaceAttributesWithMethods.new
+      visitor.visit(self)
+    end
+  end
+end
diff --git a/test/rbi/rewriters/replace_attributes_with_methods_test.rb b/test/rbi/rewriters/replace_attributes_with_methods_test.rb
new file mode 100644
index 00000000..4038c468
--- /dev/null
+++ b/test/rbi/rewriters/replace_attributes_with_methods_test.rb
@@ -0,0 +1,292 @@
+# typed: true
+# frozen_string_literal: true
+
+require "test_helper"
+
+module RBI
+  class ReplaceAttributesWithMethods < Minitest::Test
+    def test_replaces_attr_reader_with_method
+      rbi = Parser.parse_string(<<~RBI)
+        # Lorum ipsum...
+        sig { returns(Integer) }
+        attr_reader :a
+      RBI
+
+      rbi.replace_attributes_with_methods!
+
+      assert_equal(<<~RBI, rbi.string)
+        # Lorum ipsum...
+        sig { returns(Integer) }
+        def a; end
+      RBI
+    end
+
+    def test_replaces_attr_writer_with_setter_method
+      rbi = Parser.parse_string(<<~RBI)
+        # Lorum ipsum...
+        sig { params(a: Integer).void }
+        attr_writer :a
+      RBI
+
+      rbi.replace_attributes_with_methods!
+
+      assert_equal(<<~RBI, rbi.string)
+        # Lorum ipsum...
+        sig { params(a: Integer).void }
+        def a=(a); end
+      RBI
+    end
+
+    def test_replaces_attr_writer_with_return_type_with_setter_method
+      # Sorbet allows either `.void` or `.returns(TheType)`.
+      # We'll support both, until Sorbet starts to prefer one or the other.
+
+      rbi = Parser.parse_string(<<~RBI)
+        # Lorum ipsum...
+        sig { params(a: Integer).returns(Integer) }
+        attr_writer :a
+      RBI
+
+      rbi.replace_attributes_with_methods!
+
+      assert_equal(<<~RBI, rbi.string)
+        # Lorum ipsum...
+        sig { params(a: Integer).void }
+        def a=(a); end
+      RBI
+    end
+
+    def test_replaces_attr_accessor_with_getter_and_setter_methods
+      rbi = Parser.parse_string(<<~RBI)
+        # Lorum ipsum...
+        sig { returns(Integer) }
+        attr_accessor :a
+      RBI
+
+      rbi.replace_attributes_with_methods!
+
+      assert_equal(<<~RBI, rbi.string)
+        # Lorum ipsum...
+        sig { returns(Integer) }
+        def a; end
+
+        # Lorum ipsum...
+        sig { params(a: Integer).void }
+        def a=(a); end
+      RBI
+    end
+
+    ### Testing for multiple attributes defined in a single declaration
+
+    def test_replaces_multi_attr_reader_with_methods
+      rbi = Parser.parse_string(<<~RBI)
+        # Lorum ipsum...
+        sig { returns(Integer) }
+        attr_reader :a, :b, :c
+      RBI
+
+      rbi.replace_attributes_with_methods!
+
+      assert_equal(<<~RBI, rbi.string)
+        # Lorum ipsum...
+        sig { returns(Integer) }
+        def a; end
+
+        # Lorum ipsum...
+        sig { returns(Integer) }
+        def b; end
+
+        # Lorum ipsum...
+        sig { returns(Integer) }
+        def c; end
+      RBI
+    end
+
+    def test_replaces_multi_attr_writer_with_methods
+      rbi = Parser.parse_string(<<~RBI)
+        # Lorum ipsum...
+        sig { params(a: Integer).void }
+        attr_writer :a, :b, :c
+      RBI
+
+      rbi.replace_attributes_with_methods!
+
+      assert_equal(<<~RBI, rbi.string)
+        # Lorum ipsum...
+        sig { params(a: Integer).void }
+        def a=(a); end
+
+        # Lorum ipsum...
+        sig { params(b: Integer).void }
+        def b=(b); end
+
+        # Lorum ipsum...
+        sig { params(c: Integer).void }
+        def c=(c); end
+      RBI
+    end
+
+    def test_replaces_multi_attr_accessor_with_methods
+      rbi = Parser.parse_string(<<~RBI)
+        # Lorum ipsum...
+        sig { returns(Integer) }
+        attr_accessor :a, :b, :c
+      RBI
+
+      rbi.replace_attributes_with_methods!
+
+      assert_equal(<<~RBI, rbi.string)
+        # Lorum ipsum...
+        sig { returns(Integer) }
+        def a; end
+
+        # Lorum ipsum...
+        sig { returns(Integer) }
+        def b; end
+
+        # Lorum ipsum...
+        sig { returns(Integer) }
+        def c; end
+
+        # Lorum ipsum...
+        sig { params(a: Integer).void }
+        def a=(a); end
+
+        # Lorum ipsum...
+        sig { params(b: Integer).void }
+        def b=(b); end
+
+        # Lorum ipsum...
+        sig { params(c: Integer).void }
+        def c=(c); end
+      RBI
+    end
+
+    ### Testing for sig modifiers
+    # We don't test for Abstract, because attribute declarations are treated as though
+    # they have always have a method body, and so they could never be abstract.
+
+    def test_replacing_attr_reader_copies_sig_modifiers
+      rbi = Parser.parse_string(<<~RBI)
+        class GrandParent
+          sig { overridable.returns(Integer) }
+          attr_reader :a
+        end
+
+        class Parent < GrandParent
+          sig { override.overridable.returns(Integer) }
+          attr_reader :a
+        end
+
+        class Child < Parent
+          sig(:final) { override.returns(Integer) }
+          attr_reader :a
+        end
+      RBI
+
+      rbi.replace_attributes_with_methods!
+
+      assert_equal(<<~RBI, rbi.string)
+        class GrandParent
+          sig { overridable.returns(Integer) }
+          def a; end
+        end
+
+        class Parent < GrandParent
+          sig { override.overridable.returns(Integer) }
+          def a; end
+        end
+
+        class Child < Parent
+          sig(:final) { override.returns(Integer) }
+          def a; end
+        end
+      RBI
+    end
+
+    def test_replacing_attr_writer_copies_sig_modifiers
+      rbi = Parser.parse_string(<<~RBI)
+        class GrandParent
+          sig { overridable.params(a: Integer).void }
+          attr_writer :a
+        end
+
+        class Parent < GrandParent
+          sig { override.overridable.params(a: Integer).void }
+          attr_writer :a
+        end
+
+        class Child < Parent
+          sig(:final) { override.params(a: Integer).void }
+          attr_writer :a
+        end
+      RBI
+
+      rbi.replace_attributes_with_methods!
+
+      assert_equal(<<~RBI, rbi.string)
+        class GrandParent
+          sig { overridable.params(a: Integer).void }
+          def a=(a); end
+        end
+
+        class Parent < GrandParent
+          sig { override.overridable.params(a: Integer).void }
+          def a=(a); end
+        end
+
+        class Child < Parent
+          sig(:final) { override.params(a: Integer).void }
+          def a=(a); end
+        end
+      RBI
+    end
+
+    def test_replacing_attr_accessor_copies_sig_modifiers
+      rbi = Parser.parse_string(<<~RBI)
+        class GrandParent
+          sig { overridable.returns(Integer) }
+          attr_accessor :a
+        end
+
+        class Parent < GrandParent
+          sig { override.overridable.returns(Integer) }
+          attr_accessor :a
+        end
+
+        class Child < Parent
+          sig(:final) { override.returns(Integer) }
+          attr_accessor :a
+        end
+      RBI
+
+      rbi.replace_attributes_with_methods!
+
+      assert_equal(<<~RBI, rbi.string)
+        class GrandParent
+          sig { overridable.returns(Integer) }
+          def a; end
+
+          sig { overridable.params(a: Integer).void }
+          def a=(a); end
+        end
+
+        class Parent < GrandParent
+          sig { override.overridable.returns(Integer) }
+          def a; end
+
+          sig { override.overridable.params(a: Integer).void }
+          def a=(a); end
+        end
+
+        class Child < Parent
+          sig(:final) { override.returns(Integer) }
+          def a; end
+
+          sig(:final) { override.params(a: Integer).void }
+          def a=(a); end
+        end
+      RBI
+    end
+  end
+end