| 116 | |
| 117 | |
| 118 | def write_json(input_gen, file_name, output, qty, skip, id_tag, # <3> |
| 119 | gen_uuid, mongo, mfn, isis_json_type, prefix, |
| 120 | constant): |
| 121 | start = skip |
| 122 | end = start + qty |
| 123 | if id_tag: |
| 124 | id_tag = str(id_tag) |
| 125 | ids = set() |
| 126 | else: |
| 127 | id_tag = '' |
| 128 | for i, record in enumerate(input_gen): |
| 129 | if i >= end: |
| 130 | break |
| 131 | if not mongo: |
| 132 | if i == 0: |
| 133 | output.write('[') |
| 134 | elif i > start: |
| 135 | output.write(',') |
| 136 | if start <= i < end: |
| 137 | if id_tag: |
| 138 | occurrences = record.get(id_tag, None) |
| 139 | if occurrences is None: |
| 140 | msg = 'id tag #%s not found in record %s' |
| 141 | if ISIS_MFN_KEY in record: |
| 142 | msg = msg + (' (mfn=%s)' % record[ISIS_MFN_KEY]) |
| 143 | raise KeyError(msg % (id_tag, i)) |
| 144 | if len(occurrences) > 1: |
| 145 | msg = 'multiple id tags #%s found in record %s' |
| 146 | if ISIS_MFN_KEY in record: |
| 147 | msg = msg + (' (mfn=%s)' % record[ISIS_MFN_KEY]) |
| 148 | raise TypeError(msg % (id_tag, i)) |
| 149 | else: # ok, we have one and only one id field |
| 150 | if isis_json_type == 1: |
| 151 | id = occurrences[0] |
| 152 | elif isis_json_type == 2: |
| 153 | id = occurrences[0][0][1] |
| 154 | elif isis_json_type == 3: |
| 155 | id = occurrences[0]['_'] |
| 156 | if id in ids: |
| 157 | msg = 'duplicate id %s in tag #%s, record %s' |
| 158 | if ISIS_MFN_KEY in record: |
| 159 | msg = msg + (' (mfn=%s)' % record[ISIS_MFN_KEY]) |
| 160 | raise TypeError(msg % (id, id_tag, i)) |
| 161 | record['_id'] = id |
| 162 | ids.add(id) |
| 163 | elif gen_uuid: |
| 164 | record['_id'] = unicode(uuid4()) |
| 165 | elif mfn: |
| 166 | record['_id'] = record[ISIS_MFN_KEY] |
| 167 | if prefix: |
| 168 | # iterate over a fixed sequence of tags |
| 169 | for tag in tuple(record): |
| 170 | if str(tag).isdigit(): |
| 171 | record[prefix+tag] = record[tag] |
| 172 | del record[tag] # this is why we iterate over a tuple |
| 173 | # with the tags, and not directly on the record dict |
| 174 | if constant: |
| 175 | constant_key, constant_value = constant.split(':') |