From 9d882451027b3e847cd6a6878d300e146387c6bb Mon Sep 17 00:00:00 2001
From: Oliver Marriott <rktjmp@users.noreply.github.com>
Date: Sun, 6 Nov 2022 04:57:18 +1100
Subject: [PATCH] Fix `--working-dir=` and `--notebook-dir=` options parsing
 (#267)

---
 main.go                      | 67 +++++++++++++++++++++++++++++-------
 tests/flag-notebook-dir.tesh |  2 ++
 tests/flag-working-dir.tesh  | 10 ++++++
 3 files changed, 66 insertions(+), 13 deletions(-)

diff --git a/main.go b/main.go
index f58af44..d8185ef 100644
--- a/main.go
+++ b/main.go
@@ -7,6 +7,7 @@ import (
 	"os/exec"
 	"os/signal"
 	"path/filepath"
+	"regexp"
 	"runtime"
 	"strings"
 
@@ -204,9 +205,9 @@ func runAlias(container *cli.Container, args []string) (bool, error) {
 // which path arguments are relative from.
 //
 // By order of precedence:
-//   1. --notebook-dir flag
-//   2. current working directory
-//   3. ZK_NOTEBOOK_DIR environment variable
+//  1. --notebook-dir flag
+//  2. current working directory
+//  3. ZK_NOTEBOOK_DIR environment variable
 func notebookSearchDirs(dirs cli.Dirs) ([]cli.Dirs, error) {
 	wd, err := os.Getwd()
 	if err != nil {
@@ -255,23 +256,63 @@ func parseDirs(args []string) (cli.Dirs, []string, error) {
 	var d cli.Dirs
 	var err error
 
+	// Split str by first "=" if present and return the split pair, otherwise return nil
+	makeSplitPair := func(str string) (pair []string) {
+		re := regexp.MustCompile(`=`)
+		slice := re.FindStringIndex(str)
+		if slice == nil {
+			return nil
+		}
+		return []string{str[:slice[0]], str[slice[1]:]}
+	}
+
+	// Peek ahead at next value  and pair with current if it exists, otherwise return nil
+	makePeekPair := func(args []string, index int) (pair []string) {
+		if len(args) <= (index + 1) {
+			return nil
+		}
+		return []string{args[index], args[index+1]}
+	}
+
+	matchesLongOrShort := func(str string, long string, short string) bool {
+		return str == long || (short != "" && str == short)
+	}
+
 	findFlag := func(long string, short string, args []string) (string, []string, error) {
 		newArgs := []string{}
-
-		foundFlag := ""
 		for i, arg := range args {
-			if arg == long || (short != "" && arg == short) {
-				foundFlag = arg
-			} else if foundFlag != "" {
+			// We can be given "--notebook-dir x" (two args) or "--notebook-dir=x" (one arg)
+			// so we must test against the current argument split into two, and
+			// the current argument + the next.
+			splitPair := makeSplitPair(arg)
+			peekPair := makePeekPair(args, i)
+			var option string
+			var value string
+
+			if splitPair != nil && matchesLongOrShort(splitPair[0], long, short) {
+				option = splitPair[0]
+				value = splitPair[1]
+				// skip 1 ahead
 				newArgs = append(newArgs, args[i+1:]...)
-				path, err := filepath.Abs(arg)
-				return path, newArgs, err
+			} else if peekPair != nil && matchesLongOrShort(peekPair[0], long, short) {
+				option = peekPair[0]
+				value = peekPair[1]
+				// skip 2 ahead (arg and value)
+				newArgs = append(newArgs, args[i+2:]...)
 			} else {
+				// we either had no split pair or peek pair, or they didn't match the
+				// needle, so just save the given arg and keep looking.
 				newArgs = append(newArgs, arg)
 			}
-		}
-		if foundFlag != "" {
-			return "", newArgs, errors.New(foundFlag + " requires a path argument")
+
+			if option != "" && value != "" {
+				path, err := filepath.Abs(value)
+				return path, newArgs, err
+			} else if option != "" && value == "" {
+				return "", newArgs, errors.New(option + " requires a path argument")
+			} else if len(args) == (i+1) && matchesLongOrShort(arg, long, short) {
+				return "", newArgs, errors.New(arg + " requires a path argument")
+			}
 		}
 		return "", newArgs, nil
 	}
diff --git a/tests/flag-notebook-dir.tesh b/tests/flag-notebook-dir.tesh
index 671b1c0..c11e9e5 100644
--- a/tests/flag-notebook-dir.tesh
+++ b/tests/flag-notebook-dir.tesh
@@ -3,7 +3,9 @@
 2>zk: error: failed to open notebook: no notebook found in {{working-dir}} or a parent directory
 
 # Provide the notebook directory with `--notebook-dir`.
+# --notebook-dir is custom parsed, check that it handles both one and two arg forms
 $ zk index -q --notebook-dir paths
+$ zk index -q --notebook-dir=paths
 
 # Provide the notebook directory with the `ZK_NOTEBOOK_DIR` env variable.
 $ ZK_NOTEBOOK_DIR={{working-dir}}/paths zk index -q
diff --git a/tests/flag-working-dir.tesh b/tests/flag-working-dir.tesh
index ba290cb..baa861b 100644
--- a/tests/flag-working-dir.tesh
+++ b/tests/flag-working-dir.tesh
@@ -3,11 +3,21 @@ $ zk list -qfpath --working-dir paths
 >paper.md
 >brown/wood.md
 
+# working-dir is custom parsed, check that it handles both one and two arg forms
+$ zk list -qfpath --working-dir=paths
+>paper.md
+>brown/wood.md
+
 # Short flag.
 $ zk list -qfpath -W paths
 >paper.md
 >brown/wood.md
 
+# working-dir is custom parsed, check that it handles both one and two arg forms
+$ zk list -qfpath -W=paths
+>paper.md
+>brown/wood.md
+
 # The note paths and links are relative to the provided working directory
 $ zk list -qf"\{{path}} \{{link}}" --working-dir paths/brown
 >../paper.md [Paper](../paper)