improvements, bugfixes

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:])