From 5d325c83e74e3e4af2dd78b8e337da170e3b1ee4 Mon Sep 17 00:00:00 2001
From: Nicholas Wiersma <nick@wiersma.co.za>
Date: Thu, 12 Sep 2024 08:15:08 +0200
Subject: [PATCH] feat: support name customization

---
 schemagen/generator.go      | 16 +++++++++++++++-
 schemagen/generator_test.go | 25 +++++++++++++++++++++++++
 2 files changed, 40 insertions(+), 1 deletion(-)

diff --git a/schemagen/generator.go b/schemagen/generator.go
index 1c8c792..41786e4 100644
--- a/schemagen/generator.go
+++ b/schemagen/generator.go
@@ -23,6 +23,9 @@ type DocsFunc = func(obj any, sf reflect.StructField) string
 // if the field should be skipped.
 type CustomizerFunc = func(obj any, sf *reflect.StructField, typ reflect.Type, s *schema.Schema) (string, reflect.Kind, bool)
 
+// NameFunc is a function used to determine a field name.
+type NameFunc = func(name string) string
+
 var errSkip = errors.New("skip")
 
 // Generator generates Terraform schemas.
@@ -30,10 +33,16 @@ type Generator struct {
 	docsFn       DocsFunc
 	customizerFn CustomizerFunc
 	tag          string
+	nameFn       NameFunc
 }
 
 // New returns a schema generator.
 func New(docsFn DocsFunc, customizerFn CustomizerFunc, tag string) *Generator {
+	return NewWithName(docsFn, customizerFn, nil, tag)
+}
+
+// NewWithName returns a schema generator with the given nameFn.
+func NewWithName(docsFn DocsFunc, customizerFn CustomizerFunc, nameFn NameFunc, tag string) *Generator {
 	if tag == "" {
 		tag = "json"
 	}
@@ -45,10 +54,14 @@ func New(docsFn DocsFunc, customizerFn CustomizerFunc, tag string) *Generator {
 			return "", typ.Kind(), true
 		}
 	}
+	if nameFn == nil {
+		nameFn = func(name string) string { return strcase.ToSnake(name) }
+	}
 
 	return &Generator{
 		docsFn:       docsFn,
 		customizerFn: customizerFn,
+		nameFn:       nameFn,
 		tag:          tag,
 	}
 }
@@ -78,7 +91,8 @@ func (g *Generator) Struct(obj any) (map[string]string, error) {
 			}
 		}
 
-		fields[strcase.ToSnake(name)] = code
+		fieldName := g.nameFn(name)
+		fields[fieldName] = code
 	}
 	return fields, nil
 }
diff --git a/schemagen/generator_test.go b/schemagen/generator_test.go
index 2b8477f..3a27a14 100644
--- a/schemagen/generator_test.go
+++ b/schemagen/generator_test.go
@@ -4,6 +4,7 @@ import (
 	"reflect"
 	"testing"
 
+	"github.com/ettle/strcase"
 	"github.com/hashicorp/terraform-plugin-sdk/v2/helper/schema"
 	"github.com/nitrado/tfconv/schemagen"
 	"github.com/stretchr/testify/assert"
@@ -29,6 +30,30 @@ func TestGenerator_Struct(t *testing.T) {
 	assert.Equal(t, want, got)
 }
 
+func TestGenerator_StructWithName(t *testing.T) {
+	gen := schemagen.NewWithName(nil, nil, func(name string) string {
+		if name == "float" {
+			return "other"
+		}
+		return strcase.ToSnake(name)
+	}, "")
+
+	got, err := gen.Struct(&TestObject{})
+
+	require.NoError(t, err)
+	want := map[string]string{
+		"other":    "{\nType: schema.TypeFloat,\n}",
+		"int":      "{\nType: schema.TypeInt,\n}",
+		"str":      "{\nType: schema.TypeString,\n}",
+		"bool":     "{\nType: schema.TypeBool,\n}",
+		"ptr_bool": "{\nType: schema.TypeList,\nMaxItems: 1,\nElem: &schema.Resource{\nSchema: map[string]*schema.Schema{\n\"value\": {\nType: schema.TypeBool,\nRequired: true,\n},\n},\n},\n}",
+		"slice":    "{\nType: schema.TypeList,\nElem: &schema.Resource{\nSchema: map[string]*schema.Schema{\n\"a\": {\nType: schema.TypeString,\n},\n},\n},\n}",
+		"map":      "{\nType: schema.TypeMap,\nElem: &schema.Schema{Type: schema.TypeInt,},\n}",
+		"struct":   "{\nType: schema.TypeList,\nMaxItems: 1,\nElem: &schema.Resource{\nSchema: map[string]*schema.Schema{\n\"a\": {\nType: schema.TypeString,\n},\n},\n},\n}",
+	}
+	assert.Equal(t, want, got)
+}
+
 func TestGenerator_StructUsesCustomFunction(t *testing.T) {
 	c := func(_ any, _ *reflect.StructField, typ reflect.Type, _ *schema.Schema) (string, reflect.Kind, bool) {
 		if typ == reflect.TypeOf(T{}) {