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