1// Package nushell runs a notebook's code cells: each is a nushell2// pipeline run by `nu` as a separate process, reading other cells'3// outputs and ranges of sheets as variables ($sales, $selection), and4// printing what it made as NUON, which the notebook keeps as the cell's5// output. Tables reach nu as NUON in files it opens, never as text6// spliced into the pipeline. The package knows neither the terminal nor7// the notebook: the UI hands it a Job and gets back NUON.8package nushell910import (11 "bytes"12 "context"13 "encoding/base64"14 "errors"15 "fmt"16 "io"17 "os"18 "os/exec"19 "slices"20 "strconv"21 "strings"22 "time"2324 "github.com/charmbracelet/x/ansi"2526 "src.rest/cat/012/internal/nuon"27)2829// Job is a pipeline to run.30type Job struct {31 Command string32 // Tables are what the pipeline reads, by variable name, as NUON:33 // each becomes the variable of its name.34 Tables map[string][]byte35 // Config runs nu with the user's config files rather than without.36 Config bool37 // Vars are the variables the run hands back beside its output: the38 // command ends in a record of its output, as __out, and each of39 // them, and the run prints each as a line of its own (ExecVars).40 Vars []string41 // IDE asks nu about the script rather than running it: the flags42 // that ask (--ide-ast, or --ide-complete 14), nu reading the script43 // from a file named after them, without config files or the44 // standard library (see ide.go).45 IDE []string46}4748// Runner runs a job's script, writing what it prints to stdout. A test49// substitutes a fake for the real nu.50type Runner interface {51 Run(ctx context.Context, job Job, script string, stdout io.Writer) error52}5354// ErrMissing is what running a cell says when nu isn't installed.55var ErrMissing = errors.New("nu isn't installed or isn't on your PATH: see docs/nushell/notebooks.md")5657// Error is a pipeline that failed: nu's message, and all it wrote to58// standard error.59type Error struct {60 Msg string61 Stderr string62}6364func (e *Error) Error() string { return e.Msg }6566// Help is the help line of nu's message, "help: ...", or "".67func (e *Error) Help() string {68 for line := range strings.SplitSeq(e.Stderr, "\n") {69 if line = strings.TrimSpace(line); strings.HasPrefix(line, "help:") {70 return line71 }72 }73 return ""74}7576// Script is what nu runs for job: the tables bound to their names, then77// the pipeline, its output as NUON. The files are named by environment78// variables (NU012_TABLE_0 and so on) the runner sets, so no data and no79// path is written into the script.80func Script(job Job, names []string) string {81 var b strings.Builder82 for i, name := range names {83 fmt.Fprintf(&b, "let %s = (open --raw $env.NU012_TABLE_%d | from nuon)\n", name, i)84 }85 b.WriteString("do {\n" + job.Command + "\n} | ")86 if len(job.Vars) > 0 {87 b.WriteString(`items {|k, v| $"($k) ($v | to nuon | encode base64)" } | str join "\n"` + "\n")88 } else {89 b.WriteString("to nuon\n")90 }91 return b.String()92}9394// outVar names the output among a run's variables.95const outVar = "__out"9697// ExecVars runs job as Exec does, returning its output and, when it98// hands back variables (Job.Vars), their values, each as NUON. They're99// the last lines it printed, each a name and its NUON as base64, so100// what the pipeline printed itself before them doesn't mix in.101func ExecVars(ctx context.Context, r Runner, job Job, timeout time.Duration, maxBytes int) ([]byte, map[string][]byte, error) {102 out, err := Exec(ctx, r, job, timeout, maxBytes)103 if err != nil || len(job.Vars) == 0 {104 return out, nil, err105 }106 lines := strings.Split(string(out), "\n")107 vars := map[string][]byte{}108 for _, line := range lines[max(len(lines)-len(job.Vars)-1, 0):] {109 name, enc, _ := strings.Cut(strings.TrimSpace(line), " ")110 if data, err := base64.StdEncoding.DecodeString(enc); err == nil && (name == outVar || slices.Contains(job.Vars, name)) {111 vars[name] = data112 }113 }114 nuon, ok := vars[outVar]115 if !ok || len(vars) != len(job.Vars)+1 {116 return nil, nil, errors.New("the run didn't hand back its variables")117 }118 delete(vars, outVar)119 return nuon, vars, nil120}121122// ErrTooLarge is a run that printed more than it may keep.123var ErrTooLarge = errors.New("the output is too large to keep")124125// Exec runs job and returns what it printed, NUON, at most maxBytes of126// it, within timeout.127func Exec(ctx context.Context, r Runner, job Job, timeout time.Duration, maxBytes int) ([]byte, error) {128 if timeout > 0 {129 var cancel context.CancelFunc130 ctx, cancel = context.WithTimeoutCause(ctx, timeout, fmt.Errorf("stopped after %s", timeout))131 defer cancel()132 }133 ctx, cancel := context.WithCancelCause(ctx)134 defer cancel(nil)135 out := &capped{max: maxBytes, stop: func() { cancel(ErrTooLarge) }}136 err := r.Run(ctx, job, Script(job, sortedNames(job.Tables)), out)137 switch {138 case ctx.Err() != nil:139 return nil, context.Cause(ctx)140 case err != nil:141 return nil, err142 }143 return bytes.TrimSpace(out.b.Bytes()), nil144}145146// capped keeps at most max bytes, stopping the run past them.147type capped struct {148 b bytes.Buffer149 max int150 stop func()151}152153func (c *capped) Write(p []byte) (int, error) {154 if c.max > 0 && c.b.Len()+len(p) > c.max {155 c.stop()156 return 0, ErrTooLarge157 }158 return c.b.Write(p)159}160161func sortedNames(tables map[string][]byte) []string {162 names := make([]string, 0, len(tables))163 for name := range tables {164 names = append(names, name)165 }166 slices.Sort(names)167 return names168}169170// Commands lists nu's command names, for completing them.171func Commands(ctx context.Context, r Runner, timeout time.Duration) ([]string, error) {172 data, err := Exec(ctx, r, Job{Command: "help commands | get name"}, timeout, 0)173 if err != nil {174 return nil, err175 }176 v, err := nuon.Parse(data)177 if err != nil {178 return nil, err179 }180 var out []string181 for _, n := range v.List {182 if n.Kind == nuon.String {183 out = append(out, n.Str)184 }185 }186 return out, nil187}188189// Nu is the real runner: `nu -c` in a process of its own.190type Nu struct {191 // Path is the program; "" finds nu on the PATH.192 Path string193}194195// Run implements Runner.196func (n Nu) Run(ctx context.Context, job Job, script string, stdout io.Writer) error {197 path := n.Path198 if path == "" {199 var err error200 if path, err = exec.LookPath("nu"); err != nil {201 return ErrMissing202 }203 }204 dir, err := os.MkdirTemp("", "012-nu-")205 if err != nil {206 return err207 }208 defer os.RemoveAll(dir)209 env := append(os.Environ(), "NO_COLOR=1")210 for i, name := range sortedNames(job.Tables) {211 file := dir + string(os.PathSeparator) + strconv.Itoa(i) + ".nuon"212 if err := os.WriteFile(file, job.Tables[name], 0o600); err != nil {213 return err214 }215 env = append(env, "NU012_TABLE_"+strconv.Itoa(i)+"="+file)216 }217 args, err := nuArgs(job, script, dir)218 if err != nil {219 return err220 }221 cmd := exec.CommandContext(ctx, path, args...)222 cmd.Env, cmd.Stdout = env, stdout223 cmd.WaitDelay = 2 * time.Second224 var stderr bytes.Buffer225 cmd.Stderr = &limited{w: &stderr, n: 64 << 10}226 if err := cmd.Run(); err != nil {227 if ctx.Err() != nil {228 return context.Cause(ctx)229 }230 text := ansi.Strip(stderr.String())231 return &Error{Msg: Message(text, err), Stderr: text}232 }233 return nil234}235236// nuArgs is nu's command line for job: the script to run, or the flags237// that ask about it, the script in a file in dir.238func nuArgs(job Job, script, dir string) ([]string, error) {239 if len(job.IDE) > 0 {240 file := dir + string(os.PathSeparator) + "cell.nu"241 if err := os.WriteFile(file, []byte(script), 0o600); err != nil {242 return nil, err243 }244 return append(append([]string{"--no-config-file", "--no-std-lib"}, job.IDE...), file), nil245 }246 args := []string{"-c", script}247 if !job.Config {248 args = append([]string{"--no-config-file"}, args...)249 }250 return args, nil251}252253// Message is the line of nu's standard error that says what went wrong:254// the one marked × (an ASCII x when nu doesn't draw in Unicode, as when255// its output isn't a terminal), or the first, or how the process ended.256func Message(stderr string, err error) string {257 first := ""258 for line := range strings.SplitSeq(stderr, "\n") {259 line = strings.TrimSpace(line)260 if rest, ok := strings.CutPrefix(line, "×"); ok {261 return strings.TrimSpace(rest)262 }263 if rest, ok := strings.CutPrefix(line, "x "); ok && strings.HasPrefix(first, "Error:") {264 return strings.TrimSpace(rest)265 }266 if first == "" && line != "" {267 first = line268 }269 }270 if first != "" {271 return first272 }273 return err.Error()274}275276// limited writes at most n bytes to w, dropping the rest.277type limited struct {278 w io.Writer279 n int280}281282func (l *limited) Write(p []byte) (int, error) {283 if l.n > 0 {284 k := min(len(p), l.n)285 l.n -= k286 if _, err := l.w.Write(p[:k]); err != nil {287 return 0, err288 }289 }290 return len(p), nil291}