From 000f0c330c821e911a7b46e40c5c253ddf748970 Mon Sep 17 00:00:00 2001
From: Soutaro Matsumoto <matsumoto@soutaro.com>
Date: Mon, 16 Dec 2024 21:48:32 +0900
Subject: [PATCH] Attach type assertion to keyword arg

---
 lib/steep/source.rb |  8 ++++++++
 test/source_test.rb | 41 +++++++++++++++++++++++++++++++++++++++++
 2 files changed, 49 insertions(+)

diff --git a/lib/steep/source.rb b/lib/steep/source.rb
index 7cda377d0..17fbcbf00 100644
--- a/lib/steep/source.rb
+++ b/lib/steep/source.rb
@@ -496,6 +496,14 @@ def self.insert_type_node(node, comments)
             # Skip
           when :def, :defs
             # Skip
+          when :kwargs
+            # skip
+          when :pair
+            key, value = node.children
+            key = insert_type_node(key, comments.except(last_line))
+            value = insert_type_node(value, comments)
+            node = node.updated(nil, [key, value])
+            return adjust_location(node)
           when :masgn
             lhs, rhs = node.children
             node = node.updated(nil, [lhs, insert_type_node(rhs, comments)])
diff --git a/test/source_test.rb b/test/source_test.rb
index 6ad5dd2fc..7c0f378b2 100644
--- a/test/source_test.rb
+++ b/test/source_test.rb
@@ -779,6 +779,47 @@ def self.baz(x) = 123 #: Numeric
     end
   end
 
+  def test_assertion__kwargs
+    with_factory do |factory|
+      source = Steep::Source.parse(<<~RUBY, path: Pathname("foo.rb"), factory: factory)
+        foo(
+          a: "", #: ""
+          b: [], #: Array[String]
+          c: ""
+        )
+      RUBY
+
+      _receiver, _name, *args = source.node.children
+      args[0].tap do |kwargs|
+        assert_equal :kwargs, kwargs.type
+
+        kwargs.children[0].tap do |pair|
+          assert_equal :pair, pair.type
+          key, value = pair.children
+
+          assert_equal :sym, key.type
+          assert_equal :assertion, value.type
+        end
+
+        kwargs.children[1].tap do |pair|
+          assert_equal :pair, pair.type
+          key, value = pair.children
+
+          assert_equal :sym, key.type
+          assert_equal :assertion, value.type
+        end
+
+        kwargs.children[2].tap do |pair|
+          assert_equal :pair, pair.type
+          key, value = pair.children
+
+          assert_equal :sym, key.type
+          assert_equal :str, value.type
+        end
+      end
+    end
+  end
+
   def test_tapp_send
     with_factory({ Pathname("foo.rbs") => <<-RBS }) do |factory|
 class Tapp