Allow override a Pulumi token by using SetToken of infer.Annotator

infer/component.go

@@ -16,6 +16,7 @@ package infer
 
 import (
 	"fmt"
+	"reflect"
 
 	pschema "github.com/pulumi/pulumi/pkg/v3/codegen/schema"
 	"github.com/pulumi/pulumi/sdk/v3/go/common/tokens"
@@ -75,6 +76,11 @@ func (rc *derivedComponentController[R, I, O]) GetSchema(reg schema.RegisterDeri
 
 func (rc *derivedComponentController[R, I, O]) GetToken() (tokens.Type, error) {
 	var r R
+	annotator := getAnnotated(reflect.TypeOf(r))
+	if annotator.Token != "" {
+		return fnToken(tokens.Type(annotator.Token)), nil
+	}
+
 	return introspect.GetToken("pkg", r)
 }
 

infer/function.go

@@ -53,8 +53,12 @@ type derivedInvokeController[F Fn[I, O], I, O any] struct{}
 
 func (derivedInvokeController[F, I, O]) isInferredFunction() {}
 
-func (*derivedInvokeController[F, I, O]) GetToken() (tokens.Type, error) {
+func (c *derivedInvokeController[F, I, O]) GetToken() (tokens.Type, error) {
 	var f F
+	annotator := getAnnotated(reflect.TypeOf(f))
+	if annotator.Token != "" {
+		return fnToken(tokens.Type(annotator.Token)), nil
+	}
 	tk, err := introspect.GetToken("pkg", f)
 	if err != nil {
 		return "", err

infer/resource.go

@@ -138,6 +138,8 @@ type Annotator interface {
 	// Annotate a struct field with a default value. The default value must be a primitive
 	// type in the pulumi type system.
 	SetDefault(i any, defaultValue any, env ...string)
+
+	SetToken(i any, token string)
 }
 
 // Annotated is used to describe the fields of an object or a resource. Annotated can be
@@ -682,6 +684,11 @@ func (*derivedResourceController[R, I, O]) GetSchema(reg schema.RegisterDerivati
 
 func (*derivedResourceController[R, I, O]) GetToken() (tokens.Type, error) {
 	var r R
+	annotator := getAnnotated(reflect.TypeOf(r))
+	if annotator.Token != "" {
+		return fnToken(tokens.Type(annotator.Token)), nil
+	}
+
 	return introspect.GetToken("pkg", r)
 }
 

infer/schema.go

@@ -22,6 +22,7 @@ import (
 	"github.com/hashicorp/go-multierror"
 	"github.com/pulumi/pulumi/pkg/v3/codegen/schema"
 	"github.com/pulumi/pulumi/sdk/v3/go/common/resource"
+	"github.com/pulumi/pulumi/sdk/v3/go/common/tokens"
 	"github.com/pulumi/pulumi/sdk/v3/go/pulumi"
 
 	"github.com/pulumi/pulumi-go-provider/internal/introspect"
@@ -56,6 +57,7 @@ func getAnnotated(t reflect.Type) introspect.Annotator {
 		for k, v := range src.DefaultEnvs {
 			(*dst).DefaultEnvs[k] = v
 		}
+		dst.Token = src.Token
 	}
 
 	ret := introspect.Annotator{
@@ -333,10 +335,19 @@ func structReferenceToken(t reflect.Type, extTag *introspect.ExplicitType) (sche
 		t.Implements(reflect.TypeOf(new(pulumi.Output)).Elem()) {
 		return schema.TypeSpec{}, false, nil
 	}
-	tk, err := introspect.GetToken("pkg", reflect.New(t).Elem().Interface())
-	if err != nil {
-		return schema.TypeSpec{}, true, err
+
+	var tk tokens.Type
+	annotator := getAnnotated(t)
+	if annotator.Token != "" {
+		tk = fnToken(tokens.Type(annotator.Token))
+	} else {
+		var err error
+		tk, err = introspect.GetToken("pkg", reflect.New(t).Interface())
+		if err != nil {
+			return schema.TypeSpec{}, true, err
+		}
 	}
+
 	return schema.TypeSpec{
 		Ref: "#/types/" + tk.String(),
 	}, true, nil

infer/types.go

@@ -111,8 +111,17 @@ func isEnum(t reflect.Type) (enum, bool) {
 			Name:        v.FieldByName("Name").String(),
 		}
 	}
-	tk, err := introspect.GetToken("pkg", reflect.New(t).Elem().Interface())
-	contract.AssertNoErrorf(err, "failed to get token for enum: %s", t)
+
+	var tk tokens.Type
+	annotator := getAnnotated(t)
+	if annotator.Token != "" {
+		tk = fnToken(tokens.Type(annotator.Token))
+	} else {
+		var err error
+		tk, err = introspect.GetToken("pkg", reflect.New(t).Elem().Interface())
+		contract.AssertNoErrorf(err, "failed to get token for enum: %s", t)
+	}
+
 	return enum{
 		token:  tk.String(),
 		values: values,
@@ -244,10 +253,18 @@ func registerTypes[T any](reg schema.RegisterDerivativeType) error {
 			if err != nil {
 				return false, err
 			}
-			tk, err := introspect.GetToken("pkg", reflect.New(t).Interface())
-			if err != nil {
-				return false, err
+
+			var tk tokens.Type
+			annotator := getAnnotated(t)
+			if annotator.Token != "" {
+				tk = fnToken(tokens.Type(annotator.Token))
+			} else {
+				tk, err = introspect.GetToken("pkg", reflect.New(t).Interface())
+				if err != nil {
+					return false, err
+				}
 			}
+
 			return reg(tk, pschema.ComplexTypeSpec{ObjectTypeSpec: *spec}), nil
 		}
 		return true, nil

internal/introspect/annotator.go

@@ -33,6 +33,7 @@ type Annotator struct {
 	Descriptions map[string]string
 	Defaults     map[string]any
 	DefaultEnvs  map[string][]string
+	Token        string
 
 	matcher FieldMatcher
 }
@@ -83,3 +84,21 @@ func (a *Annotator) SetDefault(i any, defaultValue any, env ...string) {
 	a.Defaults[field.Name] = defaultValue
 	a.DefaultEnvs[field.Name] = append(a.DefaultEnvs[field.Name], env...)
 }
+
+func (a *Annotator) SetToken(i any, token string) {
+	_, ok, err := a.matcher.GetField(i)
+	if err != nil {
+		panic(fmt.Sprintf("Could not parse field tags: %s", err.Error()))
+	}
+	if !ok {
+		typ := reflect.TypeOf(i)
+		if typ.Kind() == reflect.Pointer && typ.Elem().Kind() == reflect.Pointer {
+			i = reflect.ValueOf(i).Elem().Interface()
+		}
+		if a.matcher.value.Addr().Interface() != i {
+			panic("A token can only be specified for a struct")
+		}
+	}
+
+	a.Token = token
+}

internal/introspect/introspect_test.go

@@ -36,6 +36,7 @@ func (m *MyStruct) Annotate(a infer.Annotator) {
 	a.Describe(&m, "This is MyStruct, but also your struct.")
 	a.Describe(&m.Fizz, "Fizz is not MyStruct.Foo.")
 	a.SetDefault(&m.Foo, "Fizz")
+	a.SetToken(&m, "MyToken")
 }
 
 func TestParseTag(t *testing.T) {
@@ -107,6 +108,7 @@ func TestAnnotate(t *testing.T) {
 	assert.Equal(t, "Fizz", a.Defaults["foo"])
 	assert.Equal(t, "Fizz is not MyStruct.Foo.", a.Descriptions["fizz"])
 	assert.Equal(t, "This is MyStruct, but also your struct.", a.Descriptions[""])
+	assert.Equal(t, "MyToken", a.Token)
 }
 
 func TestAllFields(t *testing.T) {