customtype: add ability to register any custom type (#27)

* customtype: add ability to register any custom type

* make customtype internal, add valast.RegisterType

* make reflect.Array support custom type (#1)

---------

Co-authored-by: Joe Chen <jc@unknwon.io>

README.md

@@ -30,6 +30,7 @@ This can be useful for debugging and testing, you may think of it as a more comp
 - Fully handles unexported fields, types, and values (optional.)
 - Strong emphasis on being used for producing valid Go code that can be copy & pasted directly into e.g. tests.
 - [Extensively tested](https://github.com/hexops/valast/tree/main/testdata), over 88 tests and handling numerous edge cases (such as pointers to unaddressable literal values like `&"foo"` properly, and even [finding bugs in alternative packages'](https://github.com/shurcooL/go-goon/issues/15)).
+- Provide custom AST representations for your types with `valast.RegisterType(...)`.
 
 ## Alternatives comparison
 

internal/customtype/customtype.go

@@ -0,0 +1,35 @@
+package customtype
+
+import (
+	"fmt"
+	"go/ast"
+	"reflect"
+	"sync"
+)
+
+var (
+	customTypesMux sync.Mutex
+	customTypes    = make(map[reflect.Type]func(any) ast.Expr)
+)
+
+// See the top-level valast.RegisterType docstring for details.
+func Register[T any](render func(value T) ast.Expr) {
+	customTypesMux.Lock()
+	var zero T
+	t := reflect.TypeOf(zero)
+	if _, exists := customTypes[t]; exists {
+		panic(fmt.Sprintf("%T already registered", zero))
+	}
+	customTypes[t] = func(value any) ast.Expr { return render(value.(T)) }
+	customTypesMux.Unlock()
+}
+
+// Is indicates if the given reflect.Type has a custom AST representation
+// generator registered.
+func Is(rt reflect.Type) (func(any) ast.Expr, bool) {
+	customTypesMux.Lock()
+	defer customTypesMux.Unlock()
+
+	t, ok := customTypes[rt]
+	return t, ok
+}

stdtypes.go

@@ -0,0 +1,37 @@
+package valast
+
+import (
+	"fmt"
+	"go/ast"
+	"go/token"
+	"time"
+)
+
+// Register custom reprsentations of common structs from stdlib that only
+// contain unexported fields.
+func init() {
+	// For time.Time, returns the AST expression equivalent of:
+	//
+	//	time.Date(2000, 1, 1, 0, 0, 0, 0, time.UTC)
+	RegisterType(func(t time.Time) ast.Expr {
+		return &ast.CallExpr{
+			Fun: &ast.SelectorExpr{
+				X:   &ast.Ident{Name: "time"},
+				Sel: &ast.Ident{Name: "Date"},
+			},
+			Args: []ast.Expr{
+				&ast.BasicLit{Kind: token.INT, Value: fmt.Sprintf("%d", t.Year())},
+				&ast.BasicLit{Kind: token.INT, Value: fmt.Sprintf("%d", t.Month())},
+				&ast.BasicLit{Kind: token.INT, Value: fmt.Sprintf("%d", t.Day())},
+				&ast.BasicLit{Kind: token.INT, Value: fmt.Sprintf("%d", t.Hour())},
+				&ast.BasicLit{Kind: token.INT, Value: fmt.Sprintf("%d", t.Minute())},
+				&ast.BasicLit{Kind: token.INT, Value: fmt.Sprintf("%d", t.Second())},
+				&ast.BasicLit{Kind: token.INT, Value: fmt.Sprintf("%d", t.Nanosecond())},
+				&ast.SelectorExpr{
+					X:   &ast.Ident{Name: "time"},
+					Sel: &ast.Ident{Name: t.Location().String()},
+				},
+			},
+		}
+	})
+}

types.go

@@ -5,8 +5,22 @@ import (
 	"go/ast"
 	"go/token"
 	"reflect"
+
+	"github.com/hexops/valast/internal/customtype"
 )
 
+// RegisterType registers a type that for representation in a custom manner with
+// valast. If valast encounters a value or pointer to a value of this type, it
+// will use the given render func to generate the appropriate AST representation.
+//
+// This is useful if a type's fields are private, and can only be represented
+// through a constructor - see stdtypes.go for examples.
+//
+// This mechanism currently only works with struct types.
+func RegisterType[T any](render func(value T) ast.Expr) {
+	customtype.Register(render)
+}
+
 type cacheKeyOptions struct {
 	Unqualify    bool
 	PackagePath  string

valast.go

@@ -16,6 +16,7 @@ import (
 	"time"
 
 	"github.com/hexops/valast/internal/bypass"
+	"github.com/hexops/valast/internal/customtype"
 	"golang.org/x/tools/go/packages"
 	gofumpt "mvdan.cc/gofumpt/format"
 )
@@ -360,6 +361,12 @@ func computeAST(v reflect.Value, opt *Options, cycleDetector *cycleDetector, pro
 	case reflect.Complex128:
 		return basicLit(vv, token.FLOAT, "complex128", v, opt, typeExprCache)
 	case reflect.Array:
+		if render, ok := customtype.Is(v.Type()); ok {
+			return Result{
+				AST: render(v.Interface()),
+			}, nil
+		}
+
 		var (
 			elts               []ast.Expr
 			requiresUnexported bool
@@ -559,8 +566,8 @@ func computeAST(v reflect.Value, opt *Options, cycleDetector *cycleDetector, pro
 				OmittedUnexported:  elem.OmittedUnexported,
 			}, nil
 		}
-		switch vv.Elem().Type() {
-		case reflect.TypeOf(time.Time{}):
+		// Wrap custom type representations in generic pointer.
+		if _, ok := customtype.Is(vv.Elem().Type()); ok {
 			return Result{
 				AST: pointifyASTExpr(elem.AST),
 			}, nil
@@ -608,12 +615,9 @@ func computeAST(v reflect.Value, opt *Options, cycleDetector *cycleDetector, pro
 		}
 		return basicLit(vv, token.STRING, "string", strconv.Quote(v.String()), opt.withUnqualify(), typeExprCache)
 	case reflect.Struct:
-		// special handling for common structs from stdlib
-		// that only contain unexported fields
-		switch v.Type() {
-		case reflect.TypeOf(time.Time{}):
+		if render, ok := customtype.Is(v.Type()); ok {
 			return Result{
-				AST: timeTypeASTExpr(v.Interface().(time.Time)),
+				AST: render(v.Interface()),
 			}, nil
 		}
 
@@ -721,31 +725,6 @@ func unexported(v reflect.Value) reflect.Value {
 	return bypass.UnsafeReflectValue(v)
 }
 
-// timeTypeASTExpr returns the AST expression equivalent of
-//
-// 	time.Date(2000, 1, 1, 0, 0, 0, 0, time.UTC)
-func timeTypeASTExpr(t time.Time) ast.Expr {
-	return &ast.CallExpr{
-		Fun: &ast.SelectorExpr{
-			X:   &ast.Ident{Name: "time"},
-			Sel: &ast.Ident{Name: "Date"},
-		},
-		Args: []ast.Expr{
-			&ast.BasicLit{Kind: token.INT, Value: fmt.Sprintf("%d", t.Year())},
-			&ast.BasicLit{Kind: token.INT, Value: fmt.Sprintf("%d", t.Month())},
-			&ast.BasicLit{Kind: token.INT, Value: fmt.Sprintf("%d", t.Day())},
-			&ast.BasicLit{Kind: token.INT, Value: fmt.Sprintf("%d", t.Hour())},
-			&ast.BasicLit{Kind: token.INT, Value: fmt.Sprintf("%d", t.Minute())},
-			&ast.BasicLit{Kind: token.INT, Value: fmt.Sprintf("%d", t.Second())},
-			&ast.BasicLit{Kind: token.INT, Value: fmt.Sprintf("%d", t.Nanosecond())},
-			&ast.SelectorExpr{
-				X:   &ast.Ident{Name: "time"},
-				Sel: &ast.Ident{Name: t.Location().String()},
-			},
-		},
-	}
-}
-
 // pointifyASTExpr wraps an expression in a call to the `Ptr` helper function.
 //
 //	valast.Ptr(//...)