From dcda55a0e966087f629afcee48c68494599d3a61 Mon Sep 17 00:00:00 2001
From: Jonathan Strong <jonathan.strong@gmail.com>
Date: Tue, 10 Mar 2020 23:22:17 -0400
Subject: [PATCH] improvements, bugfixes

---
 get-json-key.py | 30 ++++++++++++++++++++----------
 1 file changed, 20 insertions(+), 10 deletions(-)

diff --git a/get-json-key.py b/get-json-key.py
index 7cc0b4d..5038157 100755
--- a/get-json-key.py
+++ b/get-json-key.py
@@ -1,29 +1,39 @@
 #!/usr/bin/env python3
 
+HELP = """
+USAGE:
+    [stdin] | ./get-json-key.py [KEY]..
+
+KEY can be a string key (object) or integer index (array), the program will try both
+"""
+
+
 import fileinput, sys, json, functools
-from typing import Union, List
+from typing import *
 
-JsonKey = Union[str, int]
 
-def try_int(x: str) -> JsonKey:
+def try_as_obj_then_as_array(data: Any, key: str) -> Any:
     try:
-        return int(x)
-    except:
-        return x
+        # initial attempt, as object
+        return data[key]
+    except (KeyError, TypeError):
+        # fallback, as array
+        return data[int(key)]
+
 
-def main(indices: List[str]) -> None:
+def main(keys: List[str]) -> None:
     for line in fileinput.input(files=('-',)):
         data = json.loads(line)
         try:
-            val = functools.reduce(lambda x, i: x[try_int(i)], indices, data)
+            val = functools.reduce(lambda acc, x: try_as_obj_then_as_array(acc, x), keys, data)
         except Exception as e:
-            print('indexing json data failed with keys {}: {}\n'.format(' > '.join(indices), e), file=sys.stderr)
+            print('indexing json data failed with keys {}: {}\n'.format(' > '.join(keys), e), file=sys.stderr)
             raise e
         print(val)
 
 if __name__ == '__main__':
     if len(sys.argv) < 2:
-        print('USAGE: ./get-json-key.py [KEY].. < json-input', file=sys.stderr)
+        print(HELP, file=sys.stderr)
         sys.exit(1)
     main(sys.argv[1:])