read: retrieve specific key from specific subdocument

utils_test.go

@@ -66,6 +66,12 @@ func assertAnyErr(t *testing.T, actualValue error) {
 	}
 }
 
+func assertNilErr(t *testing.T, actualValue error) {
+	if actualValue != nil {
+		t.Error("Expected nil error, got ", actualValue)
+	}
+}
+
 func writeTempYamlFile(content string) string {
 	tmpfile, _ := ioutil.TempFile("", "testyaml")
 	defer func() {

yq.go

@@ -5,6 +5,7 @@ import (
 	"fmt"
 	"io/ioutil"
 	"os"
+	"regexp"
 	"strconv"
 	"strings"
 
@@ -180,16 +181,61 @@ func readProperty(cmd *cobra.Command, args []string) error {
 	return nil
 }
 
+var regexpSubdocPath = regexp.MustCompile(`^@(\d+)\.(.*)$`)
+
 func read(args []string) (interface{}, error) {
-	var parsedData yaml.MapSlice
-	var path = ""
+	var (
+		parsedData    yaml.MapSlice
+		path          = ""
+		multidoc      bool
+		docNumber     int
+		matches       []string
+		err           error
+		childDocument string
+	)
 
 	if len(args) < 1 {
 		return nil, errors.New("Must provide filename")
 	} else if len(args) > 1 {
 		path = args[1]
+		multidoc = regexpSubdocPath.MatchString(path)
 	}
 
+	if multidoc {
+		matches = regexpSubdocPath.FindStringSubmatch(path)
+		// matches[0] is the original path
+		// matches[1] is the document number
+		// matches[2] is the path within that document
+		if len(matches) == 3 {
+			docNumber, err = strconv.Atoi(matches[1])
+			if err == nil {
+				path = matches[2]
+			}
+		}
+		childDocument, err = readSubDocument(args[0], docNumber)
+		if err != nil {
+			return nil, err
+		}
+		tmp, err := ioutil.TempFile("", "")
+		defer os.Remove(tmp.Name())
+		if err != nil {
+			return nil, err
+		}
+		n, err := tmp.WriteString(childDocument)
+		if err != nil {
+			return nil, err
+		}
+		if n != len(childDocument) {
+			return nil, fmt.Errorf("multidoc tmp file expected to write %v bytes but wrote %v bytes", len(childDocument), n)
+		}
+
+		args[0] = tmp.Name()
+
+	}
+
+	// If we're asking for any document beyond the first,
+	// we need to provide the yaml package with that document
+
 	if err := readData(args[0], &parsedData); err != nil {
 		var generalData interface{}
 		if err = readData(args[0], &generalData); err != nil {
@@ -460,3 +506,28 @@ func readData(filename string, parsedData interface{}) error {
 
 	return yaml.Unmarshal(rawData, parsedData)
 }
+
+var regexpDocumentSeparator = regexp.MustCompile(`(?m)^---$`)
+
+// readSubDocument takes a filename and an index.
+// It looks for ^---$ markers that separate yaml documents and splits the file on those.
+// It takes the zero-indexed nth document, writes it to a temporary file, and return a handle to that file.
+// Is trailing whitespace legal yaml?
+func readSubDocument(filename string, n int) (string, error) {
+	allDocs, err := ioutil.ReadFile(filename)
+	if err != nil {
+		return "", err
+	}
+
+	subDocs := regexpDocumentSeparator.Split(string(allDocs), -1)
+	if len(subDocs) == 0 {
+		return "", errors.New("document was split into zero subdocuments")
+	}
+	if subDocs[0] == "" {
+		subDocs = subDocs[1:]
+	}
+	if len(subDocs) <= n {
+		return "", fmt.Errorf("requested document %v out of total %v documents", n, len(subDocs))
+	}
+	return strings.TrimSpace(subDocs[n]), nil
+}

yq_test.go

@@ -139,3 +139,30 @@ func TestNewYaml_WithUnknownScript(t *testing.T) {
 	expectedOutput := `open fake-unknown: no such file or directory`
 	assertResult(t, expectedOutput, err.Error())
 }
+
+func TestSplitData(t *testing.T) {
+	t.Run("document index zero", func(t *testing.T) {
+		actual, err := readSubDocument("examples/multidocument.yaml", 0)
+		assertNilErr(t, err)
+		expected := `a: Document One
+b:
+  c: 1`
+		assertResult(t, expected, actual)
+		// assert.Equal(t, expected, actual) // needed this to find whitespace problem
+	})
+
+	t.Run("document index one", func(t *testing.T) {
+		actual, err := readSubDocument("examples/multidocument.yaml", 1)
+		assertNilErr(t, err)
+		expected := `a: Document Two
+b:
+  c: 2`
+		assertResult(t, expected, actual)
+	})
+
+	t.Run("document index beyond available", func(t *testing.T) {
+		actual, err := readSubDocument("examples/multidocument.yaml", 2)
+		assertAnyErr(t, err)
+		assertResult(t, "", actual)
+	})
+}