From 3a6c868d767012846f5832773d6734b7e5830201 Mon Sep 17 00:00:00 2001
From: Alexandre Terrasa <alexandre.terrasa@shopify.com>
Date: Wed, 3 Apr 2024 17:25:17 -0400
Subject: [PATCH] Add an easier syntax to create methods and signatures

Signed-off-by: Alexandre Terrasa <alexandre.terrasa@shopify.com>
---
 lib/rbi/model.rb       | 78 ++++++++++++++++++++++++++++++++++++++++++
 test/rbi/model_test.rb | 42 +++++++++++++++++++++++
 2 files changed, 120 insertions(+)

diff --git a/lib/rbi/model.rb b/lib/rbi/model.rb
index 63153ed8..ea579190 100644
--- a/lib/rbi/model.rb
+++ b/lib/rbi/model.rb
@@ -534,6 +534,79 @@ def <<(param)
       @params << param
     end
 
+    sig { params(name: String).void }
+    def add_param(name)
+      @params << ReqParam.new(name)
+    end
+
+    sig { params(name: String, default_value: String).void }
+    def add_opt_param(name, default_value)
+      @params << OptParam.new(name, default_value)
+    end
+
+    sig { params(name: String).void }
+    def add_rest_param(name)
+      @params << RestParam.new(name)
+    end
+
+    sig { params(name: String).void }
+    def add_kw_param(name)
+      @params << KwParam.new(name)
+    end
+
+    sig { params(name: String, default_value: String).void }
+    def add_kw_opt_param(name, default_value)
+      @params << KwOptParam.new(name, default_value)
+    end
+
+    sig { params(name: String).void }
+    def add_kw_rest_param(name)
+      @params << KwRestParam.new(name)
+    end
+
+    sig { params(name: String).void }
+    def add_block_param(name)
+      @params << BlockParam.new(name)
+    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),
+        block: T.proc.params(node: Sig).void,
+      ).void
+    end
+    def add_sig(
+      params: [],
+      return_type: nil,
+      is_abstract: false,
+      is_override: false,
+      is_overridable: false,
+      is_final: false,
+      type_params: [],
+      checked: nil,
+      &block
+    )
+      sig = 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,
+        &block
+      )
+      @sigs << sig
+    end
+
     sig { returns(String) }
     def fully_qualified_name
       if is_singleton
@@ -1085,6 +1158,11 @@ def <<(param)
       @params << param
     end
 
+    sig { params(name: String, type: String).void }
+    def add_param(name, type)
+      @params << SigParam.new(name, type)
+    end
+
     sig { params(other: Object).returns(T::Boolean) }
     def ==(other)
       return false unless other.is_a?(Sig)
diff --git a/test/rbi/model_test.rb b/test/rbi/model_test.rb
index 39942533..ff8430f3 100644
--- a/test/rbi/model_test.rb
+++ b/test/rbi/model_test.rb
@@ -201,6 +201,48 @@ class Enum < ::T::Enum
       RBI
     end
 
+    def test_model_sig_builder
+      rbi = Tree.new do |tree|
+        tree << Method.new("foo") do |node|
+          node.add_param("p1")
+          node.add_opt_param("p2", "'value'")
+          node.add_rest_param("p3")
+          node.add_kw_param("p4")
+          node.add_kw_opt_param("p5", "'value'")
+          node.add_kw_rest_param("p6")
+          node.add_block_param("p7")
+
+          node.add_sig do |sig|
+            sig.add_param("p1", "T.untyped")
+            sig.add_param("p2", "String")
+            sig.return_type = "T.untyped"
+          end
+
+          node.add_sig do |sig|
+            sig.add_param("p3", "T.untyped")
+            sig.return_type = "void"
+          end
+
+          node.add_sig(type_params: ["T", "U"]) do |sig|
+            sig.is_abstract = true
+            sig.is_override = true
+            sig.is_overridable = true
+            sig.is_final = true
+            sig.checked = :never
+            sig.add_param("p4", "T.untyped")
+            sig.return_type = "void"
+          end
+        end
+      end
+
+      assert_equal(<<~RBI, rbi.string)
+        sig { params(p1: T.untyped, p2: String).returns(T.untyped) }
+        sig { params(p3: T.untyped).void }
+        sig(:final) { abstract.override.overridable.type_parameters(:T, :U).checked(:never).params(p4: T.untyped).void }
+        def foo(p1, p2 = 'value', *p3, p4:, p5: 'value', **p6, &p7); end
+      RBI
+    end
+
     def test_model_fully_qualified_names
       mod = Module.new("Foo")
       assert_equal("::Foo", mod.fully_qualified_name)