Make regsubst() function operate on arrays (feature #2491).

Allow the first argument to the regsubst() function be an array,
and perform regexp replacement on each element of the array in
that case.

This patch also adds more error checking to give better error
messages to the user when given bad parameters.

Signed-off-by: Thomas Bellman <bellman@nsc.liu.se>

1 files changed, 44 insertions, 30 deletions

diff --git a/lib/puppet/parser/functions/regsubst.rb b/lib/puppet/parser/functions/regsubst.rb
index e6b98eb2c..d02680862 100644
--- a/lib/puppet/parser/functions/regsubst.rb
+++ b/lib/puppet/parser/functions/regsubst.rb
@@ -1,22 +1,22 @@
 module Puppet::Parser::Functions
     newfunction(:regsubst, :type => :rvalue,
                 :doc => "
-                Perform regexp replacement on a string.
+                Perform regexp replacement on a string or array of strings.
 
 - **Parameters** (in order):
 
-:str:  The string to operate on.
+:target:  The string or array of strings to operate on.  If an array, the replacement will be performed on each of the elements in the array, and the return value will be an array.
 
-:regexp:  The regular expression matching the string.  If you want it anchored at the start and or end of the string, you must do that with ^ and $ yourself.
+:regexp:  The regular expression matching the target string.  If you want it anchored at the start and or end of the string, you must do that with ^ and $ yourself.
 
-:replacement:  Replacement string. Can contain back references to what was matched using 0, 1, and so on.
+:replacement:  Replacement string. Can contain back references to what was matched using \\0, \\1, and so on.
 
 :flags:  Optional. String of single letter flags for how the regexp is interpreted:
 
     - **E**         Extended regexps
     - **I**         Ignore case in regexps
     - **M**         Multiline regexps
-    - **G**         Global replacement; all occurrences of the regexp in the string will be replaced.  Without this, only the first occurrence will be replaced.
+    - **G**         Global replacement; all occurrences of the regexp in each target string will be replaced.  Without this, only the first occurrence will be replaced.
 
 :lang:  Optional.  How to handle multibyte characters.  A single-character string with the following values:
 
@@ -35,36 +35,50 @@ Put angle brackets around each octet in the node's IP address::
 
     $x = regsubst($ipaddress, '([0-9]+)', '<\\1>', 'G')
 ") \
-        do |args|
-        flag_mapping = {
-            "E" => Regexp::EXTENDED,
-            "I" => Regexp::IGNORECASE,
-            "M" => Regexp::MULTILINE,
-        }
-        if args.length < 3  or  args.length > 5
-            raise Puppet::ParseError, ("regsub(): wrong number of arguments" +
-                                       " (#{args.length}; min 3, max 5)")
+    do |args|
+        unless args.length.between?(3, 5)
+            raise(Puppet::ParseError,
+                  "regsubst(): got #{args.length} arguments, expected 3 to 5")
         end
-        str, regexp, replacement, flags, lang = args
+        target, regexp, replacement, flags, lang = args
         reflags = 0
-        global = false
-        (flags or "").each_byte do |f|
-            f = f.chr
-            if f == "G"
-                global = true
-            else
-                fvalue = flag_mapping[f]
-                if !fvalue
-                    raise Puppet::ParseError, "regsub(): bad flag `#{f}'"
-                end
-                reflags |= fvalue
+        operation = :sub
+        if flags == nil
+            flags = []
+        elsif flags.respond_to?(:split)
+            flags = flags.split('')
+        else
+            raise(Puppet::ParseError,
+                  "regsubst(): bad flags parameter #{flags.class}:`#{flags}'")
+        end
+        flags.each do |f|
+            case f
+            when 'G' then operation = :gsub
+            when 'E' then reflags |= Regexp::EXTENDED
+            when 'I' then reflags |= Regexp::IGNORECASE
+            when 'M' then reflags |= Regexp::MULTILINE
+            else raise(Puppet::ParseError, "regsubst(): bad flag `#{f}'")
             end
         end
-        re = Regexp.compile(regexp, reflags, lang)
-        if global
-            result = str.gsub(re, replacement)
+        begin
+            re = Regexp.compile(regexp, reflags, lang)
+        rescue RegexpError, TypeError
+            raise(Puppet::ParseError,
+                  "regsubst(): Bad regular expression `#{regexp}'")
+        end
+        if target.respond_to?(operation)
+            # String parameter -> string result
+            result = target.send(operation, re, replacement)
+        elsif target.respond_to?(:collect) and
+                target.respond_to?(:all?) and
+                target.all? { |e| e.respond_to?(operation) }
+            # Array parameter -> array result
+            result = target.collect { |e|
+                e.send(operation, re, replacement)
+            }
         else
-            result = str.sub(re, replacement)
+            raise(Puppet::ParseError,
+                  "regsubst(): bad target #{target.class}:`#{target}'")
         end
         return result
     end