jsonmod.go

package jsonmod

import (
	"bytes"
	"encoding/json"
	"fmt"
)

// JSONContext holds context for the current JSON path.
type JSONContext struct {
	ParentType      string // "object", "array" or ""
	LatestObjectKey string // key of the nearest parent object
	LatestArrayIdx  int    // index of the nearest parent array
	Path            string // full path to the current value (e.g. "foo.bar[3].baz")
}

// ModifyJSONValuesFunc defines a function to modify a JSON value with its context.
type ModifyJSONValuesFunc func(c JSONContext, value any) any

func ModifyJSONValues(in []byte, fn ModifyJSONValuesFunc) ([]byte, error) {
	if !json.Valid(in) {
		return nil, fmt.Errorf("invalid JSON")
	}
	rewriter := &jsonRewriter{
		dec: json.NewDecoder(bytes.NewReader(in)),
	}
	rewriter.rewriteJSON(JSONContext{}, fn)
	return rewriter.out.Bytes(), nil
}

type jsonRewriter struct {
	out bytes.Buffer
	dec *json.Decoder
}

func (r *jsonRewriter) rewriteJSONArray(c JSONContext, fn ModifyJSONValuesFunc) {
	c.ParentType = "array"
	c.LatestArrayIdx = 0
	needComma := false
	for r.dec.More() {
		if needComma {
			r.out.WriteByte(',')
		}
		needComma = true
		nc := c
		nc.Path += fmt.Sprintf("[%d]", c.LatestArrayIdx)
		r.rewriteJSON(nc, fn)
		c.LatestArrayIdx++
	}
	r.out.WriteString(fmt.Sprintf("%v", r.mustToken()))
}

func (r *jsonRewriter) rewriteJSONObject(c JSONContext, fn ModifyJSONValuesFunc) {
	c.ParentType = "object"
	needComma := false
	for r.dec.More() {
		if needComma {
			r.out.WriteByte(',')
		}
		needComma = true
		key := r.mustToken()
		r.mustWriteToken(key)
		r.out.WriteByte(':')
		nc := c
		nc.Path += "." + key.(string)
		nc.LatestObjectKey = key.(string)
		r.rewriteJSON(nc, fn)
	}
	r.out.WriteString(fmt.Sprintf("%v", r.mustToken()))
}

func (r *jsonRewriter) rewriteJSON(c JSONContext, fn ModifyJSONValuesFunc) {
	tkn := r.mustToken()
	delim, ok := tkn.(json.Delim)
	if !ok {
		r.mustWriteToken(fn(c, tkn))
		return
	}
	r.out.WriteString(delim.String())
	switch delim {
	case '{':
		r.rewriteJSONObject(c, fn)
	case '[':
		r.rewriteJSONArray(c, fn)
	}
}

func (r *jsonRewriter) mustWriteToken(tkn json.Token) {
	delim, ok := tkn.(json.Delim)
	if ok {
		r.out.WriteString(delim.String())
	}
	b, err := json.Marshal(tkn)
	if err != nil {
		panic(err)
	}
	r.out.Write(b)
}

func (r *jsonRewriter) mustToken() json.Token {
	tkn, err := r.dec.Token()
	if err != nil {
		panic(err)
	}
	return tkn
}