[ruby-dev:32255] Re: [提案] Array#tail

遠藤です。

07/11/13 に Akinori MUSHA<[e-mail]> さんは書きました:
> ary.drop(1) はどうでしょうか。

あああ、すっかり存在を忘れてました。
ありがとうございます、これで大丈夫です。


お騒がせしたついでに、take と drop のデフォルトを 1 にするパッチ
(+ テスト) を書いてみました。どうでしょうか。


Index: enum.c
===================================================================
--- enum.c	(revision 13902)
+++ enum.c	(working copy)
@@ -1405,21 +1405,31 @@

 /*
  *  call-seq:
+ *     enum.take                  => array
  *     enum.take(n)               => array
  *
- *  Returns first n elements from <i>enum</i>.
+ *  Returns first n elements from <i>enum</i>. If n is not specified,
+ *  returns an array containing only the first element.
  *
  *     a = [1, 2, 3, 4, 5, 0]
+ *     a.take                # => [1]
  *     a.take(3)             # => [1, 2, 3]
  *
  */

 static VALUE
-enum_take(VALUE obj, VALUE n)
+enum_take(int argc, VALUE *argv, VALUE obj)
 {
     VALUE args[2];
+    VALUE n;

-    args[1] = NUM2LONG(n);
+    if (argc == 0) {
+	args[1] = 1;
+    }
+    else {
+	rb_scan_args(argc, argv, "01", &n);
+	args[1] = NUM2LONG(n);
+    }
     args[0] = rb_ary_new2(args[1]);
     rb_block_call(obj, id_each, 0, 0, take_i, (VALUE)args);
     return args[0];
@@ -1470,22 +1480,31 @@

 /*
  *  call-seq:
+ *     enum.drop                  => array
  *     enum.drop(n)               => array
  *
- *  Drops first n elements from <i>enum</i>, and returns rest elements
- *  in an array.
+ *  Drops first element, or first +n+ elements, from <i>enum</i>, and
+ *  returns rest elements in an array.
  *
  *     a = [1, 2, 3, 4, 5, 0]
+ *     a.drop                # => [2, 3, 4, 5, 0]
  *     a.drop(3)             # => [4, 5, 0]
  *
  */

 static VALUE
-enum_drop(VALUE obj, VALUE n)
+enum_drop(int argc, VALUE *argv, VALUE obj)
 {
     VALUE args[2];
+    VALUE n;

-    args[1] = NUM2ULONG(n);
+    if (argc == 0) {
+	args[1] = 1;
+    }
+    else {
+	rb_scan_args(argc, argv, "01", &n);
+	args[1] = NUM2LONG(n);
+    }
     args[0] = rb_ary_new2(args[1]);
     rb_block_call(obj, id_each, 0, 0, drop_i, (VALUE)args);
     return args[0];
@@ -1620,9 +1639,9 @@
     rb_define_method(rb_mEnumerable,"include?", enum_member, 1);
     rb_define_method(rb_mEnumerable,"each_with_index",
enum_each_with_index, -1);
     rb_define_method(rb_mEnumerable, "zip", enum_zip, -1);
-    rb_define_method(rb_mEnumerable, "take", enum_take, 1);
+    rb_define_method(rb_mEnumerable, "take", enum_take, -1);
     rb_define_method(rb_mEnumerable, "take_while", enum_take_while, 0);
-    rb_define_method(rb_mEnumerable, "drop", enum_drop, 1);
+    rb_define_method(rb_mEnumerable, "drop", enum_drop, -1);
     rb_define_method(rb_mEnumerable, "drop_while", enum_drop_while, 0);
     rb_define_method(rb_mEnumerable, "cycle", enum_cycle, 0);

Index: test/ruby/test_array.rb
===================================================================
--- test/ruby/test_array.rb	(revision 13902)
+++ test/ruby/test_array.rb	(working copy)
@@ -586,6 +586,18 @@
     assert_equal(@cls[1, 2, 3], a)
   end

+  def test_drop
+    assert_equal([4, 5],    @cls[3, 4, 5].drop)
+    assert_equal([3, 4, 5], @cls[3, 4, 5].drop(0))
+    assert_equal([4, 5],    @cls[3, 4, 5].drop(1))
+    assert_equal([5],       @cls[3, 4, 5].drop(2))
+    assert_equal([],        @cls[3, 4, 5].drop(3))
+    assert_equal([],        @cls[3, 4, 5].drop(4))
+    assert_raise(ArgumentError) { @cls[3, 4, 5].drop(-1) }
+    assert_equal([],        @cls[].drop)
+    assert_equal([],        @cls[].drop(1))
+  end
+
   def test_dup
     for taint in [ false, true ]
       for frozen in [ false, true ]
@@ -1103,6 +1115,18 @@
     assert_equal(@cls[], @cls[].sort!)
   end

+  def test_take
+    assert_equal([3],       @cls[3, 4, 5].take)
+    assert_equal([],        @cls[3, 4, 5].take(0))
+    assert_equal([3],       @cls[3, 4, 5].take(1))
+    assert_equal([3, 4],    @cls[3, 4, 5].take(2))
+    assert_equal([3, 4, 5], @cls[3, 4, 5].take(3))
+    assert_equal([3, 4, 5], @cls[3, 4, 5].take(4))
+    assert_raise(ArgumentError) { @cls[3, 4, 5].take(-1) }
+    assert_equal([],        @cls[].take)
+    assert_equal([],        @cls[].take(1))
+  end
+
   def test_to_a
     a = @cls[ 1, 2, 3 ]
     a_id = a.__id__

-- 
Yusuke ENDOH <[e-mail]>