<commit_before>package main

import (
	"encoding/json"
	"io"
	"strings"
)

// JSONIn provides methods of the Input interface
type JSONIn struct {
	reader   *json.Decoder
	firstRow []string
	header   []string
	ajson    []interface{}
	count    int
}

// JSONOut provides methods of the Output interface
type JSONOut struct {
	writer  *json.Encoder
	results []map[string]string
}

func (trdsql *TRDSQL) jsonInputNew(r io.Reader) (Input, error) {
	jr := &JSONIn{}
	jr.reader = json.NewDecoder(r)
	return jr, nil
}

func (jr *JSONIn) firstRead() ([]string, error) {
	var data interface{}
	var dmap map[string]interface{}
	err := jr.reader.Decode(&data)
	if err != nil {
		return nil, err
	}
	switch data.(type) {
	case []interface{}:
		jr.ajson = data.([]interface{})
		jr.count = 0
		kv := jr.ajson[jr.count]
		dmap = kv.(map[string]interface{})
	case map[string]interface{}:
		jr.ajson = nil
		dmap = data.(map[string]interface{})
	}
	for k, v := range dmap {
		jr.header = append(jr.header, k)
		jr.firstRow = append(jr.firstRow, jsonStr(v))
	}
	debug.Printf("Column Name: [%v]", strings.Join(jr.header, ","))
	return jr.header, err
}

func jsonStr(val interface{}) string {
	switch val.(type) {
	case map[string]interface{}:
		str, _ := json.Marshal(val)
		return valString(str)
	case []interface{}:
		str, _ := json.Marshal(val)
		return valString(str)
	default:
		return valString(val)
	}
}

func (jr *JSONIn) firstRowRead(list []interface{}) []interface{} {
	for i := range jr.header {
		list[i] = jr.firstRow[i]
	}
	return list
}

func (jr *JSONIn) rowRead(list []interface{}) ([]interface{}, error) {
	var dmap map[string]interface{}
	if jr.ajson == nil {
		var data interface{}
		err := jr.reader.Decode(&data)
		if err != nil {
			return nil, err
		}
		dmap = data.(map[string]interface{})
	} else {
		jr.count++
		if len(jr.ajson) <= jr.count {
			return nil, io.EOF
		}
		kv := jr.ajson[jr.count]
		dmap = kv.(map[string]interface{})
	}
	for i := range jr.header {
		list[i] = jsonStr(dmap[jr.header[i]])
	}
	return list, nil
}

func (trdsql *TRDSQL) jsonOutNew() Output {
	js := &JSONOut{}
	js.writer = json.NewEncoder(trdsql.outStream)
	js.writer.SetIndent("", "  ")
	return js
}

func (js *JSONOut) first(columns []string) error {<commit_msg>Improve json<commit_after>package main

import (
	"encoding/json"
	"io"
	"strings"
)

// JSONIn provides methods of the Input interface
type JSONIn struct {
	reader   *json.Decoder
	firstRow []string
	header   []string
	ajson    []interface{}
	count    int
}

// JSONOut provides methods of the Output interface
type JSONOut struct {
	writer  *json.Encoder
	results []map[string]string
}

func (trdsql *TRDSQL) jsonInputNew(r io.Reader) (Input, error) {
	jr := &JSONIn{}
	jr.reader = json.NewDecoder(r)
	return jr, nil
}

func (jr *JSONIn) firstRead() ([]string, error) {
	var top interface{}
	err := jr.reader.Decode(&top)
	if err != nil {
		return nil, err
	}
	err = jr.topRow(top)
	if err != nil {
		return nil, err
	}
	debug.Printf("Column Name: [%v]", strings.Join(jr.header, ","))
	return jr.header, err
}

func (jr *JSONIn) topRow(top interface{}) error {
	switch top.(type) {
	case []interface{}:
		jr.ajson = top.([]interface{})
		val := jr.ajson[0]
		switch val.(type) {
		case map[string]interface{}:
			jr.objectFirstRow(val.(map[string]interface{}))
			return nil
		case []interface{}:
			jr.etcFirstRow(val)
		default:
			// ["a","b"]
			jr.ajson = nil
			jr.etcFirstRow(top)
		}
	case map[string]interface{}:
		// {"a":"b"} object
		jr.ajson = nil
		jr.objectFirstRow(top.(map[string]interface{}))
	}
	return nil
}

func (jr *JSONIn) objectFirstRow(obj map[string]interface{}) {
	// {"a":"b"} object
	for k, v := range obj {
		jr.header = append(jr.header, k)
		jr.firstRow = append(jr.firstRow, jsonStr(v))
	}
}

func (jr *JSONIn) etcFirstRow(val interface{}) {
	// array array
	// [["a"],
	//  ["b"]]
	jr.header = append(jr.header, "c1")
	jr.firstRow = append(jr.firstRow, jsonStr(val))
}

func jsonStr(val interface{}) string {
	switch val.(type) {
	case map[string]interface{}:
		str, _ := json.Marshal(val)
		return valString(str)
	case []interface{}:
		str, _ := json.Marshal(val)
		return valString(str)
	default:
		return valString(val)
	}
}

func (jr *JSONIn) firstRowRead(list []interface{}) []interface{} {
	for i := range jr.header {
		list[i] = jr.firstRow[i]
	}
	return list
}

func (jr *JSONIn) rowRead(list []interface{}) ([]interface{}, error) {
	if jr.ajson != nil {
		// [] array
		jr.count++
		if jr.count >= len(jr.ajson) {
			return nil, io.EOF
		}
		list = jr.rowParse(list, jr.ajson[jr.count])
	} else {
		var data interface{}
		// {} object
		err := jr.reader.Decode(&data)
		if err != nil {
			return nil, err
		}
		list = jr.rowParse(list, data)
	}
	return list, nil
}

func (jr *JSONIn) rowParse(list []interface{}, jsonRow interface{}) []interface{} {
	switch jsonRow.(type) {
	case map[string]interface{}:
		dmap := jsonRow.(map[string]interface{})
		for i := range jr.header {
			list[i] = jsonStr(dmap[jr.header[i]])
		}
	default:
		list[0] = jsonStr(jsonRow)
	}
	return list
}

func (trdsql *TRDSQL) jsonOutNew() Output {
	js := &JSONOut{}
	js.writer = json.NewEncoder(trdsql.outStream)
	js.writer.SetIndent("", "  ")
	return js
}

func (js *JSONOut) first(columns []string) error {
{'fraction_non_alphanumeric': 0.17073632405830022, 'fraction_numerical': 0.0007571455612341472, 'mean_word_length': 3.821167883211679, 'pattern_counts': {'":': 2, '<': 4, '<?xml version=': 0, '>': 4, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '1898383', 'n_tokens_mistral': 2140, 'n_tokens_neox': 2058, 'n_words': 486}