1#!/usr/local/bin/python3.8
2
3########################################################################
4## 08/2018 Justin Rajendra
5## convert : sep txt to json or the other way around
6## add stuff to json files
7##
8## Nov 23, 2018: PA Taylor
9## + expanding functionality for -txt2json stuff:
10##   - allow for strings inside double quotes to be treated as one value
11##     (will probably make string wrapping to be either " or ')
12##   - also make two new opts: '-delimiter_major ..' and
13##     '-delimiter_minor ..'  to allow for key-value and value-value
14##     separation to occur at  different chars.
15##
16########################################################################
17
18## system libraries
19import sys, os, glob, subprocess, csv, re, argparse, signal, textwrap, json
20from afnipy import abids_lib
21from collections import OrderedDict
22
23# ---------------------------------------------------------------------
24
25# This function only applies to parsing the simple, colon-separated
26# text files entered with -txt2json
27def parse_txt_value(x, delmin):
28
29    # examples of chars that we look for to open+close strings
30    str_symb = [ "'", '"' ]
31
32    y = x.rstrip().lstrip()
33    N = len(y)
34
35    olist = []
36    KEEP_GOING = 1
37    Nleft = N
38    i = 0
39    while Nleft and i<N and KEEP_GOING :
40        mysymb = ''
41
42        # see if we find the start of an enclosed str
43        for ss in str_symb:
44            if y[i] == ss :
45                i0 = i
46                mysymb = ss
47
48        # if it starts, try to find its close
49        if mysymb:
50            # search for partner, starting from next ele
51            j0 = i0+1
52            i1 = y[j0:].find(mysymb)
53            if i1 >= 0 :
54                # if partner found, save that piece of string
55                j1 = j0+i1
56                olist.append(y[j0:j1])          # inside quotes
57                # then jump to partner loc plus one, and continue
58                i = j1+1
59                Nleft = len(y[i:])
60            else:
61                # if partner NOT found, that will be it for the loop
62                KEEP_GOING = 0
63        else:
64            i+=1
65
66    # and attach the remainder as a final list, split by spaces (def)
67    # or by user-specified delimiter
68    if Nleft > 0 :
69        z = y.rstrip().lstrip().split(delmin)
70        for ele in z:
71            olist.append(ele.rstrip().lstrip())
72
73    return olist
74
75########################################################################
76## parse command line arguments / build help
77
78## make parser with help
79parser = argparse.ArgumentParser(prog=str(sys.argv[0]),add_help=False,
80                                 formatter_class=argparse.RawDescriptionHelpFormatter,
81                                 description=textwrap.dedent('''\
82------------------------------------------
83Overview ~1~
84
85    This script helps to manipulate json files in various ways.
86
87Caveats ~1~
88
89    None yet.
90
91Example ~1~
92
93    abids_json_tool.py -input out.ss_review.FT.txt \
94                       -prefix out.ss_review.FT.json \
95                       -txt2json
96
97------------------------------------------
98
99Options ~1~
100
101                                 '''),epilog=textwrap.dedent('''\
102------------------------------------------
103Justin Rajendra circa 08/2018
104Keep on keeping on!
105------------------------------------------
106                                 '''))
107
108## setup the groups
109OneRequired = parser.add_argument_group('Only one of these')
110OnlyOne = OneRequired.add_mutually_exclusive_group(required=True)
111required = parser.add_argument_group('Required arguments')
112parser._optionals.title = 'Optional arguments'
113parser._action_groups.reverse()
114
115## required
116required.add_argument('-input',type=str,metavar='FILE',required=True,
117                      help=('One file to convert. '+
118                            '(either ":" separated or json formatted.) '+
119                            'Enter NULL with -add_json to create new json file.'))
120required.add_argument('-prefix',type=str,metavar='PREFIX',required=True,
121                      help='Output file name.')
122
123## only one of these at a time
124OnlyOne.add_argument('-txt2json',action="store_true",default=False,
125                     help=('Convert from ":" separated text file to '+
126                           'json formatted file.'))
127OnlyOne.add_argument('-json2txt',action="store_true",default=False,
128                     help=('Convert from json formatted file to '+
129                           '":" separated text file.'))
130OnlyOne.add_argument('-add_json',type=str,nargs='+',metavar=('KEY','VALUE'),
131                     action="append",
132                     help=('Add an attribute to the end of the specified '+
133                            'json file. Needs exactly two arguments. '+
134                            '(e.g. Fruit Apple) '+
135                            'The KEY must not have spaces and must be only '+
136                            'one word. If the VALUE is more than one item, it '+
137                            'needs to be surrounded by single or double quotes '+
138                            'and be comma separated (e.g. Fruit "Apple,Orange")'))
139OnlyOne.add_argument('-del_json',type=str,nargs=1,metavar='KEY',
140                     help=('Remove attribute (KEY) from the -input json file.'))
141## optional
142parser.add_argument('-force_add','-f',action="store_true",default=False,
143                    help=('Use with -add_json to overwrite an existing '+
144                          'attribute in the specified json file.'))
145parser.add_argument('-overwrite',action="store_true",default=False,
146                    help=('Use caution as this will overwrite the -prefix '+
147                          'file if it exists!!'))
148parser.add_argument('-help',action='help',help='Show this help and exit.')
149# [PT: Nov 21, 2018] new opts to adjust delims under -txt2json functionality
150parser.add_argument('-delimiter_major',type=str,metavar='DELIM_MAJ',
151                    default=': | =',
152                     help=('When using "-txt2json" opt, specify the '+
153                           'new (major) delimiter to separate keys and values.'))
154parser.add_argument('-delimiter_minor',type=str,metavar='DELIM_MIN',
155                    default=None,
156                     help=('When using "-txt2json" opt, specify the '+
157                           'new (minor) delimiter to separate value items. '+
158                           'NB: pairs of quotes take priority to define '+
159                           'a single item. The default delimiter '+
160                           '(outside of quotes) is whitespace.'))
161
162## if nothing, show help
163if len(sys.argv) == 1:
164    parser.print_help()
165    sys.exit(1)
166
167########################################################################
168## collect the arguments
169args = parser.parse_args()
170input = args.input
171txt2json = args.txt2json
172json2txt = args.json2txt
173new_entry = args.add_json
174del_entry = args.del_json
175prefix = args.prefix
176overwrite = args.overwrite
177force = args.force_add
178# [PT: Nov. 21, 2018] For '-txt2json': options on what separates what
179# keys and values (DELIM_MAJ) and different values (DELIM_MIN).
180DELIM_MAJ = args.delimiter_major
181DELIM_MIN = args.delimiter_minor
182
183########################################################################
184## check stuff
185
186## check input file
187if input == "NULL":
188    input = prefix
189else:
190    if not os.path.isfile(input):
191        print("\nERROR: "+input+" not found!!\n")
192        sys.exit(1)
193
194## namey things
195infile = os.path.abspath(input)
196full_path = os.path.dirname(infile)
197infile_base = os.path.basename(infile)
198infile_ext = os.path.splitext(infile_base)[1]
199
200## check prefix
201if os.path.isfile(prefix) and not overwrite:
202    print("\nERROR: "+prefix+" exists!!\n")
203    sys.exit(1)
204
205########################################################################
206## txt2json
207if txt2json:
208    json_dict = OrderedDict()   ## preserve order
209    with open(infile) as f:
210        for line in f:
211
212            ## split on : and skip if blank line
213            field = re.split(DELIM_MAJ, line)
214            # field = line.split(":")
215            if len(field) == 1: continue
216
217            ## make dictionary key and clean up
218            key = field[0].rstrip().replace(" ","_")    ## get rid of spaces
219            key = re.sub("[()]","",key)                 ## get rid of ()
220
221            ## make value list or entry and convert to float if number
222            # [PT: Nov 21, 2018] allow strings to be a single value
223            value_list = parse_txt_value( field[1], delmin=DELIM_MIN )
224            ## older form:
225            #value_list = field[1].rstrip().lstrip().split()
226
227            value = []
228            for v in value_list:
229                try:
230                    value.append(float(v))
231                except ValueError:
232                    value.append(str(v))
233
234            ## if only one, make not a list and add to dictionary
235            if len(value) == 1: value = value[0]
236            json_dict.update({key:value})
237
238    ######################
239    ## write out json file
240    json_out = json.dumps(json_dict,indent=4)
241    f = open(prefix,"w")
242    f.write(json_out)
243    f.close()
244## end txt2json
245
246########################################################################
247## json2txt
248if json2txt:
249    ## read in json from handy abids_lib function
250    json_data = abids_lib.json_import(infile)
251
252    ## get the max characters for lining everything up
253    max_char = max([len(i) for i in json_data.keys()])
254
255    ## write out
256    with open(prefix,'wb') as csv_file:
257        writer = csv.writer(csv_file,delimiter=":")
258        for key, value in json_data.items():
259            trailing = max_char - len(key) + 2  ## spacing
260
261            ## check if list and print space separated
262            if isinstance(value, (list,)):
263                writer.writerow([key+' '*trailing,'  '+' '.join(map(str,value))])
264            else:
265                writer.writerow([key+' ' * trailing,'  '+str(value)])
266## end json2txt
267
268########################################################################
269## add entry
270if new_entry is not None:
271
272    if input is not prefix:
273        ## read in json from handy abids_lib function
274        json_data = abids_lib.json_import(infile)
275    else:
276        json_data = OrderedDict()   ## make empty one
277
278    ## get the new stuff
279    for new in new_entry:
280
281        key = new[0]
282        value_list = new[1].split(',')
283
284        ## check to see if the attribute is already there
285        if key in json_data.keys() and not force:
286            print("\nERROR: The '"+key+"' attribute is already exists in "+
287                  infile_base+"!!\n"+
288                  "       Use the -force (-f) option to overwrite attribute.\n")
289            sys.exit(1)
290
291        ## make value list or entry and convert to float if number
292        value = []
293        for v in value_list:
294            try:
295                value.append(float(v))
296            except ValueError:
297                value.append(str(v))
298
299        ## not a list if only 1 and add new entry
300        if len(value) == 1: value = value[0]
301        json_data.update({key:value})
302
303    ######################
304    ## write out json file
305    json_out = json.dumps(json_data,indent=4)
306    f = open(prefix,"w")
307    f.write(json_out)
308    f.close()
309## end add entry
310
311########################################################################
312## delete entry
313if del_entry is not None:
314    ## read in json from handy abids_lib function
315    json_data = abids_lib.json_import(infile)
316
317    ## make sure it is there and remove it
318    if del_entry[0] in json_data:
319
320        print("\nRemoving '"+del_entry[0]+"' from "+infile_base+"\n")
321        del json_data[del_entry[0]]
322
323        ## write out json file
324        json_out = json.dumps(json_data,indent=4)
325        f = open(prefix,"w")
326        f.write(json_out)
327        f.close()
328    else:
329        print("\nERROR: The '"+del_entry[0]+"' attribute does not exist in "+
330              infile_base+"!!\n"+
331              "       View the available attributes with:\n"+
332              "       abids_json_info.py -json "+infile+" -list_fields\n")
333        sys.exit(1)
334## end delete entry
335
336sys.exit(0)
337
338