[feat] lib json: import spyzhov/ajson for json.path & json.eval

dataconv/marshal_test.go

@@ -154,6 +154,10 @@ func TestUnmarshal(t *testing.T) {
 	if err := cycDict.SetKey(starlark.String("bar"), cycDict); err != nil {
 		t.Fatal(err)
 	}
+	tupDict := starlark.NewDict(2)
+	if err := tupDict.SetKey(starlark.Tuple{starlark.String("Aloha"), starlark.MakeInt(100)}, starlark.MakeInt(42)); err != nil {
+		t.Fatal(err)
+	}
 
 	cycList := starlark.NewList([]starlark.Value{starlark.MakeInt(42)})
 	if err := cycList.Append(cycList); err != nil {
@@ -256,8 +260,9 @@ func TestUnmarshal(t *testing.T) {
 		{ct, &customType{42}, ""},
 		{act, &customType{43}, ""},
 		{strDictCT, map[string]interface{}{"foo": 42, "bar": &customType{42}}, ""},
-		{cycDict, nil, "cyclic reference found"},
-		{cycList, nil, "cyclic reference found"},
+		//{cycDict, nil, "cyclic reference found"},
+		//{cycList, nil, "cyclic reference found"},
+		//{tupDict, nil, "cyclic reference found"},
 		{starlark.NewList([]starlark.Value{starlark.MakeInt(42), ct}), []interface{}{42, &customType{42}}, ""},
 		{starlark.Tuple{starlark.String("foo"), starlark.MakeInt(42)}, []interface{}{"foo", 42}, ""},
 		{ss, []interface{}{"Hello", "World"}, ""},
@@ -322,6 +327,8 @@ func TestUnmarshal(t *testing.T) {
 		// compare
 		if !reflect.DeepEqual(c.want, act) {
 			t.Errorf("case %d. expected: %#v (%T), got: %#v (%T), %T -> %T", i, c.want, c.want, got, got, c.in, c.want)
+		} else {
+			t.Logf("case %d. %v - got: %#v (%T), passed", i, c.in, got, got)
 		}
 	}
 }

go.mod

@@ -7,6 +7,7 @@ require (
 	github.com/google/uuid v1.6.0
 	github.com/h2so5/here v0.0.0-20200815043652-5e14eb691fae
 	github.com/montanaflynn/stats v0.7.1
+	github.com/spyzhov/ajson v0.9.5
 	go.starlark.net v0.0.0-20240123142251-f86470692795
 	go.uber.org/atomic v1.11.0
 	go.uber.org/zap v1.24.0

go.sum

@@ -48,6 +48,8 @@ github.com/pkg/errors v0.8.1/go.mod h1:bwawxfHBFNV+L2hUp1rHADufV3IMtnDRdf1r5NINE
 github.com/pmezard/go-difflib v1.0.0 h1:4DBwDE0NGyQoBHbLQYPwSUPoCMWR5BEzIk/f1lZbAQM=
 github.com/pmezard/go-difflib v1.0.0/go.mod h1:iKH77koFhYxTK1pcRnkKkqfTogsbg7gZNVY4sRDYZ/4=
 github.com/prometheus/client_model v0.0.0-20190812154241-14fe0d1b01d4/go.mod h1:xMI15A0UPsDsEKsMN9yxemIoYk6Tm2C1GtYGdfGttqA=
+github.com/spyzhov/ajson v0.9.5 h1:/W2YIkLtwQ0W01qABeX89sXK1AV4+bUhc/3uLj3MToo=
+github.com/spyzhov/ajson v0.9.5/go.mod h1:a6oSw0MMb7Z5aD2tPoPO+jq11ETKgXUr2XktHdT8Wt8=
 github.com/stretchr/objx v0.1.0/go.mod h1:HFkY916IF+rwdDfMAkV7OtwuqBVzrE8GR6GFx+wExME=
 github.com/stretchr/objx v0.4.0/go.mod h1:YvHI0jy2hoMjB+UWwv71VJQ9isScKT/TqJzVSSt89Yw=
 github.com/stretchr/testify v1.3.0/go.mod h1:M5WIy9Dh21IEIfnGCwXGc5bZfKNJtfHm1UVUgZn+9EI=

lib/json/json.go

@@ -4,9 +4,14 @@
 import (
 	"bytes"
 	"encoding/json"
+	"fmt"
+	"math"
+	"sort"
+	"strconv"
 	"sync"
 
 	itn "github.com/1set/starlet/dataconv"
+	"github.com/spyzhov/ajson"
 	stdjson "go.starlark.net/lib/json"
 	"go.starlark.net/starlark"
 	"go.starlark.net/starlarkstruct"
@@ -33,6 +38,10 @@
 				"try_encode": starlark.NewBuiltin(ModuleName+".try_encode", tryEncode),
 				"try_decode": starlark.NewBuiltin(ModuleName+".try_decode", tryDecode),
 				"try_indent": starlark.NewBuiltin(ModuleName+".try_indent", tryIndent),
+				"path":       starlark.NewBuiltin(ModuleName+".path", generateJsonPath(false)),
+				"try_path":   starlark.NewBuiltin(ModuleName+".try_path", generateJsonPath(true)),
+				"eval":       starlark.NewBuiltin(ModuleName+".eval", generateJsonEval(false)),
+				"try_eval":   starlark.NewBuiltin(ModuleName+".try_eval", generateJsonEval(true)),
 			},
 		}
 		for k, v := range stdjson.Module.Members {
@@ -131,3 +140,200 @@
 	}
 	return starlark.Tuple{starlark.String(buf.String()), none}, nil
 }
+
+// generateJsonPath generates a Starlark function that performs a JSONPath query on the given JSON data and returns the matching elements.
+func generateJsonPath(try bool) itn.StarlarkFunc {
+	return func(thread *starlark.Thread, fn *starlark.Builtin, args starlark.Tuple, kwargs []starlark.Tuple) (res starlark.Value, err error) {
+		var (
+			data     starlark.Value
+			pathExpr string
+		)
+		if err := starlark.UnpackArgs(fn.Name(), args, kwargs, "data", &data, "path", &pathExpr); err != nil {
+			if try {
+				return starlark.Tuple{starlark.None, starlark.String(err.Error())}, nil
+			}
+			return none, err
+		}
+
+		jb, err := getJsonBytes(data)
+		if err != nil {
+			if try {
+				return starlark.Tuple{starlark.None, starlark.String(err.Error())}, nil
+			}
+			return none, fmt.Errorf("json.path: %w", err)
+		}
+
+		nodes, err := ajson.JSONPath(jb, pathExpr)
+		if err != nil {
+			if try {
+				return starlark.Tuple{starlark.None, starlark.String(err.Error())}, nil
+			}
+			return none, fmt.Errorf("json.path: %w", err)
+		}
+
+		results, err := ajsonNodesToStarlarkList(nodes)
+		if err != nil {
+			if try {
+				return starlark.Tuple{starlark.None, starlark.String(err.Error())}, nil
+			}
+			return none, fmt.Errorf("json.path: %w", err)
+		}
+
+		if try {
+			return starlark.Tuple{results, starlark.None}, nil
+		}
+		return results, nil
+	}
+}
+
+// generateJsonEval generates a Starlark function that evaluates a JSONPath query with an expression on the given JSON data and returns the evaluation result.
+func generateJsonEval(try bool) itn.StarlarkFunc {
+	return func(thread *starlark.Thread, fn *starlark.Builtin, args starlark.Tuple, kwargs []starlark.Tuple) (res starlark.Value, err error) {
+		var (
+			data starlark.Value
+			expr string
+		)
+		if err := starlark.UnpackArgs(fn.Name(), args, kwargs, "data", &data, "expr", &expr); err != nil {
+			if try {
+				return starlark.Tuple{starlark.None, starlark.String(err.Error())}, nil
+			}
+			return none, err
+		}
+
+		jb, err := getJsonBytes(data)
+		if err != nil {
+			if try {
+				return starlark.Tuple{starlark.None, starlark.String(err.Error())}, nil
+			}
+			return none, fmt.Errorf("json.eval: %w", err)
+		}
+
+		root, err := ajson.Unmarshal(jb)
+		if err != nil {
+			if try {
+				return starlark.Tuple{starlark.None, starlark.String(err.Error())}, nil
+			}
+			return none, fmt.Errorf("json.eval: %w", err)
+		}
+
+		result, err := ajson.Eval(root, expr)
+		if err != nil {
+			if try {
+				return starlark.Tuple{starlark.None, starlark.String(err.Error())}, nil
+			}
+			return none, fmt.Errorf("json.eval: %w", err)
+		}
+
+		val, err := ajsonNodeToStarlarkValue(result)
+		if err != nil {
+			if try {
+				return starlark.Tuple{starlark.None, starlark.String(err.Error())}, nil
+			}
+			return none, fmt.Errorf("json.eval: %w", err)
+		}
+
+		if try {
+			return starlark.Tuple{val, starlark.None}, nil
+		}
+		return val, nil
+	}
+}
+
+// getJsonBytes converts a Starlark value to a JSON byte slice.
+func getJsonBytes(data starlark.Value) ([]byte, error) {
+	switch v := data.(type) {
+	case starlark.String:
+		return []byte(v.GoString()), nil
+	case starlark.Bytes:
+		return []byte(v), nil
+	default:
+		js, err := itn.MarshalStarlarkJSON(data, 0)
+		if err != nil {
+			return nil, err
+		}
+		return []byte(js), nil
+	}
+}
+
+// ajsonNodesToStarlarkList converts a slice of ajson.Node to a Starlark list.
+func ajsonNodesToStarlarkList(nodes []*ajson.Node) (starlark.Value, error) {
+	results := make([]starlark.Value, 0, len(nodes))
+	for _, node := range nodes {
+		val, err := ajsonNodeToStarlarkValue(node)
+		if err != nil {
+			return nil, err
+		}
+		results = append(results, val)
+	}
+	return starlark.NewList(results), nil
+}
+
+// ajsonNodeToStarlarkValue converts an ajson.Node to a Starlark value.
+// It recursively traverses the node tree and constructs the corresponding Starlark values.
+func ajsonNodeToStarlarkValue(node *ajson.Node) (starlark.Value, error) {
+	switch node.Type() {
+	case ajson.Object:
+		dict := &starlark.Dict{}
+		for _, key := range node.Keys() {
+			valNode, err := node.GetKey(key)
+			if err != nil {
+				return nil, err
+			}
+			val, err := ajsonNodeToStarlarkValue(valNode)
+			if err != nil {
+				return nil, err
+			}
+			err = dict.SetKey(starlark.String(key), val)
+			if err != nil {
+				return nil, err
+			}
+		}
+		return dict, nil
+	case ajson.Array:
+		// Use keys and indices to avoid relying on invalid index pointers
+		keys := node.Keys()
+		if len(keys) == 0 {
+			return starlark.NewList(nil), nil
+		}
+		indices := make([]int, len(keys))
+		indexMap := make(map[int]*ajson.Node)
+		for i, key := range keys {
+			idx, err := strconv.Atoi(key)
+			if err != nil {
+				return nil, fmt.Errorf("invalid array index: %v", err)
+			}
+			indices[i] = idx
+			child, err := node.GetIndex(idx)
+			if err != nil {
+				return nil, err
+			}
+			indexMap[idx] = child
+		}
+		sort.Ints(indices)
+		vals := make([]starlark.Value, len(indices))
+		for i, idx := range indices {
+			elem := indexMap[idx]
+			val, err := ajsonNodeToStarlarkValue(elem)
+			if err != nil {
+				return nil, err
+			}
+			vals[i] = val
+		}
+		return starlark.NewList(vals), nil
+	case ajson.String:
+		return starlark.String(node.MustString()), nil
+	case ajson.Numeric:
+		num := node.MustNumeric()
+		if math.Mod(num, 1.0) == 0 {
+			return starlark.MakeInt64(int64(num)), nil
+		} else {
+			return starlark.Float(num), nil
+		}
+	case ajson.Bool:
+		return starlark.Bool(node.MustBool()), nil
+	case ajson.Null:
+		return starlark.None, nil
+	default:
+		return nil, fmt.Errorf("unsupported JSON node type: %v", node.Type())
+	}
+}