Source code for uf3.util.json_io

"""
This module provides functions for serializing and loading nested dictionaries
with JSON files.
"""

import json
from typing import Union
import numpy as np


[docs]def dump_interaction_map(interaction_map, indent=4, filename=None, write=False): """ Utility function for writing ragged arrays to json file. e.g. {("A", "B"): [[1, 2, 3], [4, 5], [6, 7, 8, 9]]} Args: interaction_map (dict): map of interaction to ragged array containing np.ndarry vectors. indent (int): number of spaces to indent. filename (str): name of file to write. write (bool): whether to write to file. """ formatted_map = encode_interaction_map(interaction_map) text = json.dumps(formatted_map, indent=indent, cls=CompactJSONEncoder) if write: with open(filename, 'w') as f: f.write(text) else: return text
[docs]def encode_interaction_map(interaction_map): """Recursive function for converting arrays to lists and tuples into dash-joined keys for JSON serialization.""" encoded_map = {} for key, value in interaction_map.items(): if isinstance(value, list): # array to list if isinstance(value[0], np.ndarray): value = [entry.tolist() for entry in value] if isinstance(value, np.ndarray): # array to list value = value.tolist() elif isinstance(value, dict): value = encode_interaction_map(value) if isinstance(key, tuple): # tuple to joined str key = '-'.join([str(item) for item in key]) encoded_map[key] = value return encoded_map
[docs]def load_interaction_map(filename): """Parse interaction map(s) from JSON.""" with open(filename, "r") as f: formatted_map = json.load(f) interaction_map = decode_interaction_map(formatted_map) return interaction_map
[docs]def decode_interaction_map(formatted_map): """Recursive function for converting lists to arrays and dash-joined keys into tuples for JSON deserialization.""" decoded_map = {} for key, value in formatted_map.items(): if isinstance(value, list): # list to array if isinstance(value[0], list): value = [np.array(row) for row in value] else: value = np.array(value) elif isinstance(value, dict): value = decode_interaction_map(value) if '-' in key: # joined str to tuple key = key.split('-') try: key = [int(i) for i in key] except ValueError: pass key = tuple(key) decoded_map[key] = value return decoded_map
[docs]class CompactJSONEncoder(json.JSONEncoder): """ A JSON Encoder that formats vectors into single lines. Discussion on StackOverflow: https://stackoverflow.com/questions/16264515/ Original question by Saar Drimer: https://stackoverflow.com/users/458116/saar-drimer Original answer by Tim Ludwinski: https://stackoverflow.com/users/1413201/tim-ludwinski Adaptation by Jannis Mainczyk: https://stackoverflow.com/users/5379377/jmm Adapted code under CC BY-SA 3.0 license. """ CONTAINER_TYPES = (list, tuple, dict) """Container datatypes include primitives or other containers.""" INDENTATION_CHAR = " " def __init__(self, *args, **kwargs): super().__init__(*args, **kwargs) self.indentation_level = 0
[docs] def encode(self, o): """Encode JSON object *o* with respect to single line lists.""" if isinstance(o, (list, tuple)): if self._put_on_single_line(o): return "[" + ", ".join(self.encode(el) for el in o) + "]" else: self.indentation_level += 1 output = [self.indent_str + self.encode(el) for el in o] self.indentation_level -= 1 return ("[\n" + ",\n".join(output) + "\n" + self.indent_str + "]") elif isinstance(o, dict): if o: if self._put_on_single_line(o): return ("{ " + ", ".join(f"{self.encode(k)}: {self.encode(el)}" for k, el in o.items()) + " }") else: self.indentation_level += 1 output = [self.indent_str + f"{json.dumps(k)}: {self.encode(v)}" for k, v in o.items()] self.indentation_level -= 1 return ("{\n" + ",\n".join(output) + "\n" + self.indent_str + "}") else: return "{}" elif isinstance(o, float): # Use scientific notation for floats, where appropiate return format(o, "g") elif isinstance(o, str): # escape newlines o = o.replace("\n", "\\n") return f'"{o}"' else: return json.dumps(o)
def _put_on_single_line(self, o): return self._primitives_only(o) def _primitives_only(self, o: Union[list, tuple, dict]): if isinstance(o, (list, tuple)): return not any(isinstance(el, self.CONTAINER_TYPES) for el in o) elif isinstance(o, dict): return not any(isinstance(el, self.CONTAINER_TYPES) for el in o.values()) @property def indent_str(self) -> str: return self.INDENTATION_CHAR * (self.indentation_level * self.indent)