<reponame>vdjserver/agave-cli
#!/usr/bin/python

'''Python >=2.7 command line utility for querying JSON using dot notation'''

from __future__ import absolute_import
from __future__ import print_function

from easydict import EasyDict as edict
import argparse
import json
import re
import six
import sys

from operator import attrgetter
from six.moves import map

debug = False


def parse_args():
    parser = argparse.ArgumentParser(prog="pydotjson.py",
                                     description="Command line utility to query JSON with dot.notation")

    parser.add_argument("-d", "--debug", action='store_true',
                        default=False,
                        help="Show debug information")

    parser.add_argument("-s", "--stripquotes", action='store_true',
                        default=False,
                        help="Do not quote the query response")
    parser.add_argument("-q", "--query", dest="query",
                        type=str, nargs='?', default=None,
                        help="Query in JSON dot notation")
    parser.add_argument('input', nargs='?', type=argparse.FileType('r'),
                        default=sys.stdin,
                        help="The JSON file to import or stdin if provided.")

    args = parser.parse_args()

    return args


def main():

    global debug

    args = parse_args()
    debug = args.debug

    if args.query is None:

        print(json.dumps(args.input,
              sort_keys=True, indent=2, separators=(',', ': ')))

    else:

        raw = args.input.read()

        if args.query.startswith('[]') == 1:
            # Modify the JSON to be compatible with EasyDict
            raw = '{"things":' + raw + '}'

        if debug:
            print('Original JSON: ' + raw)

        raw = json.loads(raw)

        d = edict(raw)

        if d is None:

            if not args.query:
                print('')
            else:
                print('Empty object has no attribute {0}'.format(args.query))

        else:

            # skipping interpolation of the path in favor of a raw eval here.
            # not as safe, but this is an internal process
            if args.query.startswith('[]') == 1:
                json_path = 'd.things.' + args.query
            else:
                json_path = 'd.' + args.query

            json_path = json_path.strip('[]').strip('.')
            json_path = re.sub(r'\.\[(\d+)\]', '[\1]', json_path,
                               flags=re.IGNORECASE)
            json_path = re.sub(r'\.\[(.+)\]', '[\'\1\']', json_path,
                               flags=re.IGNORECASE)
            json_path = re.sub(r'\[\]\.', '.[].', json_path,
                               flags=re.IGNORECASE)
            json_path = re.sub(r'\.\.', '.', json_path,
                               flags=re.IGNORECASE)

            if json_path.find('[]') != -1:
                tokens = json_path.split('.[].', 1)
                if debug:
                    print(
                        'Path query: map(attrgetter({0}), eval({1}))'.format(
                            tokens[0], tokens[1]))
                query_result = list(map(attrgetter(
                    tokens[1]), eval(tokens[0])))
            else:
                if debug:
                        print('Path query is: ' + json_path)
                query_result = eval(json_path)

            # query produced a primary type
            if isinstance(query_result, six.string_types):

                if args.stripquotes is False:
                    print('"{0}"'.format(query_result))
                else:
                    print(query_result)

            elif isinstance(query_result, int) or \
                    isinstance(query_result, float):

                print(query_result)

            elif isinstance(query_result, bool):

                if query_result:
                    print('true')
                else:
                    print('false')

            elif isinstance(query_result, list):

                if args.stripquotes:
                    for r in query_result:
                        print(r)
                else:
                    print(json.dumps(query_result, sort_keys=True,
                                     indent=2, separators=(',', ': ')))

            else:
                print(json.dumps(query_result, sort_keys=True,
                                 indent=2, separators=(',', ': ')))


if __name__ == "__main__":
    try:
        main()
    except Exception as e:
        if debug:
            raise
        else:
            print("Error: ".format(e))
