Source code for directupsampling.io.jsonio

# This code contains parts taken from the ASE json extension, found at
# https://gitlab.com/ase/ase
import datetime
import json

import numpy as np
import pandas as pd

# Encoding and writing:




def default(obj):
    if isinstance(obj, np.ndarray) or isinstance(obj, pd.Series):
        if isinstance(obj, pd.Series):
            # Just make Series into arrays
            obj = obj.to_numpy()
        flatobj = obj.ravel()
        if np.iscomplexobj(obj):
            flatobj.dtype = obj.real.dtype
        # We use str(obj.dtype) here instead of obj.dtype.name, because
        # they are not always the same (e.g. for numpy arrays of strings).
        # Using obj.dtype.name can break the ability to recursively decode/
        # encode such arrays.
        return {"__ndarray__": (obj.shape, str(obj.dtype), flatobj.tolist())}

    if isinstance(obj, np.integer):
        return int(obj)

    if isinstance(obj, np.bool_):
        return bool(obj)

    if isinstance(obj, datetime.datetime):
        return {"__datetime__": obj.isoformat()}

    if isinstance(obj, complex):
        return {"__complex__": (obj.real, obj.imag)}

    raise TypeError(f"Cannot convert object of type {type(obj)} to JSON")





class MyEncoder(json.JSONEncoder):


    def default(self, obj):
        return default(obj)






def write_json(fd, obj):
    json.dump(obj, fd, indent=4, cls=MyEncoder)



# Decoding and reading:




def object_hook(dct):
    if "__datetime__" in dct:
        return datetime.datetime.strptime(dct["__datetime__"], "%Y-%m-%dT%H:%M:%S.%f")

    if "__complex__" in dct:
        return complex(*dct["__complex__"])

    if "__ndarray__" in dct:
        return create_ndarray(*dct["__ndarray__"])

    return dct





def create_ndarray(shape, dtype, data):
    """Create ndarray from shape, dtype and flattened data."""
    array = np.empty(shape, dtype=dtype)
    flatbuf = array.ravel()
    if np.iscomplexobj(array):
        flatbuf.dtype = array.real.dtype
    flatbuf[:] = data
    return array





def intkey(key):
    """Convert str to int if possible."""
    try:
        return int(key)
    except ValueError:
        return key





def fix_int_keys_in_dicts(obj):
    """Convert "int" keys: "1" -> 1.

    The json.dump() function will convert int keys in dicts to str keys.
    This function goes the other way.
    """
    if isinstance(obj, dict):
        return {intkey(key): fix_int_keys_in_dicts(value) for key, value in obj.items()}
    return obj





def read_json(fd):
    obj = json.load(fd, object_hook=object_hook)
    obj = fix_int_keys_in_dicts(obj)
    return obj