functions: LookupFunc

lang/funcs/collection.go

@@ -385,7 +385,6 @@ var KeysFunc = function.New(&function.Spec{
 
 		for it := args[0].ElementIterator(); it.Next(); {
 			k, _ := it.Element()
-			fmt.Printf("appending %#v to %#v\n", k, keys)
 			keys = append(keys, k)
 			if err != nil {
 				return cty.ListValEmpty(cty.String), err
@@ -439,6 +438,64 @@ var ListFunc = function.New(&function.Spec{
 	},
 })
 
+// LookupFunc contructs a function that performs dynamic lookups of map types.
+var LookupFunc = function.New(&function.Spec{
+	Params: []function.Parameter{
+		{
+			Name: "inputMap",
+			Type: cty.Map(cty.DynamicPseudoType),
+		},
+		{
+			Name: "key",
+			Type: cty.String,
+		},
+	},
+	VarParam: &function.Parameter{
+		Name:             "default",
+		Type:             cty.DynamicPseudoType,
+		AllowUnknown:     true,
+		AllowDynamicType: true,
+		AllowNull:        true,
+	},
+	Type: function.StaticReturnType(cty.DynamicPseudoType),
+	Impl: func(args []cty.Value, retType cty.Type) (ret cty.Value, err error) {
+		if len(args) < 1 || len(args) > 3 {
+			return cty.NilVal, fmt.Errorf("lookup() takes two or three arguments, got %d", len(args))
+		}
+		var defaultVal string
+		defaultValueSet := false
+
+		if len(args) == 3 {
+			defaultVal = args[2].AsString()
+			defaultValueSet = true
+		}
+
+		mapVar := args[0]
+		lookupKey := args[1].AsString()
+
+		if mapVar.HasIndex(cty.StringVal(lookupKey)) == cty.True {
+			v := mapVar.Index(cty.StringVal(lookupKey))
+			if ty := v.Type(); !ty.Equals(cty.NilType) {
+				switch {
+				case ty.Equals(cty.String):
+					return cty.StringVal(v.AsString()), nil
+				case ty.Equals(cty.Number):
+					return cty.NumberVal(v.AsBigFloat()), nil
+				default:
+					return cty.NilVal, fmt.Errorf("lookup() can only be used with flat lists")
+				}
+			}
+		}
+
+		if defaultValueSet {
+			return cty.StringVal(defaultVal), nil
+		}
+
+		return cty.UnknownVal(cty.String), fmt.Errorf(
+			"lookup failed to find '%s'", lookupKey)
+	},
+})
+
 // MapFunc contructs a function that takes an even number of arguments and
 // returns a map whose elements are constructed from consecutive pairs of arguments.
 //
@@ -653,6 +710,13 @@ func List(args ...cty.Value) (cty.Value, error) {
 	return ListFunc.Call(args)
 }
 
+// Lookup performs a dynamic lookup into a map.
+// There are two required arguments, map and key, plus an optional default,
+// which is a value to return if no key is found in map.
+func Lookup(args ...cty.Value) (cty.Value, error) {
+	return LookupFunc.Call(args)
+}
+
 // Map takes an even number of arguments and returns a map whose elements are constructed
 // from consecutive pairs of arguments.
 func Map(args ...cty.Value) (cty.Value, error) {

lang/funcs/collection_test.go

@@ -949,6 +949,116 @@ func TestList(t *testing.T) {
 	}
 }
 
+func TestLookup(t *testing.T) {
+	simpleMap := cty.MapVal(map[string]cty.Value{
+		"foo": cty.StringVal("bar"),
+	})
+	intsMap := cty.MapVal(map[string]cty.Value{
+		"foo": cty.NumberIntVal(42),
+	})
+	mapOfLists := cty.MapVal(map[string]cty.Value{
+		"foo": cty.ListVal([]cty.Value{
+			cty.StringVal("bar"),
+			cty.StringVal("baz"),
+		}),
+	})
+
+	tests := []struct {
+		Values []cty.Value
+		Want   cty.Value
+		Err    bool
+	}{
+		{
+			[]cty.Value{
+				simpleMap,
+				cty.StringVal("foo"),
+			},
+			cty.StringVal("bar"),
+			false,
+		},
+		{
+			[]cty.Value{
+				intsMap,
+				cty.StringVal("foo"),
+			},
+			cty.NumberIntVal(42),
+			false,
+		},
+		{ // Invalid key
+			[]cty.Value{
+				simpleMap,
+				cty.StringVal("bar"),
+			},
+			cty.NilVal,
+			true,
+		},
+		{ // Supplied default with valid key
+			[]cty.Value{
+				simpleMap,
+				cty.StringVal("foo"),
+				cty.StringVal(""),
+			},
+			cty.StringVal("bar"),
+			false,
+		},
+		{ // Supplied default with invalid key
+			[]cty.Value{
+				simpleMap,
+				cty.StringVal("baz"),
+				cty.StringVal(""),
+			},
+			cty.StringVal(""),
+			false,
+		},
+		{ // Supplied non-empty default with invalid key
+			[]cty.Value{
+				simpleMap,
+				cty.StringVal("bar"),
+				cty.StringVal("xyz"),
+			},
+			cty.StringVal("xyz"),
+			false,
+		},
+		{ // too many args
+			[]cty.Value{
+				simpleMap,
+				cty.StringVal("foo"),
+				cty.StringVal("bar"),
+				cty.StringVal("baz"),
+			},
+			cty.NilVal,
+			true,
+		},
+		{ // cannot search a map of lists
+			[]cty.Value{
+				mapOfLists,
+				cty.StringVal("baz"),
+			},
+			cty.NilVal,
+			true,
+		},
+	}
+
+	for _, test := range tests {
+		t.Run(fmt.Sprintf("lookup(%#v)", test.Values), func(t *testing.T) {
+			got, err := Lookup(test.Values...)
+
+			if test.Err {
+				if err == nil {
+					t.Fatal("succeeded; want error")
+				}
+				return
+			} else if err != nil {
+				t.Fatalf("unexpected error: %s", err)
+			}
+
+			if !got.RawEquals(test.Want) {
+				t.Errorf("wrong result\ngot:  %#v\nwant: %#v", got, test.Want)
+			}
+		})
+	}
+}
+
 func TestMap(t *testing.T) {
 	tests := []struct {
 		Values []cty.Value

lang/functions.go

@@ -67,7 +67,7 @@ func (s *Scope) Functions() map[string]function.Function {
 			"length":       funcs.LengthFunc,
 			"list":         funcs.ListFunc,
 			"log":          funcs.LogFunc,
-			"lookup":       unimplFunc, // TODO
+			"lookup":       funcs.LookupFunc,
 			"lower":        stdlib.LowerFunc,
 			"map":          funcs.MapFunc,
 			"matchkeys":    funcs.MatchkeysFunc,