journal/gojq-extended/cli.go
2024-11-10 08:13:25 +00:00

475 lines
12 KiB
Go

package main
import (
"errors"
"fmt"
"io"
"os"
"runtime"
"strings"
"github.com/mattn/go-isatty"
"github.com/itchyny/gojq"
)
const name = "gojq"
const version = "0.12.16"
var revision = "HEAD"
const (
exitCodeOK = iota
exitCodeFalsyErr
exitCodeFlagParseErr
exitCodeCompileErr
exitCodeNoValueErr
exitCodeDefaultErr
)
type cli struct {
inStream io.Reader
outStream io.Writer
errStream io.Writer
outputRaw bool
outputRaw0 bool
outputJoin bool
outputCompact bool
outputIndent *int
outputTab bool
inputRaw bool
inputStream bool
inputSlurp bool
argnames []string
argvalues []any
exitCodeError error
}
type flagopts struct {
OutputRaw bool `short:"r" long:"raw-output" description:"output raw strings"`
OutputRaw0 bool `long:"raw-output0" description:"implies -r with NUL character delimiter"`
OutputJoin bool `short:"j" long:"join-output" description:"implies -r with no newline delimiter"`
OutputCompact bool `short:"c" long:"compact-output" description:"output without pretty-printing"`
OutputIndent *int `long:"indent" description:"number of spaces for indentation"`
OutputTab bool `long:"tab" description:"use tabs for indentation"`
OutputColor bool `short:"C" long:"color-output" description:"output with colors even if piped"`
OutputMono bool `short:"M" long:"monochrome-output" description:"output without colors"`
InputNull bool `short:"n" long:"null-input" description:"use null as input value"`
InputRaw bool `short:"R" long:"raw-input" description:"read input as raw strings"`
InputStream bool `long:"stream" description:"parse input in stream fashion"`
InputSlurp bool `short:"s" long:"slurp" description:"read all inputs into an array"`
FromFile bool `short:"f" long:"from-file" description:"load query from file"`
ModulePaths []string `short:"L" description:"directory to search modules from"`
Arg map[string]string `long:"arg" description:"set a string value to a variable"`
ArgJSON map[string]string `long:"argjson" description:"set a JSON value to a variable"`
SlurpFile map[string]string `long:"slurpfile" description:"set the JSON contents of a file to a variable"`
RawFile map[string]string `long:"rawfile" description:"set the contents of a file to a variable"`
Args []any `long:"args" positional:"" description:"consume remaining arguments as positional string values"`
JSONArgs []any `long:"jsonargs" positional:"" description:"consume remaining arguments as positional JSON values"`
ExitStatus bool `short:"e" long:"exit-status" description:"exit 1 when the last value is false or null"`
Version bool `short:"v" long:"version" description:"display version information"`
Help bool `short:"h" long:"help" description:"display this help information"`
}
var addDefaultModulePaths = true
func (cli *cli) run(args []string) int {
if err := cli.runInternal(args); err != nil {
if _, ok := err.(interface{ isEmptyError() }); !ok {
fmt.Fprintf(cli.errStream, "%s: %s\n", name, err)
}
if err, ok := err.(interface{ ExitCode() int }); ok {
return err.ExitCode()
}
return exitCodeDefaultErr
}
return exitCodeOK
}
func (cli *cli) runInternal(args []string) (err error) {
var opts flagopts
args, err = parseFlags(args, &opts)
if err != nil {
return &flagParseError{err}
}
if opts.Help {
fmt.Fprintf(cli.outStream, `%[1]s - Go implementation of jq
Version: %s (rev: %s/%s)
Synopsis:
%% echo '{"foo": 128}' | %[1]s '.foo'
Usage:
%[1]s [OPTIONS]
`,
name, version, revision, runtime.Version())
fmt.Fprintln(cli.outStream, formatFlags(&opts))
return nil
}
if opts.Version {
fmt.Fprintf(cli.outStream, "%s %s (rev: %s/%s)\n", name, version, revision, runtime.Version())
return nil
}
cli.outputRaw, cli.outputRaw0, cli.outputJoin,
cli.outputCompact, cli.outputIndent, cli.outputTab =
opts.OutputRaw, opts.OutputRaw0, opts.OutputJoin,
opts.OutputCompact, opts.OutputIndent, opts.OutputTab
defer func(x bool) { noColor = x }(noColor)
if opts.OutputColor || opts.OutputMono {
noColor = opts.OutputMono
} else if os.Getenv("NO_COLOR") != "" || os.Getenv("TERM") == "dumb" {
noColor = true
} else {
f, ok := cli.outStream.(interface{ Fd() uintptr })
noColor = !(ok && (isatty.IsTerminal(f.Fd()) || isatty.IsCygwinTerminal(f.Fd())))
}
if !noColor {
if colors := os.Getenv("GOJQ_COLORS"); colors != "" {
if err := setColors(colors); err != nil {
return err
}
}
}
if i := cli.outputIndent; i != nil {
if *i > 9 {
return fmt.Errorf("too many indentation count: %d", *i)
} else if *i < 0 {
return fmt.Errorf("negative indentation count: %d", *i)
}
}
cli.inputRaw, cli.inputStream, cli.inputSlurp =
opts.InputRaw, opts.InputStream, opts.InputSlurp
for k, v := range opts.Arg {
cli.argnames = append(cli.argnames, "$"+k)
cli.argvalues = append(cli.argvalues, v)
}
for k, v := range opts.ArgJSON {
val, _ := newJSONInputIter(strings.NewReader(v), "$"+k).Next()
if err, ok := val.(error); ok {
return err
}
cli.argnames = append(cli.argnames, "$"+k)
cli.argvalues = append(cli.argvalues, val)
}
for k, v := range opts.SlurpFile {
val, err := slurpFile(v)
if err != nil {
return err
}
cli.argnames = append(cli.argnames, "$"+k)
cli.argvalues = append(cli.argvalues, val)
}
for k, v := range opts.RawFile {
val, err := os.ReadFile(v)
if err != nil {
return err
}
cli.argnames = append(cli.argnames, "$"+k)
cli.argvalues = append(cli.argvalues, string(val))
}
named := make(map[string]any, len(cli.argnames))
for i, name := range cli.argnames {
named[name[1:]] = cli.argvalues[i]
}
positional := opts.Args
for i, v := range opts.JSONArgs {
if v != nil {
val, _ := newJSONInputIter(strings.NewReader(v.(string)), "--jsonargs").Next()
if err, ok := val.(error); ok {
return err
}
if i < len(positional) {
positional[i] = val
} else {
positional = append(positional, val)
}
}
}
cli.argnames = append(cli.argnames, "$ARGS")
cli.argvalues = append(cli.argvalues, map[string]any{
"named": named,
"positional": positional,
})
var arg, fname string
if opts.FromFile {
if len(args) == 0 {
return errors.New("expected a query file for flag `-f'")
}
src, err := os.ReadFile(args[0])
if err != nil {
return err
}
arg, args, fname = string(src), args[1:], args[0]
} else if len(args) == 0 {
arg = "."
} else {
arg, args, fname = strings.TrimSpace(args[0]), args[1:], "<arg>"
}
if opts.ExitStatus {
cli.exitCodeError = &exitCodeError{exitCodeNoValueErr}
defer func() {
if _, ok := err.(interface{ ExitCode() int }); !ok {
err = cli.exitCodeError
}
}()
}
query, err := gojq.Parse(arg)
if err != nil {
return &queryParseError{fname, arg, err}
}
modulePaths := opts.ModulePaths
if len(modulePaths) == 0 && addDefaultModulePaths {
modulePaths = []string{"~/.jq", "$ORIGIN/../lib/gojq", "$ORIGIN/../lib"}
}
iter := cli.createInputIter(args)
defer iter.Close()
code, err := gojq.Compile(query,
gojq.WithModuleLoader(gojq.NewModuleLoader(modulePaths)),
gojq.WithEnvironLoader(os.Environ),
gojq.WithVariables(cli.argnames),
gojq.WithFunction("debug", 0, 0,
func(errStream io.Writer) func(any, []any) any {
indent := 2
if cli.outputCompact {
indent = 0
} else if cli.outputTab {
indent = 1
} else if i := cli.outputIndent; i != nil {
indent = *i
}
return func(v any, _ []any) any {
if err := newEncoder(false, indent).
marshal([]any{"DEBUG:", v}, cli.errStream); err != nil {
return err
}
if _, err := cli.errStream.Write([]byte{'\n'}); err != nil {
return err
}
return v
}
}(cli.errStream),
),
gojq.WithFunction("stderr", 0, 0,
func(errStream io.Writer) func(any, []any) any {
return func(v any, _ []any) any {
if err := (&rawMarshaler{m: newEncoder(false, 0)}).
marshal(v, cli.errStream); err != nil {
return err
}
return v
}
}(cli.errStream),
),
gojq.WithFunction("_readFile", 1, 1,
func(_input any, args []any) any {
filename := args[0].(string)
info, err := os.Stat(filename)
if err != nil {
return fmt.Errorf("file not found")
}
if info.IsDir() {
return fmt.Errorf("file is a directory")
}
content, err := os.ReadFile(filename)
if err != nil {
return fmt.Errorf("file could not be read")
}
return string(content)
},
),
gojq.WithFunction("_writeFileString", 1, 1,
func(input any, args []any) any {
filename := args[0].(string)
content := ""
switch input := input.(type) {
case string:
content = input
default:
return fmt.Errorf("invalid type passed to _writeFileString")
}
file, err := os.Create(filename)
if err != nil {
return err
}
_, err = file.Write([]byte(content))
if err != nil {
return err
}
err = file.Close()
if err != nil {
return err
}
return nil
},
),
gojq.WithFunction("input_filename", 0, 0,
func(iter inputIter) func(any, []any) any {
return func(any, []any) any {
if fname := iter.Name(); fname != "" && (len(args) > 0 || !opts.InputNull) {
return fname
}
return nil
}
}(iter),
),
gojq.WithInputIter(iter),
)
if err != nil {
if err, ok := err.(interface {
QueryParseError() (string, string, error)
}); ok {
name, query, err := err.QueryParseError()
return &queryParseError{name, query, err}
}
if err, ok := err.(interface {
JSONParseError() (string, string, error)
}); ok {
fname, contents, err := err.JSONParseError()
return &compileError{&jsonParseError{fname, contents, 0, err}}
}
return &compileError{err}
}
if opts.InputNull {
iter = newNullInputIter()
}
return cli.process(iter, code)
}
func slurpFile(name string) (any, error) {
iter := newSlurpInputIter(
newFilesInputIter(newJSONInputIter, []string{name}, nil),
)
defer iter.Close()
val, _ := iter.Next()
if err, ok := val.(error); ok {
return nil, err
}
return val, nil
}
func (cli *cli) createInputIter(args []string) (iter inputIter) {
var newIter func(io.Reader, string) inputIter
switch {
case cli.inputRaw:
if cli.inputSlurp {
newIter = newReadAllIter
} else {
newIter = newRawInputIter
}
case cli.inputStream:
newIter = newStreamInputIter
default:
newIter = newJSONInputIter
}
if cli.inputSlurp {
defer func() {
if cli.inputRaw {
iter = newSlurpRawInputIter(iter)
} else {
iter = newSlurpInputIter(iter)
}
}()
}
if len(args) == 0 {
return newIter(cli.inStream, "<stdin>")
}
return newFilesInputIter(newIter, args, cli.inStream)
}
func (cli *cli) process(iter inputIter, code *gojq.Code) error {
var err error
for {
v, ok := iter.Next()
if !ok {
break
}
if e, ok := v.(error); ok {
fmt.Fprintf(cli.errStream, "%s: %s\n", name, e)
err = e
continue
}
if e := cli.printValues(code.Run(v, cli.argvalues...)); e != nil {
if e, ok := e.(*gojq.HaltError); ok {
if v := e.Value(); v != nil {
if str, ok := v.(string); ok {
cli.errStream.Write([]byte(str))
} else {
bs, _ := gojq.Marshal(v)
cli.errStream.Write(bs)
cli.errStream.Write([]byte{'\n'})
}
}
err = e
break
}
fmt.Fprintf(cli.errStream, "%s: %s\n", name, e)
err = e
}
}
if err != nil {
return &emptyError{err}
}
return nil
}
func (cli *cli) printValues(iter gojq.Iter) error {
m := cli.createMarshaler()
for {
v, ok := iter.Next()
if !ok {
break
}
if err, ok := v.(error); ok {
return err
}
if err := m.marshal(v, cli.outStream); err != nil {
return err
}
if cli.exitCodeError != nil {
if v == nil || v == false {
cli.exitCodeError = &exitCodeError{exitCodeFalsyErr}
} else {
cli.exitCodeError = &exitCodeError{exitCodeOK}
}
}
if cli.outputRaw0 {
cli.outStream.Write([]byte{'\x00'})
} else if !cli.outputJoin {
cli.outStream.Write([]byte{'\n'})
}
}
return nil
}
func (cli *cli) createMarshaler() marshaler {
indent := 2
if cli.outputCompact {
indent = 0
} else if cli.outputTab {
indent = 1
} else if i := cli.outputIndent; i != nil {
indent = *i
}
f := newEncoder(cli.outputTab, indent)
if cli.outputRaw || cli.outputRaw0 || cli.outputJoin {
return &rawMarshaler{f, cli.outputRaw0}
}
return f
}