diff options
| -rw-r--r-- | enumerator.c | 13 | ||||
| -rw-r--r-- | test/ruby/test_enumerator.rb | 3 |
2 files changed, 14 insertions, 2 deletions
diff --git a/enumerator.c b/enumerator.c index 7695ea1a56..36e70fd084 100644 --- a/enumerator.c +++ b/enumerator.c @@ -186,6 +186,8 @@ enumerator_ptr(VALUE obj) * call-seq: * obj.to_enum(method = :each, *args) * obj.enum_for(method = :each, *args) + * obj.to_enum(method = :each, *args) {|obj, *args| block} + * obj.enum_for(method = :each, *args){|obj, *args| block} * * Creates a new Enumerator which will enumerate by on calling +method+ on * +obj+. @@ -195,6 +197,9 @@ enumerator_ptr(VALUE obj) * to the item itself. Note that the number of args * must not exceed the number expected by +method+ * + * If a block is given, it will be used to calculate the size of + * the enumerator (see Enumerator#size=). + * * === Example * * str = "xyz" @@ -213,13 +218,17 @@ enumerator_ptr(VALUE obj) static VALUE obj_to_enum(int argc, VALUE *argv, VALUE obj) { - VALUE meth = sym_each; + VALUE enumerator, meth = sym_each; if (argc > 0) { --argc; meth = *argv++; } - return rb_enumeratorize(obj, meth, argc, argv, 0); + enumerator = rb_enumeratorize(obj, meth, argc, argv, 0); + if (rb_block_given_p()) { + enumerator_ptr(enumerator)->size = rb_block_proc(); + } + return enumerator; } static VALUE diff --git a/test/ruby/test_enumerator.rb b/test/ruby/test_enumerator.rb index 0f19a4f8e7..8efc812fc5 100644 --- a/test/ruby/test_enumerator.rb +++ b/test/ruby/test_enumerator.rb @@ -410,6 +410,9 @@ class TestEnumerator < Test::Unit::TestCase assert_equal Float::INFINITY, Enumerator.new(Float::INFINITY){}.size assert_equal nil, Enumerator.new(nil){}.size assert_raise(TypeError) { Enumerator.new("42"){} } + + assert_equal nil, @obj.to_enum(:foo, 0, 1).size + assert_equal 2, @obj.to_enum(:foo, 0, 1){ 2 }.size end end |