From 46e168a6829c95e67d04f5988dc697674c8498e1 Mon Sep 17 00:00:00 2001
From: Kristin Laemmert <mildwonkey@users.noreply.github.com>
Date: Tue, 16 Oct 2018 11:04:19 -0700
Subject: [PATCH] lang/funcs: update lookup() to accept object-typed values for
 "map" arg

---
 lang/funcs/collection.go      | 22 +++++++++++++++++++---
 lang/funcs/collection_test.go | 12 ++++++++++++
 2 files changed, 31 insertions(+), 3 deletions(-)

diff --git a/lang/funcs/collection.go b/lang/funcs/collection.go
index d436204a4..8920cef59 100644
--- a/lang/funcs/collection.go
+++ b/lang/funcs/collection.go
@@ -4,6 +4,7 @@ import (
 	"fmt"
 	"sort"
 
+	"github.com/hashicorp/hcl2/hcl"
 	"github.com/zclconf/go-cty/cty"
 	"github.com/zclconf/go-cty/cty/convert"
 	"github.com/zclconf/go-cty/cty/function"
@@ -478,7 +479,7 @@ var LookupFunc = function.New(&function.Spec{
 	Params: []function.Parameter{
 		{
 			Name: "inputMap",
-			Type: cty.Map(cty.DynamicPseudoType),
+			Type: cty.DynamicPseudoType,
 		},
 		{
 			Name: "key",
@@ -496,8 +497,13 @@ var LookupFunc = function.New(&function.Spec{
 		if len(args) < 1 || len(args) > 3 {
 			return cty.NilType, fmt.Errorf("lookup() takes two or three arguments, got %d", len(args))
 		}
-
-		return args[0].Type().ElementType(), nil
+		ty := args[0].Type()
+		switch {
+		case ty.IsObjectType():
+			return cty.DynamicPseudoType, nil
+		default:
+			return args[0].Type().ElementType(), nil
+		}
 	},
 	Impl: func(args []cty.Value, retType cty.Type) (ret cty.Value, err error) {
 		var defaultVal cty.Value
@@ -515,6 +521,16 @@ var LookupFunc = function.New(&function.Spec{
 			return cty.UnknownVal(retType), nil
 		}
 
+		if retType == cty.DynamicPseudoType {
+			// we're dealing with object-typed values
+			traversal := hcl.TraverseAttr{Name: args[1].AsString()}
+			ret, diags := traversal.TraversalStep(args[0])
+			if diags.HasErrors() {
+				return cty.UnknownVal(retType), diags
+			}
+			return ret, nil
+		}
+
 		if mapVar.HasIndex(cty.StringVal(lookupKey)) == cty.True {
 			v := mapVar.Index(cty.StringVal(lookupKey))
 			if ty := v.Type(); !ty.Equals(cty.NilType) {
diff --git a/lang/funcs/collection_test.go b/lang/funcs/collection_test.go
index 77783ce1d..4589c4c2f 100644
--- a/lang/funcs/collection_test.go
+++ b/lang/funcs/collection_test.go
@@ -1123,6 +1123,10 @@ func TestLookup(t *testing.T) {
 		"foo": cty.StringVal("bar"),
 		"baz": cty.UnknownVal(cty.String),
 	})
+	mapWithObjects := cty.ObjectVal(map[string]cty.Value{
+		"foo": cty.StringVal("bar"),
+		"baz": cty.NumberIntVal(42),
+	})
 
 	tests := []struct {
 		Values []cty.Value
@@ -1137,6 +1141,14 @@ func TestLookup(t *testing.T) {
 			cty.StringVal("bar"),
 			false,
 		},
+		{
+			[]cty.Value{
+				mapWithObjects,
+				cty.StringVal("foo"),
+			},
+			cty.StringVal("bar"),
+			false,
+		},
 		{
 			[]cty.Value{
 				intsMap,