Merge pull request #573 from Yelp/pl-dig

Add dig function

4 files changed, 115 insertions, 34 deletions

diff --git a/README.markdown b/README.markdown
index 559a6a0..c2d8388 100644
--- a/README.markdown
+++ b/README.markdown
@@ -254,6 +254,42 @@ Deletes all instances of the undef value from an array or hash. For example, `$h
 
 Returns the difference between two arrays. The returned array is a copy of the original array, removing any items that also appear in the second array. For example, `difference(["a","b","c"],["b","c","d"])` returns ["a"]. *Type*: rvalue.
 
+#### `dig`
+
+*Type*: rvalue.
+
+Retrieves a value within multiple layers of hashes and arrays via an array of keys containing a path. The function goes through the structure by each path component and tries to return the value at the end of the path.
+
+In addition to the required path argument, the function accepts the default argument. It is returned if the path is not correct, if no value was found, or if any other error has occurred.
+
+~~~ruby
+$data = {
+  'a' => {
+    'b' => [
+      'b1',
+      'b2',
+      'b3',
+    ]
+  }
+}
+
+$value = dig($data, ['a', 'b', 2])
+# $value = 'b3'
+
+# with all possible options
+$value = dig($data, ['a', 'b', 2], 'not_found')
+# $value = 'b3'
+
+# using the default value
+$value = dig($data, ['a', 'b', 'c', 'd'], 'not_found')
+# $value = 'not_found'
+~~~
+
+1. **$data** The data structure we are working with.
+2. **['a', 'b', 2]** The path array.
+3. **'not_found'** The default value. It will be returned if nothing is found.
+   (optional, defaults to *undef*)
+
 #### `dirname`
 
 Returns the `dirname` of a path. For example, `dirname('/path/to/a/file.ext')` returns '/path/to/a'. *Type*: rvalue.
@@ -806,6 +842,8 @@ Converts the argument into bytes, for example "4 kB" becomes "4096". Takes a sin
 
 *Type*: rvalue.
 
+DEPRECATED: replaced by `dig()`.
+
 Retrieves a value within multiple layers of hashes and arrays via a string containing a path. The path is a string of hash keys or array indexes starting with zero, separated by the path separator character (default "/"). The function goes through the structure by each path component and tries to return the value at the end of the path.
 
 In addition to the required path argument, the function accepts the default argument. It is returned if the path is not correct, if no value was found, or if any other error has occurred. The last argument can set the path separator character.
diff --git a/lib/puppet/parser/functions/dig.rb b/lib/puppet/parser/functions/dig.rb
new file mode 100644
index 0000000..a9aa770
--- /dev/null
+++ b/lib/puppet/parser/functions/dig.rb
@@ -0,0 +1,54 @@
+#
+# dig.rb
+#
+
+module Puppet::Parser::Functions
+  newfunction(:dig, :type => :rvalue, :doc => <<-EOS
+Looks up into a complex structure of arrays and hashes and returns nil
+or the default value if nothing was found.
+
+Path is an array of keys to be looked up in data argument. The function
+will go down the structure and try to extract the required value.
+
+$data = {
+  'a' => {
+    'b' => [
+      'b1',
+      'b2',
+      'b3' ]}}
+
+$value = dig($data, ['a', 'b', '2'], 'not_found')
+=> $value = 'b3'
+
+a -> first hash key
+b -> second hash key
+2 -> array index starting with 0
+
+not_found -> (optional) will be returned if there is no value or the path
+did not match. Defaults to nil.
+
+In addition to the required "path" argument, "dig" accepts default
+argument. It will be returned if no value was found or a path component is
+missing. And the fourth argument can set a variable path separator.
+    EOS
+             ) do |arguments|
+    # Two arguments are required
+    raise(Puppet::ParseError, "dig(): Wrong number of arguments " +
+                              "given (#{arguments.size} for at least 2)") if arguments.size < 2
+
+    data, path, default = *arguments
+
+    if !(data.is_a?(Hash) || data.is_a?(Array))
+      raise(Puppet::ParseError, "dig(): first argument must be a hash or an array, " <<
+                                "given #{data.class.name}")
+    end
+
+    unless path.is_a? Array
+      raise(Puppet::ParseError, "dig(): second argument must be an array, " <<
+                                "given #{path.class.name}")
+    end
+
+    value = path.reduce(data) { |h, k| (h.is_a?(Hash) || h.is_a?(Array)) ? h[k] : break }
+    value.nil? ? default : value
+  end
+end
diff --git a/lib/puppet/parser/functions/try_get_value.rb b/lib/puppet/parser/functions/try_get_value.rb
index 0c19fd9..fc19a23 100644
--- a/lib/puppet/parser/functions/try_get_value.rb
+++ b/lib/puppet/parser/functions/try_get_value.rb
@@ -4,6 +4,8 @@ module Puppet::Parser::Functions
       :type => :rvalue,
       :arity => -2,
       :doc => <<-eos
+DEPRECATED: this function is deprecated, please use dig() instead.
+
 Looks up into a complex structure of arrays and hashes and returns a value
 or the default value if nothing was found.
 
@@ -35,43 +37,17 @@ argument. It will be returned if no value was found or a path component is
 missing. And the fourth argument can set a variable path separator.
   eos
   ) do |args|
-    path_lookup = lambda do |data, path, default|
-      debug "Try_get_value: #{path.inspect} from: #{data.inspect}"
-      if data.nil?
-        debug "Try_get_value: no data, return default: #{default.inspect}"
-        break default
-      end
-      unless path.is_a? Array
-        debug "Try_get_value: wrong path, return default: #{default.inspect}"
-        break default
-      end
-      unless path.any?
-        debug "Try_get_value: value found, return data: #{data.inspect}"
-        break data
-      end
-      unless data.is_a? Hash or data.is_a? Array
-        debug "Try_get_value: incorrect data, return default: #{default.inspect}"
-        break default
-      end
-
-      key = path.shift
-      if data.is_a? Array
-        begin
-          key = Integer key
-        rescue ArgumentError
-          debug "Try_get_value: non-numeric path for an array, return default: #{default.inspect}"
-          break default
-        end
-      end
-      path_lookup.call data[key], path, default
-    end
-
+    warning("try_get_value() DEPRECATED: this function is deprecated, please use dig() instead.")
     data = args[0]
     path = args[1] || ''
     default = args[2]
-    separator = args[3] || '/'
 
-    path = path.split separator
-    path_lookup.call data, path, default
+    if !(data.is_a?(Hash) || data.is_a?(Array)) || path == ''
+      return default || data
+    end
+
+    separator = args[3] || '/'
+    path = path.split(separator).map{ |key| key =~ /^\d+$/ ? key.to_i : key }
+    function_dig([data, path, default])
   end
 end
diff --git a/spec/functions/dig_spec.rb b/spec/functions/dig_spec.rb
new file mode 100644
index 0000000..1c5d49d
--- /dev/null
+++ b/spec/functions/dig_spec.rb
@@ -0,0 +1,13 @@
+require 'spec_helper'
+
+describe 'dig' do
+  it { is_expected.to run.with_params().and_raise_error(Puppet::ParseError) }
+  it { is_expected.to run.with_params('bad', []).and_raise_error(Puppet::ParseError) }
+  it { is_expected.to run.with_params({}, 'bad').and_raise_error(Puppet::ParseError) }
+
+  it { is_expected.to run.with_params({}, []).and_return({}) }
+  it { is_expected.to run.with_params({"a" => "b"}, ["a"]).and_return("b") }
+  it { is_expected.to run.with_params({"a" => {"b" => "c"}}, ["a", "b"]).and_return("c") }
+  it { is_expected.to run.with_params({}, ["a", "b"], "d").and_return("d") }
+  it { is_expected.to run.with_params({"a" => false}, ["a"]).and_return(false) }
+end