add --library flag to load additional Rudi scripts

Author: xrstf

cmd/rudi/cmd/console/command.go

@@ -47,7 +47,7 @@ var replCommands = map[string]replCommandFunc{
 	"help": helpCommand,
 }
 
-func Run(handler *util.SignalHandler, opts *options.Options, args []string, rudiVersion string) error {
+func Run(handler *util.SignalHandler, opts *options.Options, library rudi.Program, args []string, rudiVersion string) error {
 	rl, err := readline.New("⮞ ")
 	if err != nil {
 		return fmt.Errorf("failed to setup readline prompt: %w", err)
@@ -67,6 +67,15 @@ func Run(handler *util.SignalHandler, opts *options.Options, args []string, rudi
 	fmt.Println("Type `help` for more information, `exit` or Ctrl-D to exit, Ctrl-C to interrupt statements.")
 	fmt.Println("")
 
+	// Evaluate the library (its return value is irrelevant, as the main program has to have
+	// at least 1 statement, which will overwrite the total return value anyway).
+	if library != nil {
+		_, err = runProgram(handler, rudiCtx, library)
+		if err != nil {
+			return fmt.Errorf("failed to evaluate library: %w", err)
+		}
+	}
+
 	for {
 		line, err := rl.Readline()
 
@@ -125,13 +134,8 @@ func processInput(handler *util.SignalHandler, rudiCtx types.Context, opts *opti
 		return false, err
 	}
 
-	// allow to interrupt the statement
-	ctx, cancel := context.WithCancel(context.Background())
-	defer cancel()
-
-	handler.SetCancelFn(cancel)
-
-	evaluated, err := program.RunContext(rudiCtx.WithGoContext(ctx))
+	// run the program
+	evaluated, err := runProgram(handler, rudiCtx, program)
 	if err != nil {
 		return false, err
 	}
@@ -149,3 +153,13 @@ func processInput(handler *util.SignalHandler, rudiCtx types.Context, opts *opti
 
 	return false, nil
 }
+
+func runProgram(handler *util.SignalHandler, rudiCtx types.Context, prog rudi.Program) (any, error) {
+	// allow to interrupt the statement
+	ctx, cancel := context.WithCancel(context.Background())
+	defer cancel()
+
+	handler.SetCancelFn(cancel)
+
+	return prog.RunContext(rudiCtx.WithGoContext(ctx))
+}

cmd/rudi/cmd/script/command.go

@@ -16,7 +16,7 @@ import (
 	"go.xrstf.de/rudi/cmd/rudi/util"
 )
 
-func Run(handler *util.SignalHandler, opts *options.Options, args []string) error {
+func Run(handler *util.SignalHandler, opts *options.Options, library rudi.Program, args []string) error {
 	// determine input script to evaluate
 	var (
 		script     string
@@ -39,7 +39,7 @@ func Run(handler *util.SignalHandler, opts *options.Options, args []string) erro
 		// consume one arg for the script
 		script = args[0]
 		args = args[1:]
-		scriptName = "(stdin)"
+		scriptName = "(cli)"
 	}
 
 	// parse the script
@@ -75,6 +75,15 @@ func Run(handler *util.SignalHandler, opts *options.Options, args []string) erro
 
 	handler.SetCancelFn(cancel)
 
+	// Evaluate the library (its return value is irrelevant, as the main program has to have
+	// at least 1 statement, which will overwrite the total return value anyway).
+	if library != nil {
+		_, err = library.RunContext(rudiCtx.WithGoContext(subCtx))
+		if err != nil {
+			return fmt.Errorf("failed to evaluate script: %w", err)
+		}
+	}
+
 	// evaluate the script
 	evaluated, err := program.RunContext(rudiCtx.WithGoContext(subCtx))
 	if err != nil {

cmd/rudi/encoding/decode.go

@@ -16,6 +16,26 @@ import (
 	"gopkg.in/yaml.v3"
 )
 
+func decodeYaml(input io.Reader) ([]any, error) {
+	decoder := yaml.NewDecoder(input)
+
+	documents := []any{}
+	for {
+		var doc any
+		if err := decoder.Decode(&doc); err != nil {
+			if errors.Is(err, io.EOF) {
+				break
+			}
+
+			return nil, fmt.Errorf("failed to parse file as YAML: %w", err)
+		}
+
+		documents = append(documents, doc)
+	}
+
+	return documents, nil
+}
+
 func Decode(input io.Reader, enc types.Encoding) (any, error) {
 	var data any
 
@@ -41,6 +61,21 @@ func Decode(input io.Reader, enc types.Encoding) (any, error) {
 		}
 
 	case types.YamlEncoding:
+		decoded, err := decodeYaml(input)
+		if err != nil {
+			return nil, fmt.Errorf("failed to parse file as JSON5: %w", err)
+		}
+
+		switch len(decoded) {
+		case 0:
+			data = nil
+		case 1:
+			data = decoded[0]
+		default:
+			data = decoded
+		}
+
+	case types.YamlDocumentsEncoding:
 		decoder := yaml.NewDecoder(input)
 
 		documents := []any{}
@@ -57,11 +92,7 @@ func Decode(input io.Reader, enc types.Encoding) (any, error) {
 			documents = append(documents, doc)
 		}
 
-		if len(documents) == 1 {
-			data = documents[0]
-		} else {
-			data = documents
-		}
+		data = documents
 
 	case types.TomlEncoding:
 		decoder := toml.NewDecoder(input)

cmd/rudi/main.go

@@ -9,6 +9,7 @@ import (
 	"os"
 	"runtime"
 	"runtime/debug"
+	"strings"
 
 	"go.xrstf.de/rudi"
 	"go.xrstf.de/rudi/cmd/rudi/batteries"
@@ -114,16 +115,42 @@ func main() {
 
 	handler := util.SetupSignalHandler()
 
-	if opts.Interactive || len(args) == 0 {
-		if err := console.Run(handler, &opts, args, BuildTag); err != nil {
+	// load all --library files and assemble a single base script for both console/script mode
+	baseScripts := []string{}
+	for _, filename := range opts.LibraryFiles {
+		_, script, err := util.ParseFile(filename)
+		if err != nil {
+			fmt.Printf("Error: library %q: %v\n", filename, err)
+			os.Exit(1)
+		}
+
+		baseScripts = append(baseScripts, script)
+	}
+
+	// parse the base script
+	var baseProgram rudi.Program
+
+	if len(baseScripts) > 0 {
+		var err error
+
+		baseProgram, err = rudi.Parse("(library)", strings.Join(baseScripts, "\n"))
+		if err != nil {
+			// This should never happen, each script was already syntax-checked.
+			fmt.Printf("Error: library: %v\n", err)
+			os.Exit(1)
+		}
+	}
+
+	if opts.Interactive {
+		if err := console.Run(handler, &opts, baseProgram, args, BuildTag); err != nil {
 			fmt.Printf("Error: %v\n", err)
 			os.Exit(1)
 		}
 
 		return
 	}
 
-	if err := script.Run(handler, &opts, args); err != nil {
+	if err := script.Run(handler, &opts, baseProgram, args); err != nil {
 		parseErr := &rudi.ParseError{}
 		if errors.As(err, parseErr) {
 			fmt.Println(parseErr.Snippet())

cmd/rudi/options/options.go

@@ -21,6 +21,7 @@ type Options struct {
 	ShowHelp                 bool
 	Interactive              bool
 	ScriptFile               string
+	LibraryFiles             []string
 	StdinFormat              types.Encoding
 	OutputFormat             types.Encoding
 	PrintAst                 bool
@@ -49,6 +50,7 @@ func (o *Options) AddFlags(fs *pflag.FlagSet) {
 
 	fs.BoolVarP(&o.Interactive, "interactive", "i", o.Interactive, "Start an interactive REPL to run expressions.")
 	fs.StringVarP(&o.ScriptFile, "script", "s", o.ScriptFile, "Load Rudi script from file instead of first argument (only in non-interactive mode).")
+	fs.StringArrayVarP(&o.LibraryFiles, "library", "l", o.LibraryFiles, "Load additional Rudi file(s) to be be evaluated before the script (can be given multiple times).")
 	fs.StringArrayVar(&o.extraVariableFlags, "var", o.extraVariableFlags, "Define additional global variables (can be given multiple times).")
 	stdinFormatFlag.Add(fs, "stdin-format", "f", "What data format is used for data provided on stdin")
 	outputFormatFlag.Add(fs, "output-format", "o", "What data format to use for outputting data")
@@ -60,10 +62,6 @@ func (o *Options) AddFlags(fs *pflag.FlagSet) {
 }
 
 func (o *Options) Validate() error {
-	if o.Interactive && o.ScriptFile != "" {
-		return errors.New("cannot combine --interactive with --script")
-	}
-
 	if o.Interactive && o.PrintAst {
 		return errors.New("cannot combine --interactive with --debug-ast")
 	}
@@ -72,6 +70,10 @@ func (o *Options) Validate() error {
 		return fmt.Errorf("invalid --var flags: %w", err)
 	}
 
+	if err := o.validateLibraryFiles(); err != nil {
+		return fmt.Errorf("invalid --library flags: %w", err)
+	}
+
 	return nil
 }
 
@@ -148,3 +150,17 @@ func (o *Options) parseExtraVariable(flagValue string) (string, any, error) {
 
 	return varName, varData, nil
 }
+
+func (o *Options) validateLibraryFiles() error {
+	for _, file := range o.LibraryFiles {
+		info, err := os.Stat(file)
+		if err != nil {
+			return fmt.Errorf("invalid library %q: %w", file, err)
+		}
+		if info.IsDir() {
+			return fmt.Errorf("invalid library %q: is directory", file)
+		}
+	}
+
+	return nil
+}

cmd/rudi/types/const.go

@@ -41,14 +41,12 @@ var (
 		TomlEncoding,
 	}
 
-	// InputEncodings contains all valid encodings for reading data.
-	// Note that YAML is always read in multi-document mode, hence
-	// YamlDocumentsEncoding is not part of this list.
 	InputEncodings = []Encoding{
 		RawEncoding,
 		JsonEncoding,
 		Json5Encoding,
 		YamlEncoding,
+		YamlDocumentsEncoding,
 		TomlEncoding,
 	}
 

cmd/rudi/util/files.go

@@ -10,11 +10,28 @@ import (
 	"path/filepath"
 	"strings"
 
+	"go.xrstf.de/rudi"
 	"go.xrstf.de/rudi/cmd/rudi/encoding"
 	"go.xrstf.de/rudi/cmd/rudi/options"
 	"go.xrstf.de/rudi/cmd/rudi/types"
 )
 
+func ParseFile(filename string) (rudi.Program, string, error) {
+	content, err := os.ReadFile(filename)
+	if err != nil {
+		return nil, "", fmt.Errorf("failed to read script: %w", err)
+	}
+
+	script := strings.TrimSpace(string(content))
+
+	prog, err := rudi.Parse(filename, script)
+	if err != nil {
+		return nil, "", fmt.Errorf("failed to parse: %w", err)
+	}
+
+	return prog, script, nil
+}
+
 func LoadFiles(opts *options.Options, filenames []string) ([]any, error) {
 	results := make([]any, len(filenames))