Merge pull request #576 from YorickPeterse/walker-count-enumerable-compatible

Make Walker#count compatible with Enumerable#count

Author: arthurschreiber

ext/rugged/rugged_revwalk.c

@@ -194,15 +194,19 @@ static VALUE rb_git_walker_simplify_first_parent(VALUE self)
  *  call-seq:
  *    walker.count -> Fixnum
  *
- *  Returns the amount of objects a walker iterated over.
+ *  Returns the amount of objects a walker iterated over. If an argument or
+ *  block is given this method delegates to +Enumerable#count+.
  */
-static VALUE rb_git_walker_count(VALUE self)
+static VALUE rb_git_walker_count(int argc, VALUE *argv, VALUE self)
 {
 	git_revwalk *walk;
 	git_oid commit_oid;
 	int error = 0;
 	uint64_t count = 0;
 
+	if (argc > 0 || rb_block_given_p())
+		return rb_call_super(argc, argv);
+
 	Data_Get_Struct(self, git_revwalk, walk);
 
 	while (((error = git_revwalk_next(&commit_oid, walk)) == 0) && ++count != UINT64_MAX);
@@ -515,5 +519,5 @@ void Init_rugged_revwalk(void)
 	rb_define_method(rb_cRuggedWalker, "reset", rb_git_walker_reset, 0);
 	rb_define_method(rb_cRuggedWalker, "sorting", rb_git_walker_sorting, 1);
 	rb_define_method(rb_cRuggedWalker, "simplify_first_parent", rb_git_walker_simplify_first_parent, 0);
-	rb_define_method(rb_cRuggedWalker, "count", rb_git_walker_count, 0);
+	rb_define_method(rb_cRuggedWalker, "count", rb_git_walker_count, -1);
 }

test/walker_test.rb

@@ -149,6 +149,24 @@ def test_walk_count
     @walker.hide("5b5b025afb0b4c913b4c338a42934a3863bf3644")
     assert_equal 2, @walker.count
   end
+
+  def test_walk_count_argument
+    @walker.push("9fd738e8f7967c078dceed8190330fc8648ee56a")
+    @walker.hide("5b5b025afb0b4c913b4c338a42934a3863bf3644")
+
+    assert_equal 0, @walker.count('foo')
+  end
+
+  def test_walk_count_with_block
+    @walker.push("9fd738e8f7967c078dceed8190330fc8648ee56a")
+    @walker.hide("5b5b025afb0b4c913b4c338a42934a3863bf3644")
+
+    amount = @walker.count do |commit|
+      commit.oid == "9fd738e8f7967c078dceed8190330fc8648ee56a"
+    end
+
+    assert_equal 1, amount
+  end
 end
 
 # testrepo (the non-bare repo) is the one with non-linear history,