From 332566babeb18b58cfd736d0a64c9b723a63543e Mon Sep 17 00:00:00 2001
From: Soutaro Matsumoto <matsumoto@soutaro.com>
Date: Thu, 1 Aug 2024 16:50:10 +0900
Subject: [PATCH] Skip `#:` syntax in `Data.define` and `Struct.new`

---
 lib/steep.rb            |  2 +-
 lib/steep/source.rb     | 40 ++++++++++++++++++++++++++++++++++++++++
 sig/steep/source.rbs    |  8 ++++++++
 test/type_check_test.rb | 38 ++++++++++++++++++++++++++++++++++++++
 4 files changed, 87 insertions(+), 1 deletion(-)

diff --git a/lib/steep.rb b/lib/steep.rb
index a65039692..90e7f3691 100644
--- a/lib/steep.rb
+++ b/lib/steep.rb
@@ -76,11 +76,11 @@
 require "steep/diagnostic/signature"
 require "steep/diagnostic/lsp_formatter"
 require "steep/signature/validator"
+require "steep/module_helper"
 require "steep/source"
 require "steep/source/ignore_ranges"
 require "steep/annotation_parser"
 require "steep/typing"
-require "steep/module_helper"
 require "steep/type_construction"
 require "steep/type_inference/context"
 require "steep/type_inference/context_array"
diff --git a/lib/steep/source.rb b/lib/steep/source.rb
index db8183050..6a13ad39c 100644
--- a/lib/steep/source.rb
+++ b/lib/steep/source.rb
@@ -8,6 +8,7 @@ class Source
     attr_reader :ignores
 
     extend NodeHelper
+    extend ModuleHelper
 
     def initialize(buffer:, path:, node:, mapping:, comments:, ignores:)
       @buffer = buffer
@@ -450,6 +451,32 @@ def without_unrelated_defs(line:, column:)
       end
     end
 
+    def self.skip_arg_assertions(node)
+      send_node, _ = deconstruct_sendish_and_block_nodes(node)
+      return false unless send_node
+
+      if send_node.type == :send
+        receiver, method, args = deconstruct_send_node!(send_node)
+
+        return false unless receiver
+
+        if receiver.type == :const
+          if type_name = module_name_from_node(receiver.children[0], receiver.children[1])
+            if type_name.namespace.empty?
+              if type_name.name == :Data && method == :define
+                return true
+              end
+              if type_name.name == :Struct && method == :new
+                return true
+              end
+            end
+          end
+        end
+      end
+
+      false
+    end
+
     def self.insert_type_node(node, comments)
       if node.location.expression
         first_line = node.location.expression.first_line
@@ -573,6 +600,19 @@ def self.insert_type_node(node, comments)
         body = insert_type_node(body, comments) if body
         return adjust_location(node.updated(nil, [object, name, args, body]))
       else
+        if skip_arg_assertions(node)
+          # Data.define, Struct.new, ...??
+          if node.location.expression
+            first_line = node.location.expression.first_line
+            last_line = node.location.expression.last_line
+
+            child_assertions = comments.delete_if {|line, _ | first_line < line && line < last_line }
+            node = map_child_node(node) {|child| insert_type_node(child, child_assertions) }
+
+            return adjust_location(node)
+          end
+        end
+
         adjust_location(
           map_child_node(node, nil) {|child| insert_type_node(child, comments) }
         )
diff --git a/sig/steep/source.rbs b/sig/steep/source.rbs
index 3939961ae..59bbb6dff 100644
--- a/sig/steep/source.rbs
+++ b/sig/steep/source.rbs
@@ -2,6 +2,8 @@ module Steep
   class Source
     extend NodeHelper
 
+    extend ModuleHelper
+
     attr_reader buffer: RBS::Buffer
 
     attr_reader path: Pathname
@@ -92,6 +94,12 @@ module Steep
 
     def self.insert_type_node: (Parser::AST::Node node, Hash[Integer, type_comment]) -> Parser::AST::Node
 
+    # Skip type assertions on arguments
+    #
+    # `Data.define` and `Struct.new` are examples of methods that have type assertions on arguments.
+    #
+    def self.skip_arg_assertions: (Parser::AST::Node) -> bool
+
     def self.adjust_location: (Parser::AST::Node) -> Parser::AST::Node
 
     # Returns an `:assertion` node with `TypeAssertion`
diff --git a/test/type_check_test.rb b/test/type_check_test.rb
index 99310de85..bf10f0a96 100644
--- a/test/type_check_test.rb
+++ b/test/type_check_test.rb
@@ -2103,4 +2103,42 @@ def foo(x) #: Integer
       YAML
     )
   end
+
+  def test_data_struct_annotation
+    run_type_check_test(
+      signatures: {
+        "a.rbs" => <<~RBS
+          class Hello < Data
+            attr_reader name(): String
+
+            def self.new: (String name) -> instance
+                        | (name: String) -> instance
+          end
+
+          class World < Struct[untyped]
+            attr_accessor size(): Integer
+
+            def self.new: (Integer size) -> instance
+                        | (size: Integer) -> instance
+          end
+        RBS
+      },
+      code: {
+        "a.rb" => <<~RUBY
+          Hello = Data.define(
+            :name #: String
+          )
+
+          World = Struct.new(
+            :size #: Integer
+          )
+        RUBY
+      },
+      expectations: <<~YAML
+        ---
+        - file: a.rb
+          diagnostics: []
+      YAML
+    )
+  end
 end