From 16ae8ff5c78e811e2019f270eb493c6a8d434767 Mon Sep 17 00:00:00 2001
From: itchyny <itchyny@hatena.ne.jp>
Date: Thu, 24 Dec 2020 21:19:57 +0900
Subject: [PATCH] implement and export Marshal for jq-flavored encoding

---
 encoder.go      | 26 +++++++++++++++++-----
 encoder_test.go | 58 +++++++++++++++++++++++++++++++++++++++++++++++++
 2 files changed, 79 insertions(+), 5 deletions(-)
 create mode 100644 encoder_test.go

diff --git a/encoder.go b/encoder.go
index 089cd126..2a5e90cd 100644
--- a/encoder.go
+++ b/encoder.go
@@ -1,23 +1,39 @@
 package gojq
 
 import (
+	"bytes"
 	"fmt"
 	"math"
 	"math/big"
 	"sort"
 	"strconv"
-	"strings"
 	"unicode/utf8"
 )
 
+// Marshal returns the jq-flavored JSON encoding of v.
+//
+// This method only accepts limited types (nil, bool, int, float64, *big.Int,
+// string, []interface{} and map[string]interface{}) because these are the
+// possible types a gojq iterator can emit. This method marshals NaN to null,
+// truncates infinities to (+|-) math.MaxFloat64 and does not escape '<' and
+// '>' for embedding in HTML. These behaviors are based on the marshaler of jq
+// command and different from the standard library method json.Marshal.
+func Marshal(v interface{}) ([]byte, error) {
+	return jsonMarshalBytes(v), nil
+}
+
 func jsonMarshal(v interface{}) string {
-	var s strings.Builder
-	(&encoder{w: &s}).encode(v)
-	return s.String()
+	return string(jsonMarshalBytes(v))
+}
+
+func jsonMarshalBytes(v interface{}) []byte {
+	var b bytes.Buffer
+	(&encoder{w: &b}).encode(v)
+	return b.Bytes()
 }
 
 type encoder struct {
-	w   *strings.Builder
+	w   *bytes.Buffer
 	buf [64]byte
 }
 
diff --git a/encoder_test.go b/encoder_test.go
new file mode 100644
index 00000000..f1366ab1
--- /dev/null
+++ b/encoder_test.go
@@ -0,0 +1,58 @@
+package gojq
+
+import (
+	"math"
+	"math/big"
+	"testing"
+)
+
+func TestMarshal(t *testing.T) {
+	testCases := []struct {
+		name     string
+		value    interface{}
+		expected string
+	}{
+		{
+			name:     "nil",
+			value:    nil,
+			expected: "null",
+		},
+		{
+			name:     "booleans",
+			value:    []interface{}{false, true},
+			expected: "[false,true]",
+		},
+		{
+			name: "numbers",
+			value: []interface{}{42, 3.14, 1e-6, 1e-7, -1e-9, 1e-10, math.NaN(), math.Inf(1), math.Inf(-1),
+				new(big.Int).SetBytes([]byte("\xff\xff\xff\xff\xff\xff\xff\xff\xff\xff\xff\xff\xff\xff\xff\xff"))},
+			expected: "[42,3.14,0.000001,1e-7,-1e-9,1e-10,null,1.7976931348623157e+308,-1.7976931348623157e+308,340282366920938463463374607431768211455]",
+		},
+		{
+			name:     "strings",
+			value:    []interface{}{"", "abcde", "foo\x00\x1f\r\t\n\f\b<=>!\"#$%'& \\\x7fbar"},
+			expected: `["","abcde","foo\u0000\u001f\r\t\n\u000c\u0008<=>!\"#$%'& \\\u007fbar"]`,
+		},
+		{
+			name:     "arrays",
+			value:    []interface{}{1, []interface{}{2, []interface{}{3, []interface{}{map[string]interface{}{}}}}},
+			expected: `[1,[2,[3,[{}]]]]`,
+		},
+		{
+			name:     "objects",
+			value:    map[string]interface{}{"x": []interface{}{100}, "y": map[string]interface{}{"z": 42}},
+			expected: `{"x":[100],"y":{"z":42}}`,
+		},
+	}
+	for _, tc := range testCases {
+		t.Run(tc.name, func(t *testing.T) {
+			got, err := Marshal(tc.value)
+			if err != nil {
+				t.Fatal(err)
+			}
+			if string(got) != tc.expected {
+				t.Errorf("expected: %s, got: %s", tc.expected, string(got))
+			}
+		})
+	}
+}