feat: add test init

internal/analysis/check.go

@@ -135,32 +135,32 @@ type CheckResult struct {
 	Program                parser.Program
 
 	stmtType  Type
-	ExprTypes map[parser.ValueExpr]Type
-	VarTypes  map[parser.VarDeclaration]Type
+	exprTypes map[parser.ValueExpr]Type
+	varTypes  map[parser.VarDeclaration]Type
 }
 
-func (r *CheckResult) getExprType(expr parser.ValueExpr) Type {
-	exprType, ok := r.ExprTypes[expr]
+func (r *CheckResult) GetExprType(expr parser.ValueExpr) Type {
+	exprType, ok := r.exprTypes[expr]
 	if !ok {
 		t := TVar{}
-		r.ExprTypes[expr] = &t
+		r.exprTypes[expr] = &t
 		return &t
 	}
-	return exprType
+	return exprType.Resolve()
 }
 
-func (r *CheckResult) getVarDeclType(decl parser.VarDeclaration) Type {
-	exprType, ok := r.VarTypes[decl]
+func (r *CheckResult) GetVarDeclType(decl parser.VarDeclaration) Type {
+	exprType, ok := r.varTypes[decl]
 	if !ok {
 		t := TVar{}
-		r.VarTypes[decl] = &t
+		r.varTypes[decl] = &t
 		return &t
 	}
-	return exprType
+	return exprType.Resolve()
 }
 
 func (r *CheckResult) unifyNodeWith(expr parser.ValueExpr, t Type) {
-	exprT := r.getExprType(expr)
+	exprT := r.GetExprType(expr)
 	r.unify(expr.GetRange(), exprT, t)
 }
 
@@ -224,8 +224,8 @@ func newCheckResult(program parser.Program) CheckResult {
 		unusedVars:       make(map[string]parser.Range),
 		varResolution:    make(map[*parser.Variable]parser.VarDeclaration),
 		fnCallResolution: make(map[*parser.FnCallIdentifier]FnCallResolution),
-		ExprTypes:        make(map[parser.ValueExpr]Type),
-		VarTypes:         make(map[parser.VarDeclaration]Type),
+		exprTypes:        make(map[parser.ValueExpr]Type),
+		varTypes:         make(map[parser.VarDeclaration]Type),
 	}
 }
 
@@ -242,7 +242,7 @@ func (res *CheckResult) check() {
 
 			if varDecl.Origin != nil {
 				res.checkExpression(*varDecl.Origin, varDecl.Type.Name)
-				res.unifyNodeWith(*varDecl.Origin, res.getVarDeclType(varDecl))
+				res.unifyNodeWith(*varDecl.Origin, res.GetVarDeclType(varDecl))
 			}
 		}
 	}
@@ -357,10 +357,10 @@ func (res *CheckResult) checkFnCallArity(fnCall *parser.FnCall) {
 		case FnVarOriginBalance, FnVarOriginOverdraft:
 			// we run unify(<expr>, <asset>) in:
 			// <expr> := balance(@acc, <asset>)
-			res.unifyNodeWith(fnCall, res.getExprType(validArgs[1]))
+			res.unifyNodeWith(fnCall, res.GetExprType(validArgs[1]))
 
 		case FnVarOriginGetAsset:
-			res.unifyNodeWith(fnCall, res.getExprType(validArgs[0]))
+			res.unifyNodeWith(fnCall, res.GetExprType(validArgs[0]))
 		}
 	} else {
 		for _, arg := range validArgs {
@@ -421,7 +421,7 @@ func (res *CheckResult) checkTypeOf(lit parser.ValueExpr, typeHint string) strin
 	case *parser.Variable:
 		if varDeclaration, ok := res.DeclaredVars[lit.Name]; ok {
 			res.varResolution[lit] = varDeclaration
-			res.unifyNodeWith(lit, res.getVarDeclType(varDeclaration))
+			res.unifyNodeWith(lit, res.GetVarDeclType(varDeclaration))
 		} else {
 			res.pushDiagnostic(lit.Range, UnboundVariable{Name: lit.Name, Type: typeHint})
 		}
@@ -440,11 +440,11 @@ func (res *CheckResult) checkTypeOf(lit parser.ValueExpr, typeHint string) strin
 			we unify $mon and $asset in:
 			`let $mon := [$asset 42]`
 		*/
-		res.unifyNodeWith(lit, res.getExprType(lit.Asset))
+		res.unifyNodeWith(lit, res.GetExprType(lit.Asset))
 		return TypeMonetary
 
 	case *parser.BinaryInfix:
-		res.unifyNodeWith(lit.Left, res.getExprType(lit.Right))
+		res.unifyNodeWith(lit.Left, res.GetExprType(lit.Right))
 
 		switch lit.Operator {
 		case parser.InfixOperatorPlus:

internal/analysis/check_test.go

@@ -1066,9 +1066,8 @@ send $mon1 (
 
 	res := analysis.CheckSource(input)
 
-	t1 := res.VarTypes[res.DeclaredVars["mon1"]]
-
-	t2 := res.VarTypes[res.DeclaredVars["mon2"]]
+	t1 := res.GetVarDeclType(res.DeclaredVars["mon1"])
+	t2 := res.GetVarDeclType(res.DeclaredVars["mon2"])
 
 	require.Same(t, t1.Resolve(), t2.Resolve())
 }
@@ -1085,7 +1084,7 @@ vars {
 	res := analysis.CheckSource(input)
 
 	v := res.DeclaredVars["ass"]
-	t1 := res.VarTypes[v]
+	t1 := res.GetVarDeclType(v)
 
 	expected := analysis.TAsset("USD/2")
 	require.Equal(t, &expected, t1.Resolve())

internal/cmd/root.go

@@ -26,6 +26,7 @@ func Execute(options CliOptions) {
 	rootCmd.AddCommand(lspCmd)
 	rootCmd.AddCommand(checkCmd)
 	rootCmd.AddCommand(getTestCmd())
+	rootCmd.AddCommand(getTestInitCmd())
 	rootCmd.AddCommand(getRunCmd())
 
 	if err := rootCmd.Execute(); err != nil {

internal/cmd/test.go

@@ -11,9 +11,7 @@ type testArgs struct {
 	paths []string
 }
 
-var opts = testArgs{}
-
-func runTestCmd() {
+func runTestCmd(opts testArgs) {
 	files, err := specs_format.ReadSpecsFiles(opts.paths)
 	if err != nil {
 		_, _ = os.Stderr.Write([]byte(err.Error()))
@@ -37,13 +35,12 @@ and tests the corresponding <file>.num file (if any).
 Defaults to "." if there are no given paths`,
 		Args: cobra.MatchAll(),
 		Run: func(cmd *cobra.Command, paths []string) {
-
 			if len(paths) == 0 {
 				paths = []string{"."}
 			}
-
-			opts.paths = paths
-			runTestCmd()
+			runTestCmd(testArgs{
+				paths: paths,
+			})
 		},
 	}
 

internal/cmd/test_init.go

@@ -0,0 +1,239 @@
+package cmd
+
+import (
+	"context"
+	"encoding/json"
+	"fmt"
+	"math/big"
+	"os"
+
+	"github.com/formancehq/numscript/internal/analysis"
+	"github.com/formancehq/numscript/internal/interpreter"
+	"github.com/formancehq/numscript/internal/parser"
+	"github.com/formancehq/numscript/internal/specs_format"
+	"github.com/formancehq/numscript/internal/utils"
+	"github.com/spf13/cobra"
+)
+
+type testInitArgs struct {
+	path string
+}
+
+func getTestInitCmd() *cobra.Command {
+
+	cmd := &cobra.Command{
+		Use:   "test-init path",
+		Short: "Create a specs file for the given numscript",
+		Args:  cobra.ExactArgs(1),
+		RunE: func(cmd *cobra.Command, args []string) error {
+			err := runTestInitCmd(testInitArgs{
+				path: args[0],
+			})
+
+			if err != nil {
+				cmd.SilenceErrors = true
+				cmd.SilenceUsage = true
+				return err
+			}
+
+			return nil
+
+		},
+	}
+
+	return cmd
+}
+
+func mkDefaultVar(decl parser.VarDeclaration, checkResult analysis.CheckResult) string {
+	defaultAmt := 100
+	defaultCurr := "USD/2"
+
+	asset := checkResult.GetVarDeclType(decl)
+	switch asset := asset.(type) {
+	case *analysis.TAsset:
+		defaultCurr = string(*asset)
+	}
+
+	switch decl.Type.Name {
+
+	case analysis.TypeMonetary:
+		return fmt.Sprintf("%s %d", defaultCurr, defaultAmt)
+
+	case analysis.TypeNumber:
+		return fmt.Sprintf("%d", defaultAmt)
+
+	case analysis.TypeAccount:
+		return decl.Name.Name
+
+	case analysis.TypeString:
+		return decl.Name.Name
+
+	case analysis.TypePortion:
+		return "1/2"
+
+	case analysis.TypeAsset:
+		return string(defaultCurr)
+	}
+
+	return ""
+}
+
+func MakeSpecsFile(source string) (specs_format.Specs, error) {
+	parseResult := parser.Parse(source)
+	if len(parseResult.Errors) != 0 {
+		fmt.Fprint(os.Stderr, parser.ParseErrorsToString(parseResult.Errors, source))
+		return specs_format.Specs{}, fmt.Errorf("parsing failed")
+	}
+
+	checkResult := analysis.CheckProgram(parseResult.Value)
+
+	vars := map[string]string{}
+	if parseResult.Value.Vars != nil {
+		for _, decl := range parseResult.Value.Vars.Declarations {
+			if decl.Origin != nil {
+				continue
+			}
+
+			value := mkDefaultVar(decl, checkResult)
+			vars[decl.Name.Name] = value
+		}
+	}
+
+	return makeSpecsFile(
+		parseResult.Value,
+		vars,
+		map[string]struct{}{},
+		big.NewInt(100),
+	)
+}
+
+func makeSpecsFile(
+	program parser.Program,
+	vars map[string]string,
+	featureFlags map[string]struct{},
+	defaultBalance *big.Int,
+) (specs_format.Specs, error) {
+
+	store := TestInitStore{
+		DefaultBalance: defaultBalance,
+		Balances:       make(interpreter.Balances),
+		Meta:           make(interpreter.AccountsMetadata),
+	}
+
+	res, iErr := interpreter.RunProgram(
+		context.Background(),
+		program,
+		vars,
+		store,
+		featureFlags,
+	)
+
+	if iErr != nil {
+		missingFundsErr, missingFunds := iErr.(interpreter.MissingFundsErr)
+		if missingFunds {
+			// TODO we could have better heuristics with a balance for each account/asset pair
+			return makeSpecsFile(
+				program,
+				vars,
+				featureFlags,
+				&missingFundsErr.Needed,
+			)
+		}
+
+		expFeatErr, missingFeatureFlag := iErr.(interpreter.ExperimentalFeature)
+		if missingFeatureFlag {
+			featureFlags[expFeatErr.FlagName] = struct{}{}
+			return makeSpecsFile(
+				program,
+				vars,
+				featureFlags,
+				&missingFundsErr.Needed,
+			)
+		}
+
+		return specs_format.Specs{}, iErr
+	}
+
+	var featureFlags_ []string
+	for k := range featureFlags {
+		featureFlags_ = append(featureFlags_, k)
+	}
+
+	specs := specs_format.Specs{
+		Schema:       "https://raw.githubusercontent.com/formancehq/numscript/main/specs.schema.json",
+		Balances:     store.Balances,
+		Vars:         vars,
+		FeatureFlags: featureFlags_,
+		TestCases: []specs_format.TestCase{
+			{
+				It:             "example spec",
+				ExpectPostings: res.Postings,
+			},
+		},
+	}
+
+	return specs, nil
+}
+
+func runTestInitCmd(opts testInitArgs) error {
+	// TODO check there isn't a specsfile already
+
+	numscriptContent, err := os.ReadFile(opts.path)
+	if err != nil {
+		return err
+	}
+
+	specs, err := MakeSpecsFile(string(numscriptContent))
+
+	if err != nil {
+		return err
+	}
+
+	marshaled, _ := json.MarshalIndent(specs, "", "  ")
+
+	os.WriteFile(opts.path+".specs.json", marshaled, 0644)
+
+	fmt.Printf("✅ Created specs file: %s.specs.json\n", opts.path)
+
+	return nil
+}
+
+type TestInitStore struct {
+	DefaultBalance *big.Int
+	Balances       interpreter.Balances
+	Meta           interpreter.AccountsMetadata
+}
+
+func (s TestInitStore) GetBalances(_ context.Context, q interpreter.BalanceQuery) (interpreter.Balances, error) {
+	outputBalance := interpreter.Balances{}
+	for queriedAccount, queriedCurrencies := range q {
+
+		for _, curr := range queriedCurrencies {
+			amt := utils.NestedMapGetOrPutDefault(s.Balances, queriedAccount, curr, func() *big.Int {
+				return new(big.Int).Set(s.DefaultBalance)
+			})
+
+			outputAccountBalance := utils.MapGetOrPutDefault(outputBalance, queriedAccount, func() interpreter.AccountBalance {
+				return interpreter.AccountBalance{}
+			})
+
+			outputAccountBalance[curr] = new(big.Int).Set(amt)
+		}
+	}
+
+	return outputBalance, nil
+}
+
+func (s TestInitStore) GetAccountsMetadata(c context.Context, q interpreter.MetadataQuery) (interpreter.AccountsMetadata, error) {
+	outputMeta := interpreter.AccountsMetadata{}
+	for queriedAccount, queriedCurrencies := range q {
+		for _, curr := range queriedCurrencies {
+			outputAccountMeta := utils.MapGetOrPutDefault(outputMeta, queriedAccount, func() interpreter.AccountMetadata {
+				return interpreter.AccountMetadata{}
+			})
+			outputAccountMeta[curr] = ""
+		}
+	}
+
+	return outputMeta, nil
+}