Allow setting a different output stream than os.Stderr for std.trace (google#514)

Allow std.trace output to other streams

Author: hanyucui

builtins.go

@@ -24,7 +24,6 @@ import (
 	"encoding/json"
 	"fmt"
 	"math"
-	"os"
 	"reflect"
 	"sort"
 	"strconv"
@@ -248,10 +247,10 @@ func builtinTrace(i *interpreter, x value, y value) (value, error) {
 		return nil, err
 	}
 	trace := i.stack.currentTrace
-	filename := trace.loc.FileName
+	filename := trace.loc.File.DiagnosticFileName
 	line := trace.loc.Begin.Line
 	fmt.Fprintf(
-		os.Stderr, "TRACE: %s:%d %s\n", filename, line, xStr.getGoString())
+		i.traceOut, "TRACE: %s:%d %s\n", filename, line, xStr.getGoString())
 	return y, nil
 }
 

interpreter.go

@@ -19,6 +19,7 @@ package jsonnet
 import (
 	"bytes"
 	"fmt"
+	"io"
 	"math"
 	"reflect"
 	"sort"
@@ -257,6 +258,9 @@ type interpreter struct {
 
 	// Keeps imports
 	importCache *importCache
+
+	// Output stream for trace() for
+	traceOut io.Writer
 }
 
 // Map union, b takes precedence when keys collide.
@@ -1215,10 +1219,11 @@ func buildObject(hide ast.ObjectFieldHide, fields map[string]value) *valueObject
 	return makeValueSimpleObject(bindingFrame{}, fieldMap, nil, nil)
 }
 
-func buildInterpreter(ext vmExtMap, nativeFuncs map[string]*NativeFunction, maxStack int, ic *importCache) (*interpreter, error) {
+func buildInterpreter(ext vmExtMap, nativeFuncs map[string]*NativeFunction, maxStack int, ic *importCache, traceOut io.Writer) (*interpreter, error) {
 	i := interpreter{
 		stack:       makeCallStack(maxStack),
 		importCache: ic,
+		traceOut:    traceOut,
 		nativeFuncs: nativeFuncs,
 	}
 
@@ -1288,9 +1293,9 @@ func evaluateAux(i *interpreter, node ast.Node, tla vmExtMap) (value, error) {
 
 // TODO(sbarzowski) this function takes far too many arguments - build interpreter in vm instead
 func evaluate(node ast.Node, ext vmExtMap, tla vmExtMap, nativeFuncs map[string]*NativeFunction,
-	maxStack int, ic *importCache, stringOutputMode bool) (string, error) {
+	maxStack int, ic *importCache, traceOut io.Writer, stringOutputMode bool) (string, error) {
 
-	i, err := buildInterpreter(ext, nativeFuncs, maxStack, ic)
+	i, err := buildInterpreter(ext, nativeFuncs, maxStack, ic, traceOut)
 	if err != nil {
 		return "", err
 	}
@@ -1317,9 +1322,9 @@ func evaluate(node ast.Node, ext vmExtMap, tla vmExtMap, nativeFuncs map[string]
 
 // TODO(sbarzowski) this function takes far too many arguments - build interpreter in vm instead
 func evaluateMulti(node ast.Node, ext vmExtMap, tla vmExtMap, nativeFuncs map[string]*NativeFunction,
-	maxStack int, ic *importCache, stringOutputMode bool) (map[string]string, error) {
+	maxStack int, ic *importCache, traceOut io.Writer, stringOutputMode bool) (map[string]string, error) {
 
-	i, err := buildInterpreter(ext, nativeFuncs, maxStack, ic)
+	i, err := buildInterpreter(ext, nativeFuncs, maxStack, ic, traceOut)
 	if err != nil {
 		return nil, err
 	}
@@ -1337,9 +1342,9 @@ func evaluateMulti(node ast.Node, ext vmExtMap, tla vmExtMap, nativeFuncs map[st
 
 // TODO(sbarzowski) this function takes far too many arguments - build interpreter in vm instead
 func evaluateStream(node ast.Node, ext vmExtMap, tla vmExtMap, nativeFuncs map[string]*NativeFunction,
-	maxStack int, ic *importCache) ([]string, error) {
+	maxStack int, ic *importCache, traceOut io.Writer) ([]string, error) {
 
-	i, err := buildInterpreter(ext, nativeFuncs, maxStack, ic)
+	i, err := buildInterpreter(ext, nativeFuncs, maxStack, ic, traceOut)
 	if err != nil {
 		return nil, err
 	}

jsonnet_test.go

@@ -3,6 +3,7 @@ package jsonnet
 import (
 	"bytes"
 	"encoding/json"
+	"fmt"
 	"reflect"
 	"strings"
 	"testing"
@@ -327,3 +328,24 @@ func TestTLATypes(t *testing.T) {
 	}
 	assertVarOutput(t, jsonStr)
 }
+
+func TestSetTraceOut(t *testing.T) {
+	traceOut := &strings.Builder{}
+	vm := MakeVM()
+	vm.SetTraceOut(traceOut)
+
+	const filename = "blah.jsonnet"
+	const msg = "TestSetTraceOut Trace Message"
+	expected := fmt.Sprintf("TRACE: %s:1 %s", filename, msg)
+	input := fmt.Sprintf("std.trace('%s', 'rest')", msg)
+
+	_, err := vm.EvaluateAnonymousSnippet(filename, input)
+	if err != nil {
+		t.Errorf("Unexpected error: %v", err)
+	}
+
+	actual := removeExcessiveWhitespace(traceOut.String())
+	if actual != expected {
+		t.Errorf("Expected %q, but got %q", expected, actual)
+	}
+}

vm.go

@@ -19,6 +19,7 @@ package jsonnet
 import (
 	"errors"
 	"fmt"
+	"io"
 	"os"
 	"path/filepath"
 	"runtime/debug"
@@ -44,6 +45,7 @@ type VM struct {
 	ErrorFormatter ErrorFormatter
 	StringOutput   bool
 	importCache    *importCache
+	traceOut       io.Writer
 }
 
 // extKind indicates the kind of external variable that is being initialized for the VM
@@ -78,6 +80,7 @@ func MakeVM() *VM {
 		ErrorFormatter: &termErrorFormatter{pretty: false, maxStackTraceSize: 20},
 		importer:       &FileImporter{},
 		importCache:    makeImportCache(defaultImporter),
+		traceOut:       os.Stderr,
 	}
 }
 
@@ -93,6 +96,11 @@ func (vm *VM) flushValueCache() {
 	vm.importCache.flushValueCache()
 }
 
+// SetTraceOut sets the output stream for the builtin function std.trace().
+func (vm *VM) SetTraceOut(traceOut io.Writer) {
+	vm.traceOut = traceOut
+}
+
 // ExtVar binds a Jsonnet external var to the given value.
 func (vm *VM) ExtVar(key string, val string) {
 	vm.ext[key] = vmExt{value: val, kind: extKindVar}
@@ -174,7 +182,7 @@ func (vm *VM) Evaluate(node ast.Node) (val string, err error) {
 			err = fmt.Errorf("(CRASH) %v\n%s", r, debug.Stack())
 		}
 	}()
-	return evaluate(node, vm.ext, vm.tla, vm.nativeFuncs, vm.MaxStack, vm.importCache, vm.StringOutput)
+	return evaluate(node, vm.ext, vm.tla, vm.nativeFuncs, vm.MaxStack, vm.importCache, vm.traceOut, vm.StringOutput)
 }
 
 // EvaluateStream evaluates a Jsonnet program given by an Abstract Syntax Tree
@@ -185,7 +193,7 @@ func (vm *VM) EvaluateStream(node ast.Node) (output []string, err error) {
 			err = fmt.Errorf("(CRASH) %v\n%s", r, debug.Stack())
 		}
 	}()
-	return evaluateStream(node, vm.ext, vm.tla, vm.nativeFuncs, vm.MaxStack, vm.importCache)
+	return evaluateStream(node, vm.ext, vm.tla, vm.nativeFuncs, vm.MaxStack, vm.importCache, vm.traceOut)
 }
 
 // EvaluateMulti evaluates a Jsonnet program given by an Abstract Syntax Tree
@@ -197,7 +205,7 @@ func (vm *VM) EvaluateMulti(node ast.Node) (output map[string]string, err error)
 			err = fmt.Errorf("(CRASH) %v\n%s", r, debug.Stack())
 		}
 	}()
-	return evaluateMulti(node, vm.ext, vm.tla, vm.nativeFuncs, vm.MaxStack, vm.importCache, vm.StringOutput)
+	return evaluateMulti(node, vm.ext, vm.tla, vm.nativeFuncs, vm.MaxStack, vm.importCache, vm.traceOut, vm.StringOutput)
 }
 
 func (vm *VM) evaluateSnippet(diagnosticFileName ast.DiagnosticFileName, filename string, snippet string, kind evalKind) (output interface{}, err error) {
@@ -212,11 +220,11 @@ func (vm *VM) evaluateSnippet(diagnosticFileName ast.DiagnosticFileName, filenam
 	}
 	switch kind {
 	case evalKindRegular:
-		output, err = evaluate(node, vm.ext, vm.tla, vm.nativeFuncs, vm.MaxStack, vm.importCache, vm.StringOutput)
+		output, err = evaluate(node, vm.ext, vm.tla, vm.nativeFuncs, vm.MaxStack, vm.importCache, vm.traceOut, vm.StringOutput)
 	case evalKindMulti:
-		output, err = evaluateMulti(node, vm.ext, vm.tla, vm.nativeFuncs, vm.MaxStack, vm.importCache, vm.StringOutput)
+		output, err = evaluateMulti(node, vm.ext, vm.tla, vm.nativeFuncs, vm.MaxStack, vm.importCache, vm.traceOut, vm.StringOutput)
 	case evalKindStream:
-		output, err = evaluateStream(node, vm.ext, vm.tla, vm.nativeFuncs, vm.MaxStack, vm.importCache)
+		output, err = evaluateStream(node, vm.ext, vm.tla, vm.nativeFuncs, vm.MaxStack, vm.importCache, vm.traceOut)
 	}
 	if err != nil {
 		return "", err