* enumerator.c (obj_to_enum): Have #to_enum accept a block

  [Feature #6636]

git-svn-id: svn+ssh://ci.ruby-lang.org/ruby/trunk@37498 b2dd03c8-39d4-4d8f-98ff-823fe69b080e

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