Move code generation to a package

generate.go → gen/generate.go

@@ -1,8 +1,10 @@
-package main
+// Package gen generates Go code from an XGBoost model.
+package gen
 
 import (
 	"fmt"
 	"go/format"
+	"os"
 )
 
 type treeFunction struct {
@@ -65,3 +67,36 @@ func codegenTree(r *renderer, tree *node, level int) (string, error) {
 
 	return r.executeDecisionNode(tree, level, left, right)
 }
+
+// GenerateFile generates a .go file containing a function that implements the XGB model.
+func GenerateFile(
+	inputJSON string,
+	packageName,
+	funcName,
+	outputFile string,
+) error {
+	x, err := readModel(inputJSON)
+	if err != nil {
+		return err
+	}
+
+	trees, err := readTrees(x)
+	if err != nil {
+		return err
+	}
+
+	r, err := newRenderer()
+	if err != nil {
+		return err
+	}
+
+	code, err := generateSource(packageName, funcName, trees, r)
+	if err != nil {
+		return err
+	}
+
+	if err := os.WriteFile(outputFile, []byte(code), 0o644); err != nil {
+		return fmt.Errorf("error writing file: %w", err)
+	}
+	return nil
+}

parse.go → gen/parse.go

@@ -1,4 +1,4 @@
-package main
+package gen
 
 import (
 	"encoding/json"

parse_test.go → gen/parse_test.go

@@ -1,4 +1,4 @@
-package main
+package gen
 
 import (
 	"path/filepath"

templates.go → gen/templates.go

@@ -1,4 +1,4 @@
-package main
+package gen
 
 import (
 	"bytes"

main.go

@@ -1,11 +1,13 @@
-// xgb2code generates code for an XGB model.
+// This program runs a command line program that generates Go code from an xgb model in JSON format.
 package main
 
 import (
 	"flag"
 	"fmt"
 	"log"
 	"os"
+
+	"github.com/maxmind/xgb2code/gen"
 )
 
 type language string
@@ -57,42 +59,8 @@ func main() {
 		os.Exit(1)
 	}
 
-	err := GenerateFile(*inputJSON, *packageName, *funcName, *outputFile)
+	err := gen.GenerateFile(*inputJSON, *packageName, *funcName, *outputFile)
 	if err != nil {
 		log.Fatal(err)
 	}
 }
-
-// GenerateFile generates a .go file containing a function that implements the XGB model.
-func GenerateFile(
-	inputJSON string,
-	packageName,
-	funcName,
-	outputFile string,
-) error {
-	x, err := readModel(inputJSON)
-	if err != nil {
-		return err
-	}
-
-	trees, err := readTrees(x)
-	if err != nil {
-		return err
-	}
-
-	r, err := newRenderer()
-	if err != nil {
-		return err
-	}
-
-	code, err := generateSource(packageName, funcName, trees, r)
-	if err != nil {
-		return err
-	}
-
-	if err := os.WriteFile(outputFile, []byte(code), 0o644); err != nil {
-		return fmt.Errorf("error writing file: %w", err)
-	}
-
-	return nil
-}

main_test.go

@@ -7,6 +7,7 @@ import (
 	"path/filepath"
 	"testing"
 
+	"github.com/maxmind/xgb2code/gen"
 	"github.com/stretchr/testify/require"
 )
 
@@ -21,8 +22,7 @@ func TestGenerateAndRunModels(t *testing.T) {
 	for _, test := range tests {
 		t.Run(test.model, func(t *testing.T) {
 			// Generate the code.
-
-			modelDir := filepath.Join("testdata", test.model)
+			modelDir := filepath.Join("gen", "testdata", test.model)
 			modelFile := filepath.Join(modelDir, "model.json")
 
 			packageName := "main"
@@ -31,13 +31,13 @@ func TestGenerateAndRunModels(t *testing.T) {
 			outputDir := t.TempDir()
 			funcFile := filepath.Join(outputDir, "predict.go")
 
-			err := GenerateFile(modelFile, packageName, functionName, funcFile)
+			err := gen.GenerateFile(modelFile, packageName, functionName, funcFile)
 			require.NoError(t, err)
 
 			// Copy the test program and test data into place.
 
 			files := []string{
-				filepath.Join("testdata", "main.go"),
+				filepath.Join("gen", "testdata", "main.go"),
 				filepath.Join(modelDir, "xtest.csv"),
 				filepath.Join(modelDir, "preds.csv"),
 			}