Merge pull request #519 from soutaro/better-splat-in-array-typing

Better splat in array typing

lib/steep/type_construction.rb

@@ -3945,6 +3945,20 @@ def try_tuple_type(node, hint)
       constr.add_typing(node, type: AST::Types::Tuple.new(types: element_types))
     end
 
+    def try_convert(type, method)
+      interface = checker.factory.interface(type, private: false)
+      if entry = interface.methods[method]
+        method_type = entry.method_types.find do |method_type|
+          method_type.type.params.optional?
+        end
+
+        method_type.type.return_type
+      end
+    rescue => exn
+      Steep.log_error(exn, message: "Unexpected error when converting #{type.to_s} with #{method}")
+      nil
+    end
+
     def try_array_type(node, hint)
       element_hint = hint ? hint.args[0] : nil
 
@@ -3955,8 +3969,14 @@ def try_array_type(node, hint)
         case child.type
         when :splat
           type, constr = constr.synthesize(child.children[0], hint: hint)
-          if AST::Builtin::Array.instance_type?(type)
+
+          type = try_convert(type, :to_a) || type
+
+          case
+          when AST::Builtin::Array.instance_type?(type)
             element_types << type.args[0]
+          when type.is_a?(AST::Types::Tuple)
+            element_types.push(*type.types)
           else
             element_types.push(*flatten_array_elements(type))
           end

test/type_construction_test.rb

@@ -8564,4 +8564,31 @@ def foo: (?untyped, *untyped, ?a: untyped, **untyped) -> String
       end
     end
   end
+
+  def test_splat_in_array
+    with_checker(<<-RBS) do |checker|
+class Range[T]
+  def to_a: () -> Array[T]
+end
+    RBS
+
+      source = parse_ruby(<<-'RUBY')
+a = [*'0'..'9']
+b = [*123]
+
+# @type var x: [Integer, String]
+x = [1, "a"]
+c = [*x]
+      RUBY
+
+      with_standard_construction(checker, source) do |construction, typing|
+        _, constr = construction.synthesize(source.node)
+
+        assert_no_error typing
+        assert_equal parse_type("::Array[::String]"), constr.context.lvar_env[:a]
+        assert_equal parse_type("::Array[::Integer]"), constr.context.lvar_env[:b]
+        assert_equal parse_type("::Array[::Integer | ::String]"), constr.context.lvar_env[:c]
+      end
+    end
+  end
 end