diff --git a/README.md b/README.md index 91c047c395..a91089e7a3 100644 --- a/README.md +++ b/README.md @@ -122,6 +122,7 @@ d2 --watch in.d2 out.svg ``` A browser window will open with `out.svg` and live-reload on changes to `in.d2`. +If you want the generated output removed on exit, use `d2 --watch --watch-cleanup in.d2 out.svg`. ## Install diff --git a/d2cli/main.go b/d2cli/main.go index 6e49ad0615..c439e97ff3 100644 --- a/d2cli/main.go +++ b/d2cli/main.go @@ -75,6 +75,10 @@ func Run(ctx context.Context, ms *xmain.State) (err error) { if err != nil { return err } + watchCleanupFlag, err := ms.Opts.Bool("D2_WATCH_CLEANUP", "watch-cleanup", "", false, "when used with watch, remove generated output files when exiting") + if err != nil { + return err + } layoutFlag := ms.Opts.String("D2_LAYOUT", "layout", "l", "dagre", `the layout engine used`) themeFlag, err := ms.Opts.Int64("D2_THEME", "theme", "t", 0, "the diagram theme ID") if err != nil { @@ -369,6 +373,7 @@ func Run(ctx context.Context, ms *xmain.State) (err error) { port: *portFlag, inputPath: inputPath, outputPath: outputPath, + cleanupOutput: *watchCleanupFlag, bundle: *bundleFlag, forceAppendix: *forceAppendixFlag, pw: pw, diff --git a/d2cli/watch.go b/d2cli/watch.go index 3e3db95b99..4089fc8c69 100644 --- a/d2cli/watch.go +++ b/d2cli/watch.go @@ -51,6 +51,7 @@ type watcherOpts struct { port string inputPath string outputPath string + cleanupOutput bool boardPath string pwd string bundle bool @@ -88,6 +89,8 @@ type watcher struct { resMu sync.Mutex res *compileResult + + cleanupOutputDir bool } type compileResult struct { @@ -118,6 +121,13 @@ func newWatcher(ctx context.Context, ms *xmain.State, opts watcherOpts) (*watche } func (w *watcher) init() error { + if w.cleanupOutput && w.outputPath != "-" { + outputDir := watchOutputDir(w.outputPath) + if outputDir != "" { + _, err := os.Stat(outputDir) + w.cleanupOutputDir = errors.Is(err, os.ErrNotExist) + } + } fw, err := fsnotify.NewWatcher() if err != nil { return err @@ -151,20 +161,35 @@ func (w *watcher) initStaticFileServer() error { return nil } -func (w *watcher) run() error { - defer w.close() +func (w *watcher) run() (err error) { + defer func() { + w.close() + if !w.cleanupOutput || w.outputPath == "-" { + return + } + cleanupErr := cleanupWatchOutput(w.outputPath, w.cleanupOutputDir) + if cleanupErr == nil { + return + } + if err == nil { + err = cleanupErr + return + } + w.ms.Log.Warn.Printf("failed to clean up watch output %s: %v", w.ms.HumanPath(w.outputPath), cleanupErr) + }() w.goFunc(w.watchLoop) w.goFunc(w.compileLoop) - err := w.goServe() + err = w.goServe() if err != nil { return err } w.wg.Wait() w.close() - return w.err + err = w.err + return err } func (w *watcher) close() { @@ -485,9 +510,35 @@ func (w *watcher) listen() error { } w.l = l w.ms.Log.Success.Printf("listening on http://%v", w.l.Addr()) + if w.cleanupOutput { + w.ms.Log.Info.Printf("press Ctrl-C to quit and remove %s", w.ms.HumanPath(w.outputPath)) + } else { + w.ms.Log.Info.Printf("press Ctrl-C to quit") + } return nil } +func cleanupWatchOutput(outputPath string, removeDir bool) error { + err := os.Remove(outputPath) + if err != nil && !errors.Is(err, os.ErrNotExist) { + return err + } + + outputDir := watchOutputDir(outputPath) + if !removeDir || outputDir == "" { + return nil + } + return os.RemoveAll(outputDir) +} + +func watchOutputDir(outputPath string) string { + ext := filepath.Ext(outputPath) + if ext == "" { + return "" + } + return strings.TrimSuffix(outputPath, ext) +} + func (w *watcher) goServe() error { m := http.NewServeMux() // TODO: Add cmdlog logging and error reporting middleware diff --git a/e2etests-cli/main_test.go b/e2etests-cli/main_test.go index 8d2913f6e9..3385ed567b 100644 --- a/e2etests-cli/main_test.go +++ b/e2etests-cli/main_test.go @@ -1342,6 +1342,52 @@ layers: { assert.Success(t, err) }, }, + { + name: "watch-cleanup", + serial: true, + run: func(t *testing.T, ctx context.Context, dir string, env *xos.Env) { + writeFile(t, dir, "index.d2", `a -> b`) + + stderr := &stderrWrapper{} + tms := testMain(dir, env, "--watch", "--watch-cleanup", "--browser=0", "index.d2", "out.svg") + tms.Stderr = stderr + + tms.Start(t, ctx) + interrupted := false + defer func() { + if interrupted { + return + } + err := tms.Signal(ctx, os.Interrupt) + assert.Success(t, err) + }() + + doneRE := regexp.MustCompile(`successfully compiled index.d2`) + _, err := waitLogs(ctx, stderr, doneRE) + assert.Success(t, err) + + outputPath := filepath.Join(dir, "out.svg") + _, err = os.Stat(outputPath) + assert.Success(t, err) + + err = tms.Signal(ctx, os.Interrupt) + assert.Success(t, err) + interrupted = true + + deadline := time.Now().Add(2 * time.Second) + for { + _, err = os.Stat(outputPath) + if errors.Is(err, os.ErrNotExist) { + break + } + assert.Success(t, err) + if time.Now().After(deadline) { + t.Fatalf("expected %s to be removed after watch exit", outputPath) + } + time.Sleep(10 * time.Millisecond) + } + }, + }, { name: "watch-ok-link", serial: true,