Support nil, #call and #to_int

ruby · Jun 21, 2024 · 8191f90 · 8191f90
1 parent 11a6714
commit 8191f90
Show file tree

Hide file tree

Showing 2 changed files with 20 additions and 1 deletion.
diff --git a/core/enumerator.rbs b/core/enumerator.rbs
@@ -269,6 +269,10 @@ class Enumerator[unchecked out Elem, out Return] < Object
   #
   def feed: (Elem arg0) -> NilClass
 
+  interface _Call
+    def call: () -> untyped
+  end
+
   # <!--
   #   rdoc-file=enumerator.c
   #   - Enumerator.new(size = nil) { |yielder| ... }
@@ -293,7 +297,7 @@ class Enumerator[unchecked out Elem, out Return] < Object
   # lazy fashion (see Enumerator#size). It can either be a value or a callable
   # object.
   #
-  def initialize: (?Integer arg0) { (Enumerator::Yielder arg0) -> Return } -> void
+  def initialize: (?(nil | _Call | ::_ToInt) arg0) { (Enumerator::Yielder arg0) -> Return } -> void
 
   # <!--
   #   rdoc-file=enumerator.c

diff --git a/test/stdlib/Enumerator_test.rb b/test/stdlib/Enumerator_test.rb
@@ -34,6 +34,21 @@ class EnumeratorSingletonTest < Test::Unit::TestCase
   def test_new
     assert_send_type "() { (Enumerator::Yielder) -> 12345 } -> Enumerator[untyped, 12345]",
                      Enumerator, :new do 12345 end
+
+    assert_send_type "(123) { (Enumerator::Yielder) -> nil } -> Enumerator[untyped, nil]",
+                     Enumerator, :new, 123 do nil end
+
+    o = Object.new
+    def o.call = 'hi'
+    assert_send_type "(Enumerator::_Call) { (Enumerator::Yielder) -> nil } -> Enumerator[untyped, nil]",
+                     Enumerator, :new, o do nil end
+    assert_send_type "(Enumerator::_Call) { (Enumerator::Yielder) -> nil } -> Enumerator[untyped, nil]",
+                     Enumerator, :new, ->{'hi'} do nil end
+
+    with_int(123) do |int|
+      assert_send_type "(_ToInt) { (Enumerator::Yielder) -> nil } -> Enumerator[untyped, nil]",
+                       Enumerator, :new, int do nil end
+    end
   end
 
   def test_produce