Source code for torx.fileio.json_io_m

"""
Routines for writing and reading dictionaries to JSON format files.

Automatically handles numpy array I/O.
"""
import json
import numpy as np
from pathlib import Path
from torx.autodoc_decorators_m import autodoc_class, autodoc_function

@autodoc_class
class NumpyEncoder(json.JSONEncoder):
    """Converts numpy arrays into lists."""

    def default(self, obj):
        """Return the default encoding."""
        return obj.tolist()

def recursive_convert_list_to_array(data: dict) -> dict:
    """Convert all list elements in a nested dictionary to numpy array."""
    for key, value in data.items():

        if isinstance(value, list):
            as_array = np.asarray(value)
            if np.issubdtype(as_array.dtype, np.number):
                # Only convert numerical arrays to numpy -- assume that non-numeric
                # lists should be stored as lists
                data[key] = as_array

        elif isinstance(value, dict):
            data[key] = recursive_convert_list_to_array(value)

    return data



@autodoc_function
def write_to_json(data: dict, filepath: Path, allow_overwrite: bool = False):
    """Write a dictionary of data to a JSON file at filepath."""
    assert (
        filepath.suffix == ".json"
    ), f"JSON files should be identified with the .json suffix"

    if not allow_overwrite:
        assert not filepath.exists()

    with open(filepath, "w", encoding="utf-8") as f:
        json.dump(data, f, ensure_ascii=False, indent=4, cls=NumpyEncoder)




@autodoc_function
def read_from_json(filepath: Path) -> dict:
    """
    Read a dictionary of data from a JSON file at filepath.

    If convert_list_to_array is True will convert all lists to numpy array.
    """
    assert filepath.exists(), f"File {filepath} not found"

    with open(filepath, "r", encoding="utf-8") as f:
        data = json.load(f)

    data = recursive_convert_list_to_array(data)

    return data