Allow walking types

Author: t0yv0

pkg/valueshim/cty.go

@@ -149,6 +149,25 @@ func (t ctyTypeShim) IsObjectType() bool {
 	return t.ty().IsObjectType()
 }
 
+func (t ctyTypeShim) AttributeType(name string) (Type, bool) {
+	tt := t.ty()
+	if !tt.IsObjectType() {
+		return nil, false
+	}
+	if !tt.HasAttribute(name) {
+		return nil, false
+	}
+	return FromHCtyType(tt.AttributeType(name)), true
+}
+
+func (t ctyTypeShim) ElementType() (Type, bool) {
+	tt := t.ty()
+	if !tt.IsCollectionType() {
+		return nil, false
+	}
+	return FromHCtyType(tt.ElementType()), true
+}
+
 func (t ctyTypeShim) GoString() string {
 	return t.ty().GoString()
 }

pkg/valueshim/cty_test.go

@@ -314,3 +314,31 @@ func Test_HCty_ToX(t *testing.T) {
 	assert.Equal(t, 42.41, valueshim.FromHCtyValue(cty.NumberFloatVal(42.41)).NumberValue())
 	assert.Equal(t, true, valueshim.FromHCtyValue(cty.BoolVal(true)).BoolValue())
 }
+
+func Test_HCtyType_AttributeType(t *testing.T) {
+	objTy := cty.Object(map[string]cty.Type{"x": cty.String})
+
+	ty, ok := valueshim.FromHCtyType(objTy).AttributeType("x")
+	assert.True(t, ok)
+	assert.Equal(t, valueshim.FromHCtyType(cty.String), ty)
+
+	_, ok = valueshim.FromHCtyType(objTy).AttributeType("y")
+	assert.False(t, ok)
+}
+
+func Test_HCtyType_ElementType(t *testing.T) {
+	ty, ok := valueshim.FromHCtyType(cty.Set(cty.Number)).ElementType()
+	assert.True(t, ok)
+	assert.Equal(t, valueshim.FromHCtyType(cty.Number), ty)
+
+	ty, ok = valueshim.FromHCtyType(cty.List(cty.Number)).ElementType()
+	assert.True(t, ok)
+	assert.Equal(t, valueshim.FromHCtyType(cty.Number), ty)
+
+	ty, ok = valueshim.FromHCtyType(cty.Map(cty.Number)).ElementType()
+	assert.True(t, ok)
+	assert.Equal(t, valueshim.FromHCtyType(cty.Number), ty)
+
+	ty, ok = valueshim.FromHCtyType(cty.String).ElementType()
+	assert.False(t, ok)
+}

pkg/valueshim/shim.go

@@ -55,5 +55,7 @@ type Type interface {
 	IsMapType() bool
 	IsSetType() bool
 	IsObjectType() bool
+	AttributeType(name string) (Type, bool)
+	ElementType() (Type, bool)
 	GoString() string
 }

pkg/valueshim/tfvalue.go

@@ -190,3 +190,28 @@ func (t tTypeShim) IsObjectType() bool {
 func (t tTypeShim) GoString() string {
 	return t.ty().String()
 }
+
+func (t tTypeShim) AttributeType(name string) (Type, bool) {
+	if !t.IsObjectType() {
+		return nil, false
+	}
+	tt := t.ty()
+	attr, ok := tt.(tftypes.Object).AttributeTypes[name]
+	if !ok {
+		return nil, false
+	}
+	return FromTType(attr), true
+}
+
+func (t tTypeShim) ElementType() (Type, bool) {
+	tt := t.ty()
+	switch {
+	case tt.Is(tftypes.Map{}):
+		return FromTType(tt.(tftypes.Map).ElementType), true
+	case tt.Is(tftypes.List{}):
+		return FromTType(tt.(tftypes.List).ElementType), true
+	case tt.Is(tftypes.Set{}):
+		return FromTType(tt.(tftypes.Set).ElementType), true
+	}
+	return nil, false
+}

pkg/valueshim/tfvalue_test.go

@@ -313,6 +313,34 @@ func Test_TType(t *testing.T) {
 	assert.Equal(t, "tftypes.Object[]", valueshim.FromTType(tftypes.Object{}).GoString())
 }
 
+func Test_TType_AttributeType(t *testing.T) {
+	objTy := tftypes.Object{AttributeTypes: map[string]tftypes.Type{"x": tftypes.String}}
+
+	ty, ok := valueshim.FromTType(objTy).AttributeType("x")
+	assert.True(t, ok)
+	assert.Equal(t, valueshim.FromTType(tftypes.String), ty)
+
+	_, ok = valueshim.FromTType(objTy).AttributeType("y")
+	assert.False(t, ok)
+}
+
+func Test_TType_ElementType(t *testing.T) {
+	ty, ok := valueshim.FromTType(tftypes.Set{ElementType: tftypes.Number}).ElementType()
+	assert.True(t, ok)
+	assert.Equal(t, valueshim.FromTType(tftypes.Number), ty)
+
+	ty, ok = valueshim.FromTType(tftypes.List{ElementType: tftypes.Number}).ElementType()
+	assert.True(t, ok)
+	assert.Equal(t, valueshim.FromTType(tftypes.Number), ty)
+
+	ty, ok = valueshim.FromTType(tftypes.Map{ElementType: tftypes.Number}).ElementType()
+	assert.True(t, ok)
+	assert.Equal(t, valueshim.FromTType(tftypes.Number), ty)
+
+	ty, ok = valueshim.FromTType(tftypes.String).ElementType()
+	assert.False(t, ok)
+}
+
 func Test_TValue_ToX(t *testing.T) {
 	t.Parallel()