snippets/get-json-key.py

#!/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 *


def try_as_obj_then_as_array(data: Any, key: str) -> Any:
    try:
        # initial attempt, as object
        return data[key]
    except (KeyError, TypeError):
        # fallback, as array
        return data[int(key)]


def main(keys: List[str]) -> None:
    for line in fileinput.input(files=('-',)):
        data = json.loads(line)
        try:
            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(keys), e), file=sys.stderr)
            raise e
        print(val)

if __name__ == '__main__':
    if len(sys.argv) < 2:
        print(HELP, file=sys.stderr)
        sys.exit(1)
    main(sys.argv[1:])