From 9b185a4409872afc6bfc7ff8505420054d1f9617 Mon Sep 17 00:00:00 2001
From: Mike Farah <mikefarah@gmail.com>
Date: Tue, 5 Jan 2021 13:28:37 +1100
Subject: [PATCH] Added shorthand document index selection

---
 pkg/yqlib/doc/Document Index.md           | 34 +++++++++++++++++++++++
 pkg/yqlib/operator_document_index_test.go | 17 ++++++++++++
 pkg/yqlib/path_tokeniser.go               |  1 +
 3 files changed, 52 insertions(+)

diff --git a/pkg/yqlib/doc/Document Index.md b/pkg/yqlib/doc/Document Index.md
index 8268b77c..024e4916 100644
--- a/pkg/yqlib/doc/Document Index.md	
+++ b/pkg/yqlib/doc/Document Index.md	
@@ -17,6 +17,24 @@ will output
 1
 ```
 
+## Retrieve a document index, shorthand
+Given a sample.yml file of:
+```yaml
+a: cat
+---
+a: frog
+```
+then
+```bash
+yq eval '.a | di' sample.yml
+```
+will output
+```yaml
+0
+---
+1
+```
+
 ## Filter by document index
 Given a sample.yml file of:
 ```yaml
@@ -33,6 +51,22 @@ will output
 a: frog
 ```
 
+## Filter by document index shorthand
+Given a sample.yml file of:
+```yaml
+a: cat
+---
+a: frog
+```
+then
+```bash
+yq eval 'select(di == 1)' sample.yml
+```
+will output
+```yaml
+a: frog
+```
+
 ## Print Document Index with matches
 Given a sample.yml file of:
 ```yaml
diff --git a/pkg/yqlib/operator_document_index_test.go b/pkg/yqlib/operator_document_index_test.go
index 81898ed2..11e4c772 100644
--- a/pkg/yqlib/operator_document_index_test.go
+++ b/pkg/yqlib/operator_document_index_test.go
@@ -14,6 +14,15 @@ var documentIndexScenarios = []expressionScenario{
 			"D1, P[a], (!!int)::1\n",
 		},
 	},
+	{
+		description: "Retrieve a document index, shorthand",
+		document:    "a: cat\n---\na: frog\n",
+		expression:  `.a | di`,
+		expected: []string{
+			"D0, P[a], (!!int)::0\n",
+			"D1, P[a], (!!int)::1\n",
+		},
+	},
 	{
 		description: "Filter by document index",
 		document:    "a: cat\n---\na: frog\n",
@@ -22,6 +31,14 @@ var documentIndexScenarios = []expressionScenario{
 			"D1, P[], (doc)::a: frog\n",
 		},
 	},
+	{
+		description: "Filter by document index shorthand",
+		document:    "a: cat\n---\na: frog\n",
+		expression:  `select(di == 1)`,
+		expected: []string{
+			"D1, P[], (doc)::a: frog\n",
+		},
+	},
 	{
 		description: "Print Document Index with matches",
 		document:    "a: cat\n---\na: frog\n",
diff --git a/pkg/yqlib/path_tokeniser.go b/pkg/yqlib/path_tokeniser.go
index 1cf3d93c..b01c74c0 100644
--- a/pkg/yqlib/path_tokeniser.go
+++ b/pkg/yqlib/path_tokeniser.go
@@ -192,6 +192,7 @@ func initLexer() (*lex.Lexer, error) {
 	lexer.Add([]byte(`\/\/`), opToken(Alternative))
 
 	lexer.Add([]byte(`documentIndex`), opToken(GetDocumentIndex))
+	lexer.Add([]byte(`di`), opToken(GetDocumentIndex))
 
 	lexer.Add([]byte(`style`), opAssignableToken(GetStyle, AssignStyle))