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

infer/component.go

@@ -74,8 +74,7 @@ func (rc *derivedComponentController[R, I, O]) GetSchema(reg schema.RegisterDeri
 }
 
 func (rc *derivedComponentController[R, I, O]) GetToken() (tokens.Type, error) {
-	var r R
-	return introspect.GetToken("pkg", r)
+	return getToken[R](nil)
 }
 
 func (rc *derivedComponentController[R, I, O]) Construct(

infer/function.go

@@ -24,7 +24,6 @@ import (
 
 	p "github.com/pulumi/pulumi-go-provider"
 	"github.com/pulumi/pulumi-go-provider/infer/internal/ende"
-	"github.com/pulumi/pulumi-go-provider/internal/introspect"
 	t "github.com/pulumi/pulumi-go-provider/middleware"
 	"github.com/pulumi/pulumi-go-provider/middleware/schema"
 )
@@ -54,12 +53,15 @@ 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) {
-	var f F
-	tk, err := introspect.GetToken("pkg", f)
-	if err != nil {
-		return "", err
-	}
-	return fnToken(tk), nil
+	// By default, we get resource style tokens:
+	//
+	//	pkg:index:FizzBuzz
+	//
+	// Functions use a different capitalization convection, so we need to convert:
+	//
+	//	pkg:index:FizzBuzz
+	//
+	return getToken[F](fnToken)
 }
 
 func fnToken(tk tokens.Type) tokens.Type {

infer/resource.go

@@ -16,6 +16,7 @@ package infer
 
 import (
 	"fmt"
+	"reflect"
 
 	"github.com/hashicorp/go-multierror"
 	pschema "github.com/pulumi/pulumi/pkg/v3/codegen/schema"
@@ -138,6 +139,21 @@ 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)
+
+	// Set the token of the annotated type.
+	//
+	// module and name should be valid Pulumi token segments. The package name will be
+	// inferred from the provider.
+	//
+	// For example:
+	//
+	//	a.SetToken("mymodule", "MyResource")
+	//
+	// On a provider created with the name "mypkg" will have the token:
+	//
+	//	mypkg:mymodule:MyResource
+	//
+	SetToken(module, name string)
 }
 
 // Annotated is used to describe the fields of an object or a resource. Annotated can be
@@ -680,9 +696,27 @@ func (*derivedResourceController[R, I, O]) GetSchema(reg schema.RegisterDerivati
 	return r, errs.ErrorOrNil()
 }
 
-func (*derivedResourceController[R, I, O]) GetToken() (tokens.Type, error) {
+func getToken[R any](transform func(tokens.Type) tokens.Type) (tokens.Type, error) {
 	var r R
-	return introspect.GetToken("pkg", r)
+	return getTokenOf(reflect.TypeOf(r), transform)
+}
+
+func getTokenOf(t reflect.Type, transform func(tokens.Type) tokens.Type) (tokens.Type, error) {
+	annotator := getAnnotated(t)
+	if annotator.Token != "" {
+		return tokens.Type(annotator.Token), nil
+	}
+
+	tk, err := introspect.GetToken("pkg", t)
+	if transform == nil || err != nil {
+		return tk, err
+	}
+
+	return transform(tk), nil
+}
+
+func (*derivedResourceController[R, I, O]) GetToken() (tokens.Type, error) {
+	return getToken[R](nil)
 }
 
 func (*derivedResourceController[R, I, O]) getInstance() *R {

infer/schema.go

@@ -56,6 +56,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 +334,12 @@ 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())
+
+	tk, err := getTokenOf(t, nil)
 	if err != nil {
 		return schema.TypeSpec{}, true, err
 	}
+
 	return schema.TypeSpec{
 		Ref: "#/types/" + tk.String(),
 	}, true, nil

infer/tests/token_test.go

@@ -0,0 +1,198 @@
+// Copyright 2023, Pulumi Corporation.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package tests
+
+import (
+	"testing"
+
+	"github.com/blang/semver"
+	"github.com/pulumi/pulumi/sdk/v3/go/common/tokens"
+	"github.com/pulumi/pulumi/sdk/v3/go/pulumi"
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+
+	p "github.com/pulumi/pulumi-go-provider"
+	"github.com/pulumi/pulumi-go-provider/infer"
+	"github.com/pulumi/pulumi-go-provider/integration"
+)
+
+type CustomToken struct{}
+
+func (c *CustomToken) Annotate(a infer.Annotator) { a.SetToken("overwritten", "Tk") }
+
+func (*CustomToken) Create(
+	ctx p.Context, name string, inputs TokenArgs, preview bool,
+) (string, TokenResult, error) {
+	panic("unimplemented")
+}
+
+type TokenArgs struct {
+	Array []ObjectToken `pulumi:"arr"`
+
+	Single ObjectToken `pulumi:"single"`
+}
+type TokenResult struct {
+	Map map[string]ObjectToken `pulumi:"m"`
+}
+
+type TokenComponent struct{ pulumi.ResourceState }
+
+type ComponentToken struct{}
+
+// Check that we allow other capitalization schemes
+func (c *ComponentToken) Annotate(a infer.Annotator) { a.SetToken("cmp", "tK") }
+
+func (*ComponentToken) Construct(
+	ctx *pulumi.Context, name, typ string, inputs TokenArgs, opts pulumi.ResourceOption,
+) (*TokenComponent, error) {
+	panic("unimplemented")
+}
+
+type FnToken struct{}
+
+func (c *FnToken) Annotate(a infer.Annotator) { a.SetToken("fn", "TK") }
+
+func (*FnToken) Call(ctx p.Context, input TokenArgs) (output TokenResult, err error) {
+	panic("unimplemented")
+}
+
+type ObjectToken struct {
+	Value string `pulumi:"value"`
+}
+
+func (c *ObjectToken) Annotate(a infer.Annotator) { a.SetToken("obj", "Customized") }
+
+func TestTokens(t *testing.T) {
+	provider := infer.Provider(infer.Options{
+		Resources: []infer.InferredResource{
+			infer.Resource[*CustomToken, TokenArgs, TokenResult](),
+		},
+		Components: []infer.InferredComponent{
+			infer.Component[*ComponentToken, TokenArgs, *TokenComponent](),
+		},
+		Functions: []infer.InferredFunction{
+			infer.Function[*FnToken, TokenArgs, TokenResult](),
+		},
+		ModuleMap: map[tokens.ModuleName]tokens.ModuleName{"overwritten": "index"},
+	})
+	server := integration.NewServer("test", semver.MustParse("1.0.0"), provider)
+
+	schema, err := server.GetSchema(p.GetSchemaRequest{})
+	require.NoError(t, err)
+
+	assert.JSONEq(t, `{
+  "name": "test",
+  "version": "1.0.0",
+  "config": {},
+  "types": {
+    "test:obj:Customized": {
+      "properties": {
+        "value": {
+          "type": "string"
+        }
+      },
+      "type": "object",
+      "required": [
+        "value"
+      ]
+    }
+  },
+  "provider": {},
+  "resources": {
+    "test:cmp:tK": {
+      "inputProperties": {
+        "arr": {
+          "type": "array",
+          "items": {
+            "$ref": "#/types/test:obj:Customized"
+          }
+        },
+        "single": {
+          "$ref": "#/types/test:obj:Customized"
+        }
+      },
+      "requiredInputs": [
+        "arr",
+        "single"
+      ],
+      "isComponent": true
+    },
+    "test:index:Tk": {
+      "properties": {
+        "m": {
+          "type": "object",
+          "additionalProperties": {
+            "$ref": "#/types/test:obj:Customized"
+          }
+        }
+      },
+      "required": [
+        "m"
+      ],
+      "inputProperties": {
+        "arr": {
+          "type": "array",
+          "items": {
+            "$ref": "#/types/test:obj:Customized"
+          }
+        },
+        "single": {
+          "$ref": "#/types/test:obj:Customized"
+        }
+      },
+      "requiredInputs": [
+        "arr",
+        "single"
+      ]
+    }
+  },
+  "functions": {
+    "test:fn:TK": {
+      "inputs": {
+        "properties": {
+          "arr": {
+            "type": "array",
+            "items": {
+              "$ref": "#/types/test:obj:Customized"
+            }
+          },
+          "single": {
+            "$ref": "#/types/test:obj:Customized"
+          }
+        },
+        "type": "object",
+        "required": [
+          "arr",
+          "single"
+        ]
+      },
+      "outputs": {
+        "properties": {
+          "m": {
+            "type": "object",
+            "additionalProperties": {
+              "$ref": "#/types/test:obj:Customized"
+            }
+          }
+        },
+        "type": "object",
+        "required": [
+          "m"
+        ]
+      }
+    }
+  }
+}`, schema.Schema)
+}

infer/types.go

@@ -111,8 +111,10 @@ func isEnum(t reflect.Type) (enum, bool) {
 			Name:        v.FieldByName("Name").String(),
 		}
 	}
-	tk, err := introspect.GetToken("pkg", reflect.New(t).Elem().Interface())
+
+	tk, err := getTokenOf(t, nil)
 	contract.AssertNoErrorf(err, "failed to get token for enum: %s", t)
+
 	return enum{
 		token:  tk.String(),
 		values: values,
@@ -244,10 +246,12 @@ func registerTypes[T any](reg schema.RegisterDerivativeType) error {
 			if err != nil {
 				return false, err
 			}
-			tk, err := introspect.GetToken("pkg", reflect.New(t).Interface())
+
+			tk, err := getTokenOf(t, nil)
 			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,7 @@ 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(module, token string) {
+	a.Token = fmt.Sprintf("pkg:%s:%s", module, token)
+}