Support different output streams than os.Stderr for std.trace (C API) (google#520)

Support different output streams than os.Stderr for std.trace (C API)

Author: hanyucui

c-bindings-tests/compat_test.py

@@ -1,8 +1,9 @@
 #!/usr/bin/env python3
-import unittest
 import ctypes
-import re
+import io
 import os
+import re
+import unittest
 
 
 lib = ctypes.CDLL('../c-bindings/libgojsonnet.so')
@@ -128,6 +129,19 @@
 ]
 lib.jsonnet_import_callback.restype = None
 
+IO_WRITER_CALLBACK = ctypes.CFUNCTYPE(
+    ctypes.c_int,
+    ctypes.c_void_p,
+    ctypes.c_size_t,
+    ctypes.POINTER(ctypes.c_int)
+)
+
+lib.jsonnet_set_trace_out_callback.argtypes = [
+    ctypes.c_void_p,
+    IO_WRITER_CALLBACK,
+]
+lib.jsonnet_set_trace_out_callback.restype = None
+
 # json declaration
 
 lib.jsonnet_json_make_string.argtypes = [
@@ -339,6 +353,15 @@ def import_callback(ctx, dir, rel, found_here, success):
 
     return ctypes.addressof(dst.contents)
 
+io_writer_buf = None
+
+@IO_WRITER_CALLBACK
+def io_writer_callback(buf, nbytes, success):
+    global io_writer_buf
+    io_writer_buf = ctypes.string_at(buf, nbytes)
+    success[0] = ctypes.c_int(1)
+    return nbytes
+
 #  Returns content if worked, None if file not found, or throws an exception
 def jsonnet_try_path(dir, rel):
     if not rel:
@@ -490,6 +513,15 @@ def test_jsonnet_import_callback(self):
         self.assertEqual(b'42\n', to_bytes(res))
         free_buffer(self.vm, res)
 
+    def test_jsonnet_set_trace_out_callback(self):
+        lib.jsonnet_set_trace_out_callback(self.vm, io_writer_callback)
+        fname = b"vm1"
+        msg = b"test_jsonnet_set_trace_out_callback trace message"
+        expected = b"TRACE: " + fname + b":1 " + msg + b"\n"
+        snippet = b"std.trace('" + msg + b"', 'rest')"
+        lib.jsonnet_evaluate_snippet(self.vm, fname, snippet, self.err_ref)
+        self.assertEqual(io_writer_buf,  expected)
+
     def tearDown(self):
         lib.jsonnet_destroy(self.vm)
 

c-bindings/c-bindings.go

@@ -639,5 +639,29 @@ func formatSnippet(vmRef *C.struct_JsonnetVm, filename string, code string, e *C
 	return result
 }
 
+type traceOut struct {
+	cb *C.JsonnetIoWriterCallback
+}
+
+func (o *traceOut) Write(p []byte) (int, error) {
+	if len(p) == 0 {
+		return 0, nil
+	}
+
+	success := C.int(0)
+	var n C.int = C.jsonnet_internal_execute_writer(o.cb, unsafe.Pointer(&p[0]),
+				                                    C.size_t(len(p)), &success)
+	if success != 1 {
+		return int(n), errors.New("std.trace() failed to write to output stream")
+	}
+	return int(n), nil
+}
+
+//export jsonnet_set_trace_out_callback
+func jsonnet_set_trace_out_callback(vmRef *C.struct_JsonnetVm, cb *C.JsonnetIoWriterCallback) {
+	vm := getVM(vmRef)
+	vm.SetTraceOut(&traceOut{cb})
+}
+
 func main() {
 }

c-bindings/internal.h

@@ -38,4 +38,11 @@ char* jsonnet_internal_execute_import(JsonnetImportCallback *cb,
                                       char **found_here,
                                       int *success);
 
+typedef int JsonnetIoWriterCallback(const void *buf, size_t nbytes, int *success);
+
+int jsonnet_internal_execute_writer(JsonnetIoWriterCallback *cb,
+                                    const void *buf,
+                                    size_t nbytes,
+                                    int *success);
+
 void jsonnet_internal_free_string(char *str);

c-bindings/libjsonnet.cpp

@@ -44,6 +44,14 @@ char* jsonnet_internal_execute_import(JsonnetImportCallback *cb,
     return (cb)(ctx, base, rel, found_here, success);
 }
 
+int jsonnet_internal_execute_writer(JsonnetIoWriterCallback *cb,
+                                    const void *buf,
+                                    size_t nbytes,
+                                    int *success)
+{
+    return (cb)(buf, nbytes, success);
+}
+
 void jsonnet_internal_free_string(char *str) {
     if (str != nullptr) {
         ::free(str);