kinda working

commands/resource.go

@@ -112,7 +112,8 @@ func (d ResourceData) generate() error {
 		return err
 	}
 
-	_ = helpers.UpdateRegistration(config.ServicePackagesPath+"/"+d.ServicePackage, d.Name, helpers.TypeResource, helpers.Register, d.Typed)
+	//_ = helpers.UpdateRegistration(config.ServicePackagesPath+"/"+d.ServicePackage, d.Name, helpers.TypeResource, helpers.Register, d.Typed)
+	_ = helpers.UpdateRegistrationByNode(config.ServicePackagesPath+"/"+d.ServicePackage, d.Name, helpers.TypeResource, helpers.Register, d.Typed)
 
 	return nil
 }

go.mod

@@ -7,6 +7,7 @@ require (
 	github.com/hashicorp/hcl/v2 v2.10.1
 	github.com/iancoleman/strcase v0.2.0
 	github.com/mitchellh/cli v1.1.2
+	golang.org/x/tools v0.0.0-20180917221912-90fa682c2a6e
 )
 
 require (

go.sum

@@ -163,6 +163,7 @@ golang.org/x/text v0.3.2/go.mod h1:bEr9sfX3Q8Zfm5fL9x+3itogRgK3+ptLWKqgva+5dAk=
 golang.org/x/text v0.3.3/go.mod h1:5Zoc/QRtKVWzQhOtBMvqHzDpF6irO9z98xDceosuGiQ=
 golang.org/x/text v0.3.5 h1:i6eZZ+zk0SOf0xgBpEpPD18qWcJda6q1sxt3S0kzyUQ=
 golang.org/x/text v0.3.5/go.mod h1:5Zoc/QRtKVWzQhOtBMvqHzDpF6irO9z98xDceosuGiQ=
+golang.org/x/tools v0.0.0-20180917221912-90fa682c2a6e h1:FDhOuMEY4JVRztM/gsbk+IKUQ8kj74bxZrgw87eMMVc=
 golang.org/x/tools v0.0.0-20180917221912-90fa682c2a6e/go.mod h1:n7NCudcB/nEzxVGmLbDWY5pfWTLqBcC2KZ6jyYvM4mQ=
 google.golang.org/appengine v1.1.0/go.mod h1:EbEs0AVv82hx2wNQdGPgUI5lhzA/G0D9YwlJXL52JkM=
 google.golang.org/appengine v1.6.5/go.mod h1:8WjMMxjGQR8xUklV/ARdw2HLXBOI7O7uCIDZVag1xfc=

helpers/manage_registration.go

@@ -7,6 +7,8 @@ import (
 	"go/format"
 	"go/parser"
 	"go/token"
+	"golang.org/x/tools/go/ast/astutil"
+	"log"
 	"os"
 	"strings"
 )
@@ -25,6 +27,13 @@ const (
 	TypeDatasource resourceType = "datasource"
 )
 
+type astKey struct {
+	KeyValuePos int
+	KeyKind     token.Token
+	KeyValue    string
+	ValFun      string
+}
+
 // UpdateRegistration modifies the `registration.go` file to add (or remove?) items from the relevant blocks when
 // a user adds (or removes?) a reource or datasource via `tfpdk resource` or `tfpdk datasource`
 //
@@ -46,18 +55,45 @@ func UpdateRegistration(servicePackagePath string, resourceName string, resource
 		fn, ok := node.(*ast.FuncDecl)
 		if ok {
 			if fn.Name.Name == nodeName {
+				offset := int(fn.Body.List[0].(*ast.ReturnStmt).Results[0].(*ast.CompositeLit).Lbrace) + 5
 				ast.Inspect(fn.Body.List[0], func(r ast.Node) bool {
 					ret, ok := r.(*ast.ReturnStmt)
 					if ok {
 						ast.Inspect(ret.Results[0], func(n ast.Node) bool {
 							out, ok := n.(*ast.CompositeLit)
 							if ok {
-								pos := 0
-								if count := len(out.Elts); count > 0 {
-									pos = int(out.Elts[count-1].(*ast.KeyValueExpr).Value.(*ast.CallExpr).Rparen) + 5
+								resourceList := make([]astKey, 0)
+								maxOffset := 0
+								for _, v := range out.Elts {
+									offset = offset + maxOffset
+									key := v.(*ast.KeyValueExpr).Key.(*ast.BasicLit)
+									val := v.(*ast.KeyValueExpr).Value.(*ast.CallExpr)
+									resourceList = append(resourceList, astKey{
+										KeyValuePos: offset,
+										KeyKind:     key.Kind,
+										KeyValue:    key.Value,
+										ValFun:      val.Fun.(*ast.Ident).Name,
+									})
+									newOffset := len(string(key.Kind)) + len(key.Value) + len(val.Fun.(*ast.Ident).Name) + 4
+									if newOffset > maxOffset {
+										maxOffset = newOffset
+									}
 								}
-								newReturnEntry := newUnTypedASTReturnEntry(resourceName, "testEntry", pos)
-								out.Elts = append(out.Elts, newReturnEntry)
+
+								resourceList = append(resourceList, astKey{
+									KeyValuePos: offset + maxOffset,
+									KeyKind:     token.STRING,
+									KeyValue:    TerraformResourceName("azurerm", resourceName),
+									ValFun:      "testEntry",
+								})
+
+								newElts := make([]ast.Expr, 0)
+
+								for _, v := range resourceList {
+									newElts = append(newElts, newUnTypedASTReturnEntry(v.KeyValue, v.ValFun, v.KeyValuePos))
+								}
+
+								out.Elts = newElts
 								ret.Results[0] = out
 							}
 							return true
@@ -104,7 +140,7 @@ func newUnTypedASTReturnEntry(key string, value string, pos int) *ast.KeyValueEx
 		Key: &ast.BasicLit{
 			ValuePos: token.Pos(pos),
 			Kind:     token.STRING,
-			Value:    fmt.Sprintf("%q", key),
+			Value:    fmt.Sprintf("%q", strings.Trim(key, "\"")),
 		},
 		Value: &ast.CallExpr{
 			Fun: &ast.Ident{
@@ -113,3 +149,49 @@ func newUnTypedASTReturnEntry(key string, value string, pos int) *ast.KeyValueEx
 		},
 	}
 }
+
+func UpdateRegistrationByNode(servicePackagePath string, resourceName string, resource resourceType, _ operation, isTyped bool) error {
+	fSet := token.NewFileSet()
+	regFilePath := fmt.Sprintf("%s/registration.go", strings.TrimSuffix(servicePackagePath, "/"))
+	regFile, err := parser.ParseFile(fSet, regFilePath, nil, 0)
+	if err != nil {
+		return err
+	}
+
+	nodeName := normaliseNodeName(resource, isTyped)
+	newKeyValue := TerraformResourceName("azurerm", resourceName)
+	n := astutil.Apply(regFile, func(c *astutil.Cursor) bool {
+		if d, ok := c.Parent().(*ast.FuncDecl); ok && d.Name.Name == nodeName && c.Name() == "Body" {
+			node := c.Node()
+			log.Printf("%+v", node)
+			if _, ok := node.(*ast.BlockStmt); ok {
+				elts := node.(*ast.BlockStmt).List[0].(*ast.ReturnStmt).Results[0].(*ast.CompositeLit).Elts
+				log.Printf("%+v", elts)
+				elts = append(elts, newUnTypedASTReturnEntry(newKeyValue, "testEntry", 0))
+				node.(*ast.BlockStmt).List[0].(*ast.ReturnStmt).Results[0].(*ast.CompositeLit).Elts = elts
+				c.Replace(node)
+			}
+		}
+		return true
+	}, nil)
+
+	outBuf := new(bytes.Buffer)
+	if err := format.Node(outBuf, fSet, n); err != nil {
+		return err
+	}
+
+	if err := os.WriteFile(regFilePath, outBuf.Bytes(), 0755); err != nil {
+		return err
+	}
+
+	return nil
+}
+
+//
+//n := astutil.Apply(out.Elts[len(out.Elts)-1], func(c *astutil.Cursor) bool {
+//	if _, ok := c.Parent().(*ast.CompositeLit); ok {
+//		c.InsertAfter(newUnTypedASTReturnEntry(resourceName, TerraformResourceName("azurerm", resourceName), 0))
+//	}
+//	return true
+//}, nil)
+//