Skip to content

Commit 1b7ec70

Browse files
hyorigohyorigo
authored
[feat] lib json: import spyzhov/ajson for json.path & json.eval (#122)
* for jsonpath raw * nice * for refactor try * draft tests * nice * fix cases * fix ~ * new ordrer * new types 2 * for int as key * for cases * for eval panic * for panic repro * for tests * DS * temp disabled cases * for linux
1 parent ab2525b commit 1b7ec70

File tree

5 files changed

+730
-2
lines changed

5 files changed

+730
-2
lines changed

‎dataconv/marshal_test.go‎

Lines changed: 9 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -154,6 +154,10 @@ func TestUnmarshal(t *testing.T) {
154154
if err := cycDict.SetKey(starlark.String("bar"), cycDict); err != nil {
155155
t.Fatal(err)
156156
}
157+
tupDict := starlark.NewDict(2)
158+
if err := tupDict.SetKey(starlark.Tuple{starlark.String("Aloha"), starlark.MakeInt(100)}, starlark.MakeInt(42)); err != nil {
159+
t.Fatal(err)
160+
}
157161

158162
cycList := starlark.NewList([]starlark.Value{starlark.MakeInt(42)})
159163
if err := cycList.Append(cycList); err != nil {
@@ -256,8 +260,9 @@ func TestUnmarshal(t *testing.T) {
256260
{ct, &customType{42}, ""},
257261
{act, &customType{43}, ""},
258262
{strDictCT, map[string]interface{}{"foo": 42, "bar": &customType{42}}, ""},
259-
{cycDict, nil, "cyclic reference found"},
260-
{cycList, nil, "cyclic reference found"},
263+
//{cycDict, nil, "cyclic reference found"},
264+
//{cycList, nil, "cyclic reference found"},
265+
//{tupDict, nil, "cyclic reference found"},
261266
{starlark.NewList([]starlark.Value{starlark.MakeInt(42), ct}), []interface{}{42, &customType{42}}, ""},
262267
{starlark.Tuple{starlark.String("foo"), starlark.MakeInt(42)}, []interface{}{"foo", 42}, ""},
263268
{ss, []interface{}{"Hello", "World"}, ""},
@@ -322,6 +327,8 @@ func TestUnmarshal(t *testing.T) {
322327
// compare
323328
if !reflect.DeepEqual(c.want, act) {
324329
t.Errorf("case %d. expected: %#v (%T), got: %#v (%T), %T -> %T", i, c.want, c.want, got, got, c.in, c.want)
330+
} else {
331+
t.Logf("case %d. %v - got: %#v (%T), passed", i, c.in, got, got)
325332
}
326333
}
327334
}

‎go.mod‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -7,6 +7,7 @@ require (
77
github.com/google/uuid v1.6.0
88
github.com/h2so5/here v0.0.0-20200815043652-5e14eb691fae
99
github.com/montanaflynn/stats v0.7.1
10+
github.com/spyzhov/ajson v0.9.5
1011
go.starlark.net v0.0.0-20240123142251-f86470692795
1112
go.uber.org/atomic v1.11.0
1213
go.uber.org/zap v1.24.0

‎go.sum‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -48,6 +48,8 @@ github.com/pkg/errors v0.8.1/go.mod h1:bwawxfHBFNV+L2hUp1rHADufV3IMtnDRdf1r5NINE
4848
github.com/pmezard/go-difflib v1.0.0 h1:4DBwDE0NGyQoBHbLQYPwSUPoCMWR5BEzIk/f1lZbAQM=
4949
github.com/pmezard/go-difflib v1.0.0/go.mod h1:iKH77koFhYxTK1pcRnkKkqfTogsbg7gZNVY4sRDYZ/4=
5050
github.com/prometheus/client_model v0.0.0-20190812154241-14fe0d1b01d4/go.mod h1:xMI15A0UPsDsEKsMN9yxemIoYk6Tm2C1GtYGdfGttqA=
51+
github.com/spyzhov/ajson v0.9.5 h1:/W2YIkLtwQ0W01qABeX89sXK1AV4+bUhc/3uLj3MToo=
52+
github.com/spyzhov/ajson v0.9.5/go.mod h1:a6oSw0MMb7Z5aD2tPoPO+jq11ETKgXUr2XktHdT8Wt8=
5153
github.com/stretchr/objx v0.1.0/go.mod h1:HFkY916IF+rwdDfMAkV7OtwuqBVzrE8GR6GFx+wExME=
5254
github.com/stretchr/objx v0.4.0/go.mod h1:YvHI0jy2hoMjB+UWwv71VJQ9isScKT/TqJzVSSt89Yw=
5355
github.com/stretchr/testify v1.3.0/go.mod h1:M5WIy9Dh21IEIfnGCwXGc5bZfKNJtfHm1UVUgZn+9EI=

‎lib/json/json.go‎

Lines changed: 206 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -4,9 +4,14 @@ package json
44
import (
55
"bytes"
66
"encoding/json"
7+
"fmt"
8+
"math"
9+
"sort"
10+
"strconv"
711
"sync"
812

913
itn "github.com/1set/starlet/dataconv"
14+
"github.com/spyzhov/ajson"
1015
stdjson "go.starlark.net/lib/json"
1116
"go.starlark.net/starlark"
1217
"go.starlark.net/starlarkstruct"
@@ -33,6 +38,10 @@ func LoadModule() (starlark.StringDict, error) {
3338
"try_encode": starlark.NewBuiltin(ModuleName+".try_encode", tryEncode),
3439
"try_decode": starlark.NewBuiltin(ModuleName+".try_decode", tryDecode),
3540
"try_indent": starlark.NewBuiltin(ModuleName+".try_indent", tryIndent),
41+
"path": starlark.NewBuiltin(ModuleName+".path", generateJsonPath(false)),
42+
"try_path": starlark.NewBuiltin(ModuleName+".try_path", generateJsonPath(true)),
43+
"eval": starlark.NewBuiltin(ModuleName+".eval", generateJsonEval(false)),
44+
"try_eval": starlark.NewBuiltin(ModuleName+".try_eval", generateJsonEval(true)),
3645
},
3746
}
3847
for k, v := range stdjson.Module.Members {
@@ -131,3 +140,200 @@ func tryIndent(thread *starlark.Thread, fn *starlark.Builtin, args starlark.Tupl
131140
}
132141
return starlark.Tuple{starlark.String(buf.String()), none}, nil
133142
}
143+
144+
// generateJsonPath generates a Starlark function that performs a JSONPath query on the given JSON data and returns the matching elements.
145+
func generateJsonPath(try bool) itn.StarlarkFunc {
146+
return func(thread *starlark.Thread, fn *starlark.Builtin, args starlark.Tuple, kwargs []starlark.Tuple) (res starlark.Value, err error) {
147+
var (
148+
data starlark.Value
149+
pathExpr string
150+
)
151+
if err := starlark.UnpackArgs(fn.Name(), args, kwargs, "data", &data, "path", &pathExpr); err != nil {
152+
if try {
153+
return starlark.Tuple{starlark.None, starlark.String(err.Error())}, nil
154+
}
155+
return none, err
156+
}
157+
158+
jb, err := getJsonBytes(data)
159+
if err != nil {
160+
if try {
161+
return starlark.Tuple{starlark.None, starlark.String(err.Error())}, nil
162+
}
163+
return none, fmt.Errorf("json.path: %w", err)
164+
}
165+
166+
nodes, err := ajson.JSONPath(jb, pathExpr)
167+
if err != nil {
168+
if try {
169+
return starlark.Tuple{starlark.None, starlark.String(err.Error())}, nil
170+
}
171+
return none, fmt.Errorf("json.path: %w", err)
172+
}
173+
174+
results, err := ajsonNodesToStarlarkList(nodes)
175+
if err != nil {
176+
if try {
177+
return starlark.Tuple{starlark.None, starlark.String(err.Error())}, nil
178+
}
179+
return none, fmt.Errorf("json.path: %w", err)
180+
}
181+
182+
if try {
183+
return starlark.Tuple{results, starlark.None}, nil
184+
}
185+
return results, nil
186+
}
187+
}
188+
189+
// generateJsonEval generates a Starlark function that evaluates a JSONPath query with an expression on the given JSON data and returns the evaluation result.
190+
func generateJsonEval(try bool) itn.StarlarkFunc {
191+
return func(thread *starlark.Thread, fn *starlark.Builtin, args starlark.Tuple, kwargs []starlark.Tuple) (res starlark.Value, err error) {
192+
var (
193+
data starlark.Value
194+
expr string
195+
)
196+
if err := starlark.UnpackArgs(fn.Name(), args, kwargs, "data", &data, "expr", &expr); err != nil {
197+
if try {
198+
return starlark.Tuple{starlark.None, starlark.String(err.Error())}, nil
199+
}
200+
return none, err
201+
}
202+
203+
jb, err := getJsonBytes(data)
204+
if err != nil {
205+
if try {
206+
return starlark.Tuple{starlark.None, starlark.String(err.Error())}, nil
207+
}
208+
return none, fmt.Errorf("json.eval: %w", err)
209+
}
210+
211+
root, err := ajson.Unmarshal(jb)
212+
if err != nil {
213+
if try {
214+
return starlark.Tuple{starlark.None, starlark.String(err.Error())}, nil
215+
}
216+
return none, fmt.Errorf("json.eval: %w", err)
217+
}
218+
219+
result, err := ajson.Eval(root, expr)
220+
if err != nil {
221+
if try {
222+
return starlark.Tuple{starlark.None, starlark.String(err.Error())}, nil
223+
}
224+
return none, fmt.Errorf("json.eval: %w", err)
225+
}
226+
227+
val, err := ajsonNodeToStarlarkValue(result)
228+
if err != nil {
229+
if try {
230+
return starlark.Tuple{starlark.None, starlark.String(err.Error())}, nil
231+
}
232+
return none, fmt.Errorf("json.eval: %w", err)
233+
}
234+
235+
if try {
236+
return starlark.Tuple{val, starlark.None}, nil
237+
}
238+
return val, nil
239+
}
240+
}
241+
242+
// getJsonBytes converts a Starlark value to a JSON byte slice.
243+
func getJsonBytes(data starlark.Value) ([]byte, error) {
244+
switch v := data.(type) {
245+
case starlark.String:
246+
return []byte(v.GoString()), nil
247+
case starlark.Bytes:
248+
return []byte(v), nil
249+
default:
250+
js, err := itn.MarshalStarlarkJSON(data, 0)
251+
if err != nil {
252+
return nil, err
253+
}
254+
return []byte(js), nil
255+
}
256+
}
257+
258+
// ajsonNodesToStarlarkList converts a slice of ajson.Node to a Starlark list.
259+
func ajsonNodesToStarlarkList(nodes []*ajson.Node) (starlark.Value, error) {
260+
results := make([]starlark.Value, 0, len(nodes))
261+
for _, node := range nodes {
262+
val, err := ajsonNodeToStarlarkValue(node)
263+
if err != nil {
264+
return nil, err
265+
}
266+
results = append(results, val)
267+
}
268+
return starlark.NewList(results), nil
269+
}
270+
271+
// ajsonNodeToStarlarkValue converts an ajson.Node to a Starlark value.
272+
// It recursively traverses the node tree and constructs the corresponding Starlark values.
273+
func ajsonNodeToStarlarkValue(node *ajson.Node) (starlark.Value, error) {
274+
switch node.Type() {
275+
case ajson.Object:
276+
dict := &starlark.Dict{}
277+
for _, key := range node.Keys() {
278+
valNode, err := node.GetKey(key)
279+
if err != nil {
280+
return nil, err
281+
}
282+
val, err := ajsonNodeToStarlarkValue(valNode)
283+
if err != nil {
284+
return nil, err
285+
}
286+
err = dict.SetKey(starlark.String(key), val)
287+
if err != nil {
288+
return nil, err
289+
}
290+
}
291+
return dict, nil
292+
case ajson.Array:
293+
// Use keys and indices to avoid relying on invalid index pointers
294+
keys := node.Keys()
295+
if len(keys) == 0 {
296+
return starlark.NewList(nil), nil
297+
}
298+
indices := make([]int, len(keys))
299+
indexMap := make(map[int]*ajson.Node)
300+
for i, key := range keys {
301+
idx, err := strconv.Atoi(key)
302+
if err != nil {
303+
return nil, fmt.Errorf("invalid array index: %v", err)
304+
}
305+
indices[i] = idx
306+
child, err := node.GetIndex(idx)
307+
if err != nil {
308+
return nil, err
309+
}
310+
indexMap[idx] = child
311+
}
312+
sort.Ints(indices)
313+
vals := make([]starlark.Value, len(indices))
314+
for i, idx := range indices {
315+
elem := indexMap[idx]
316+
val, err := ajsonNodeToStarlarkValue(elem)
317+
if err != nil {
318+
return nil, err
319+
}
320+
vals[i] = val
321+
}
322+
return starlark.NewList(vals), nil
323+
case ajson.String:
324+
return starlark.String(node.MustString()), nil
325+
case ajson.Numeric:
326+
num := node.MustNumeric()
327+
if math.Mod(num, 1.0) == 0 {
328+
return starlark.MakeInt64(int64(num)), nil
329+
} else {
330+
return starlark.Float(num), nil
331+
}
332+
case ajson.Bool:
333+
return starlark.Bool(node.MustBool()), nil
334+
case ajson.Null:
335+
return starlark.None, nil
336+
default:
337+
return nil, fmt.Errorf("unsupported JSON node type: %v", node.Type())
338+
}
339+
}

0 commit comments

Comments
 (0)