Add a count function

Similar to the ruby count method on arrays.

2 files changed, 49 insertions, 0 deletions

diff --git a/lib/puppet/parser/functions/count.rb b/lib/puppet/parser/functions/count.rb
new file mode 100644
index 0000000..52de1b8
--- /dev/null
+++ b/lib/puppet/parser/functions/count.rb
@@ -0,0 +1,22 @@
+module Puppet::Parser::Functions
+  newfunction(:count, :type => :rvalue, :arity => -2, :doc => <<-EOS
+Takes an array as first argument and an optional second argument.
+Count the number of elements in array that matches second argument.
+If called with only an array it counts the number of elements that are not nil/undef.
+    EOS
+  ) do |args|
+
+    if (args.size > 2) then
+      raise(ArgumentError, "count(): Wrong number of arguments "+
+        "given #{args.size} for 1 or 2.")
+    end
+
+    collection, item = args
+
+    if item then
+      collection.count item
+    else
+      collection.count { |obj| obj != nil && obj != :undef && obj != '' }
+    end
+  end
+end
diff --git a/spec/unit/puppet/parser/functions/count_spec.rb b/spec/unit/puppet/parser/functions/count_spec.rb
new file mode 100644
index 0000000..5e4a8b9
--- /dev/null
+++ b/spec/unit/puppet/parser/functions/count_spec.rb
@@ -0,0 +1,27 @@
+#! /usr/bin/env ruby -S rspec
+
+require 'spec_helper'
+
+describe "the count function" do
+  let(:scope) { PuppetlabsSpec::PuppetInternals.scope }
+
+  it "should exist" do
+    Puppet::Parser::Functions.function("count").should == "function_count"
+  end
+
+  it "should raise a ArgumentError if there is more than 2 arguments" do
+    lambda { scope.function_count(['foo', 'bar', 'baz']) }.should( raise_error(ArgumentError))
+  end
+
+  it "should be able to count arrays" do
+    scope.function_count([["1","2","3"]]).should(eq(3))
+  end
+
+  it "should be able to count matching elements in arrays" do
+    scope.function_count([["1", "2", "2"], "2"]).should(eq(2))
+  end
+
+  it "should not count :undef, nil or empty strings" do
+    scope.function_count([["foo","bar",:undef,nil,""]]).should(eq(2))
+  end
+end