diff options
| -rw-r--r-- | enumerator.c | 7 | ||||
| -rw-r--r-- | test/ruby/test_enumerator.rb | 6 |
2 files changed, 11 insertions, 2 deletions
diff --git a/enumerator.c b/enumerator.c index 36e70fd084..a5eb999d8d 100644 --- a/enumerator.c +++ b/enumerator.c @@ -489,6 +489,9 @@ enumerator_with_index_i(VALUE val, VALUE m, int argc, VALUE *argv) return rb_yield_values(2, rb_ary_new4(argc, argv), idx); } +static VALUE +enumerator_size(VALUE obj); + /* * call-seq: * e.with_index(offset = 0) {|(*args), idx| ... } @@ -507,7 +510,7 @@ enumerator_with_index(int argc, VALUE *argv, VALUE obj) VALUE memo; rb_scan_args(argc, argv, "01", &memo); - RETURN_ENUMERATOR(obj, argc, argv); + RETURN_SIZED_ENUMERATOR(obj, argc, argv, enumerator_size); memo = NIL_P(memo) ? 0 : (VALUE)NUM2LONG(memo); return enumerator_block_call(obj, enumerator_with_index_i, (VALUE)&memo); } @@ -567,7 +570,7 @@ enumerator_with_object_i(VALUE val, VALUE memo, int argc, VALUE *argv) static VALUE enumerator_with_object(VALUE obj, VALUE memo) { - RETURN_ENUMERATOR(obj, 1, &memo); + RETURN_SIZED_ENUMERATOR(obj, 1, &memo, enumerator_size); enumerator_block_call(obj, enumerator_with_object_i, memo); return memo; diff --git a/test/ruby/test_enumerator.rb b/test/ruby/test_enumerator.rb index 8efc812fc5..3143be21a0 100644 --- a/test/ruby/test_enumerator.rb +++ b/test/ruby/test_enumerator.rb @@ -414,5 +414,11 @@ class TestEnumerator < Test::Unit::TestCase assert_equal nil, @obj.to_enum(:foo, 0, 1).size assert_equal 2, @obj.to_enum(:foo, 0, 1){ 2 }.size end + + def test_size_for_enum_created_by_enumerators + enum = to_enum{ 42 } + assert_equal 42, enum.with_index.size + assert_equal 42, enum.with_object(:foo).size + end end |