diff --git a/pyerrors/input/json.py b/pyerrors/input/json.py index 6f5cf470..6b854874 100644 --- a/pyerrors/input/json.py +++ b/pyerrors/input/json.py @@ -6,6 +6,7 @@ import socket import datetime import platform import warnings +import re from ..obs import Obs from ..covobs import Covobs from ..correlators import Corr @@ -20,7 +21,7 @@ def create_json_string(ol, description='', indent=1): Parameters ---------- ol : list - List of objects that will be exported. At the moments, these objects can be + List of objects that will be exported. At the moment, these objects can be either of: Obs, list, numpy.ndarray, Corr. All Obs inside a structure have to be defined on the same set of configurations. description : str @@ -246,7 +247,7 @@ def dump_to_json(ol, fname, description='', indent=1, gz=True): Parameters ---------- ol : list - List of objects that will be exported. At the moments, these objects can be + List of objects that will be exported. At the moment, these objects can be either of: Obs, list, numpy.ndarray, Corr. All Obs inside a structure have to be defined on the same set of configurations. fname : str @@ -476,7 +477,7 @@ def import_json_string(json_string, verbose=True, full_output=False): def load_json(fname, verbose=True, gz=True, full_output=False): - """Import a list of Obs or structures containing Obs from a .json.gz file. + """Import a list of Obs or structures containing Obs from a .json(.gz) file. The following structures are supported: Obs, list, numpy.ndarray, Corr If the list contains only one element, it is unpacked from the list. @@ -507,3 +508,215 @@ def load_json(fname, verbose=True, gz=True, full_output=False): d = fin.read() return import_json_string(d, verbose, full_output) + + +def _ol_from_dict(ind, reps='DICTOBS'): + """Convert a dictionary of Obs objects to a list and a dictionary that contains + placeholders instead of the Obs objects. + + Parameters + ---------- + ind : dict + Dict of JSON valid structures and objects that will be exported. + At the moment, these object can be either of: Obs, list, numpy.ndarray, Corr. + All Obs inside a structure have to be defined on the same set of configurations. + reps : str + Specify the structure of the placeholder in exported dict to be reps[0-9]+. + """ + + obstypes = (Obs, Corr, np.ndarray) + + if not reps.isalnum(): + raise Exception('Placeholder string has to be alphanumeric!') + ol = [] + counter = 0 + + def dict_replace_obs(d): + nonlocal ol + nonlocal counter + x = {} + for k, v in d.items(): + if isinstance(v, dict): + v = dict_replace_obs(v) + elif isinstance(v, list) and all([isinstance(o, Obs) for o in v]): + v = obslist_replace_obs(v) + elif isinstance(v, list): + v = list_replace_obs(v) + elif isinstance(v, obstypes): + ol.append(v) + v = reps + '%d' % (counter) + counter += 1 + elif isinstance(v, str): + if bool(re.match(r'%s[0-9]+' % (reps), v)): + raise Exception('Dict contains string %s that matches the placeholder! %s Cannot be savely exported.' % (v, reps)) + x[k] = v + return x + + def list_replace_obs(li): + nonlocal ol + nonlocal counter + x = [] + for e in li: + if isinstance(e, list): + e = list_replace_obs(e) + elif isinstance(e, list) and all([isinstance(o, Obs) for o in e]): + e = obslist_replace_obs(e) + elif isinstance(e, dict): + e = dict_replace_obs(e) + elif isinstance(e, obstypes): + ol.append(e) + e = reps + '%d' % (counter) + counter += 1 + elif isinstance(e, str): + if bool(re.match(r'%s[0-9]+' % (reps), e)): + raise Exception('Dict contains string %s that matches the placeholder! %s Cannot be savely exported.' % (e, reps)) + x.append(e) + return x + + def obslist_replace_obs(li): + nonlocal ol + nonlocal counter + il = [] + for e in li: + il.append(e) + + ol.append(il) + x = reps + '%d' % (counter) + counter += 1 + return x + + nd = dict_replace_obs(ind) + + return ol, nd + + +def dump_dict_to_json(od, fname, description='', indent=1, reps='DICTOBS', gz=True): + """Export a dict of Obs or structures containing Obs to a .json(.gz) file + + Parameters + ---------- + od : dict + Dict of JSON valid structures and objects that will be exported. + At the moment, these objects can be either of: Obs, list, numpy.ndarray, Corr. + All Obs inside a structure have to be defined on the same set of configurations. + fname : str + Filename of the output file. + description : str + Optional string that describes the contents of the json file. + indent : int + Specify the indentation level of the json file. None or 0 is permissible and + saves disk space. + reps : str + Specify the structure of the placeholder in exported dict to be reps[0-9]+. + gz : bool + If True, the output is a gzipped json. If False, the output is a json file. + """ + + if not isinstance(od, dict): + raise Exception('od has to be a dictionary. Did you want to use dump_to_json?') + + infostring = ('This JSON file contains a python dictionary that has been parsed to a list of structures. ' + 'OBSDICT contains the dictionary, where Obs or other structures have been replaced by ' + '' + reps + '[0-9]+. The field description contains the additional description of this JSON file. ' + 'This file may be parsed to a dict with the pyerrors routine load_json_dict.') + + desc_dict = {'INFO': infostring, 'OBSDICT': {}, 'description': description} + ol, desc_dict['OBSDICT'] = _ol_from_dict(od, reps=reps) + + dump_to_json(ol, fname, description=desc_dict, indent=indent, gz=gz) + + +def _od_from_list_and_dict(ol, ind, reps='DICTOBS'): + """Parse a list of Obs or structures containing Obs and an accompanying + dict, where the structures have been replaced by placeholders to a + dict that contains the structures. + + The following structures are supported: Obs, list, numpy.ndarray, Corr + + Parameters + ---------- + ol : list + List of objects - + At the moment, these objects can be either of: Obs, list, numpy.ndarray, Corr. + All Obs inside a structure have to be defined on the same set of configurations. + ind : dict + Dict that defines the structure of the resulting dict and contains placeholders + reps : str + Specify the structure of the placeholder in imported dict to be reps[0-9]+. + """ + if not reps.isalnum(): + raise Exception('Placeholder string has to be alphanumeric!') + + counter = 0 + + def dict_replace_string(d): + nonlocal counter + nonlocal ol + x = {} + for k, v in d.items(): + if isinstance(v, dict): + v = dict_replace_string(v) + elif isinstance(v, list): + v = list_replace_string(v) + elif isinstance(v, str) and bool(re.match(r'%s[0-9]+' % (reps), v)): + index = int(v[len(reps):]) + v = ol[index] + counter += 1 + x[k] = v + return x + + def list_replace_string(li): + nonlocal counter + nonlocal ol + x = [] + for e in li: + if isinstance(e, list): + e = list_replace_string(e) + elif isinstance(e, dict): + e = dict_replace_string(e) + elif isinstance(e, str) and bool(re.match(r'%s[0-9]+' % (reps), e)): + index = int(e[len(reps):]) + e = ol[index] + counter += 1 + x.append(e) + return x + + nd = dict_replace_string(ind) + + if counter == 0: + raise Exception('No placeholder has been replaced! Check if reps is set correctly.') + + return nd + + +def load_json_dict(fname, verbose=True, gz=True, full_output=False, reps='DICTOBS'): + """Import a dict of Obs or structures containing Obs from a .json(.gz) file. + + The following structures are supported: Obs, list, numpy.ndarray, Corr + + Parameters + ---------- + fname : str + Filename of the input file. + verbose : bool + Print additional information that was written to the file. + gz : bool + If True, assumes that data is gzipped. If False, assumes JSON file. + full_output : bool + If True, a dict containing auxiliary information and the data is returned. + If False, only the data is returned. + reps : str + Specify the structure of the placeholder in imported dict to be reps[0-9]+. + """ + indata = load_json(fname, verbose=verbose, gz=gz, full_output=True) + description = indata['description']['description'] + indict = indata['description']['OBSDICT'] + ol = indata['obsdata'] + od = _od_from_list_and_dict(ol, indict, reps=reps) + + if full_output: + indata['description'] = description + indata['obsdata'] = od + return indata + else: + return od diff --git a/tests/io_test.py b/tests/io_test.py index 88f0a3ae..f5546d9f 100644 --- a/tests/io_test.py +++ b/tests/io_test.py @@ -3,6 +3,7 @@ import gzip import numpy as np import pyerrors as pe import pyerrors.input.json as jsonio +import pytest def test_jsonio(): @@ -136,3 +137,108 @@ def test_json_corr_2d_io(): assert recover[index] is None assert my_corr.tag == recover.tag assert my_corr.prange == recover.prange + + +def test_json_dict_io(): + def check_dict_equality(d1, d2): + def dict_check_obs(d1, d2): + for k, v in d1.items(): + if isinstance(v, dict): + v = dict_check_obs(v, d2[k]) + elif isinstance(v, list) and all([isinstance(o, pe.Obs) for o in v]): + for i in range(len(v)): + assert((v[i] - d2[k][i]).is_zero()) + elif isinstance(v, list): + v = list_check_obs(v, d2[k]) + elif isinstance(v, pe.Obs): + assert((v - d2[k]).is_zero()) + elif isinstance(v, pe.Corr): + for i in range(v.T): + assert((v[i] - d2[k][i]).is_zero()) + elif isinstance(v, np.ndarray): + a1 = np.ravel(v) + a2 = np.ravel(d2[k]) + for i in range(len(a1)): + assert((a1[i] - a2[i]).is_zero()) + + def list_check_obs(l1, l2): + for ei in range(len(l1)): + e = l1[ei] + if isinstance(e, list): + e = list_check_obs(e, l2[ei]) + elif isinstance(e, list) and all([isinstance(o, pe.Obs) for o in e]): + for i in range(len(e)): + assert((e[i] - l2[ei][i]).is_zero()) + elif isinstance(e, dict): + e = dict_check_obs(e, l2[ei]) + elif isinstance(e, pe.Obs): + assert((e - l2[ei]).is_zero()) + elif isinstance(e, pe.Corr): + for i in range(e.T): + assert((e[i] - l2[ei][i]).is_zero()) + elif isinstance(e, np.ndarray): + a1 = np.ravel(e) + a2 = np.ravel(l2[ei]) + for i in range(len(a1)): + assert((a1[i] - a2[i]).is_zero()) + dict_check_obs(d1, d2) + return True + + od = { + 'l': + { + 'a': pe.pseudo_Obs(1, .2, 'testa', samples=10), + 'b': [pe.pseudo_Obs(1.1, .1, 'test', samples=10), pe.pseudo_Obs(1.2, .1, 'test', samples=10), pe.pseudo_Obs(1.3, .1, 'test', samples=10)], + 'c': { + 'd': 1, + 'e': pe.pseudo_Obs(.2, .01, 'teste', samples=10), + 'f': pe.Corr([pe.pseudo_Obs(.1, .01, 'a', samples=10), pe.pseudo_Obs(.1, .01, 'a', samples=10), pe.pseudo_Obs(.1, .01, 'a', samples=10), pe.pseudo_Obs(.1, .01, 'a', samples=10)]), + 'g': np.reshape(np.asarray([pe.pseudo_Obs(.1, .01, 'a', samples=10), pe.pseudo_Obs(.1, .01, 'a', samples=10), pe.pseudo_Obs(.1, .01, 'a', samples=10), pe.pseudo_Obs(.1, .01, 'a', samples=10)]), (2, 2)), + } + }, + 's': + { + 'a': 'Infor123', + 'b': ['Some', 'list'], + 'd': pe.pseudo_Obs(.01, .001, 'testd', samples=10) * pe.cov_Obs(1, .01, 'cov1'), + 'se': None, + 'sf': 1.2, + } + } + + fname = 'test_rw' + + desc = 'This is a random description' + + with pytest.raises(Exception): + jsonio.dump_dict_to_json(od, fname, description=desc, reps='|Test') + + jsonio.dump_dict_to_json(od, fname, description=desc, reps='TEST') + nd = jsonio.load_json_dict(fname, full_output=True, reps='TEST') + + with pytest.raises(Exception): + nd = jsonio.load_json_dict(fname, full_output=True) + + jsonio.dump_dict_to_json(od, fname, description=desc) + nd = jsonio.load_json_dict(fname, full_output=True) + assert (desc == nd['description']) + + assert(check_dict_equality(od, nd['obsdata'])) + nd = jsonio.load_json_dict(fname, full_output=False) + assert(check_dict_equality(od, nd)) + + nl = jsonio.load_json(fname, full_output=True) + nl = jsonio.load_json(fname, full_output=False) + + with pytest.raises(Exception): + jsonio.dump_dict_to_json(nl, fname, description=desc) + + od['k'] = 'DICTOBS2' + with pytest.raises(Exception): + jsonio.dump_dict_to_json(od, fname, description=desc) + + od['k'] = ['DICTOBS2'] + with pytest.raises(Exception): + jsonio.dump_dict_to_json(od, fname, description=desc) + + os.remove(fname + '.json.gz')