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