from copy import deepcopy from enum import Enum import numpy as np from dataclasses import is_dataclass, fields, dataclass import graphlib import json from importlib import import_module from typing import Dict, List, Literal, Optional, TypedDict, Union from types import GeneratorType import logging import warnings import inspect from base64 import b64encode, b64decode from . import plugin from .util import SCHEMA_ATTRIBUTE_NAME, get_libraries DEBUG = False class SCHEMA_VERSIONS(str, Enum): BUMPS_DRAFT_O1 = "bumps-draft-01" BUMPS_DRAFT_02 = "bumps-draft-02" BUMPS_DRAFT_03 = "bumps-draft-03" SCHEMA = SCHEMA_VERSIONS.BUMPS_DRAFT_03 REFERENCES_KEY = "references" REFERENCE_IDENTIFIER = "$ref" MISSING = object() REFERENCE_TYPE_NAME = "Reference" REFERENCE_TYPE = Literal["Reference"] TYPE_KEY = "__class__" @dataclass class Reference: id: str type: REFERENCE_TYPE JSON = Union[Dict[str, "JSON"], List["JSON"], str, int, float, bool, None] class SerializedObject(TypedDict, total=True): schema: str object: JSON references: Dict[str, JSON] def deserialize(serialized: SerializedObject, migration: bool = True): """rehydrate all items in serialzed['references'] then - reydrate all objects in serialized['object'] - replacing `Reference` types with python objects from `references` """ if migration: # first apply built-in migrations: _, serialized = migrate(serialized) # then apply plugin migrations: serialized = plugin.migrate_serialized(serialized) serialized_references = serialized[REFERENCES_KEY] references = {} dependency_graph = {} for ref_id, ref_obj in serialized_references.items(): dependencies = set() _find_ref_dependencies(ref_obj, dependencies) dependency_graph[ref_id] = dependencies sorter = graphlib.TopologicalSorter(dependency_graph) for ref_id in sorter.static_order(): # deserialize and put all references into self.references references[ref_id] = _rehydrate(serialized_references[ref_id], references) # references is now full of deserialized objects, # and we're ready to rehydrate the entire tree... obj = _rehydrate(serialized["object"], references) return obj #### deserializer helpers: def _rehydrate(obj, references: Dict[str, object]): if isinstance(obj, dict): obj = obj.copy() t: str = obj.pop(TYPE_KEY, MISSING) if t == REFERENCE_TYPE_NAME: obj_id: str = obj.get("id", MISSING) if obj_id is MISSING: raise ValueError("object id is required for Reference type") # requires that self.references is populated with rehydrated objects: return references[obj_id] elif t == "bumps.util.NumpyArray": # skip processing values list in ndarray return _to_ndarray(obj) elif t == "Callable": return deserialize_function(obj) else: for key in obj: obj[key] = _rehydrate(obj[key], references) if t is MISSING: # no "type" provided, so no class to instantiate: return hydrated object return obj # obj values are now rehydrated: instantiate the class from 'type' else: try: module_name, class_name = t.rsplit(".", 1) # print(module_name, class_name) klass = getattr(import_module(module_name), class_name) except Exception as e: # there is a type, but it is not found... logging.exception(e) raise RuntimeError(f"Error importing {t}: {e}") hydrated = _instantiate(klass, t, obj) # print("returning instantiated object", t, hydrated) return hydrated elif isinstance(obj, list): # rehydrate all the items return [_rehydrate(v, references) for v in obj] else: # it's a bare value - just return return obj def _instantiate(klass: type, typename: str, serialized: dict): # TODO: why are we copying the top-level dict? serialized = serialized.copy() # if klass provides 'from_dict' method, use it - # otherwise use klass.__init__ directly. class_factory = getattr(klass, "from_dict", klass) try: hydrated = class_factory(**serialized) except Exception as e: logging.exception(e) warnings.warn(f"Error restoring {typename}: {e}") # Note: users of the failed deserialization may complain that it is None hydrated = None return hydrated def _to_ndarray(obj: dict): return np.asarray(obj["values"], dtype=np.dtype(obj.get("dtype", float))) def _find_ref_dependencies(obj, dependencies: set): if isinstance(obj, dict): if obj.get(TYPE_KEY, None) == REFERENCE_TYPE: dependencies.add(obj["id"]) else: for v in obj.values(): _find_ref_dependencies(v, dependencies) elif isinstance(obj, list): for v in obj: _find_ref_dependencies(v, dependencies) #### end deserializer helpers def serialize(obj, use_refs=True, add_libraries=True): references = {} def make_ref(obj_id: str): return {"id": obj_id, TYPE_KEY: REFERENCE_TYPE_NAME} def dataclass_to_dict(dclass, include=None, exclude=None): all_fields = fields(dclass) if include is not None: all_fields = [f for f in all_fields if f.name in include] elif exclude is not None: all_fields = [f for f in all_fields if not f.name in exclude and not f.name.startswith("_")] else: all_fields = [f for f in all_fields if not f.name.startswith("_")] cls = dclass.__class__ fqn = f"{cls.__module__}.{cls.__qualname__}" output = dict([(f.name, obj_to_dict(getattr(dclass, f.name))) for f in all_fields]) output[TYPE_KEY] = fqn return output def obj_to_dict(obj): if hasattr(obj, SCHEMA_ATTRIBUTE_NAME): schema_opts = getattr(obj, SCHEMA_ATTRIBUTE_NAME) include = schema_opts.get("include", None) exclude = schema_opts.get("exclude", None) use_ref = use_refs and hasattr(obj, "id") if (not use_ref) or (obj.id not in references): # only calculate dict if it's not already in refs, or if not using refs obj_dict = dataclass_to_dict(obj, include=include, exclude=exclude) if use_ref: references.setdefault(obj.id, obj_dict) return make_ref(obj.id) if use_ref else obj_dict elif is_dataclass(obj): return dataclass_to_dict(obj) elif isinstance(obj, (list, tuple, GeneratorType)): return list(obj_to_dict(v) for v in obj) # elif isinstance(obj, GeneratorType) and issubclass(obj_type, Tuple): # return tuple(to_dict(v) for v in obj) elif isinstance(obj, dict): return type(obj)((obj_to_dict(k), obj_to_dict(v)) for k, v in obj.items()) elif isinstance(obj, np.ndarray) and obj.dtype.kind in ["f", "i", "U"]: return {TYPE_KEY: "bumps.util.NumpyArray", "dtype": str(obj.dtype), "values": obj.tolist()} elif isinstance(obj, np.ndarray) and obj.dtype.kind == "O": return obj_to_dict(obj.tolist()) elif isinstance(obj, Enum): return obj_to_dict(obj.value) elif isinstance(obj, float): return str(obj) if np.isinf(obj) else obj elif isinstance(obj, int) or isinstance(obj, str) or obj is None: return obj elif callable(obj): return serialize_function(obj) else: raise ValueError("obj %s is not serializable" % str(obj)) serialized = {"$schema": SCHEMA, "object": obj_to_dict(obj), REFERENCES_KEY: references} if add_libraries: serialized["libraries"] = get_libraries(obj) # print("serialized as", serialized) return serialized def deserialize_function(obj): import dill import cloudpickle try: # CRUFT: older versions did not specify the pickler pickler = obj.get("pickler", "dill") data = deserialize_bytes(obj["pickle"]) if pickler == "dill": return dill.loads(data) elif pickler == "cloudpickle": return cloudpickle.loads(deserialize_bytes(obj["pickle"])) else: raise ValueError(f"unrecognized pickler {pickler}") except Exception as e: logging.exception(e) warnings.warn(f"Error loading function: {e}") return None def serialize_bytes(b): return b64encode(b).decode("ascii") def deserialize_bytes(s): return b64decode(s) def serialize_function(fn): from cloudpickle import dumps name = getattr(fn, "__name__", "unknown") # print("type fn", type(fn)) # Note: need dedent to handle decorator syntax. Dedent will fail when there are # triple-quoted strings. Alternative: if first character is a space, then wrap # the code in an "if True:" block # source = dedent(inspect.getsource(fn)) #.strip() try: source = inspect.getsource(fn) except Exception: source = None # print("source =>", source) pickle = serialize_bytes(dumps(fn)) res = {TYPE_KEY: "Callable", "name": name, "source": source, "pickle": pickle, "pickler": "cloudpickle"} # print(f"serializing {fn} to {res}") return res def save_file(filename, problem): try: p = serialize(problem) with open(filename, "w") as fid: json.dump(p, fid) except Exception as e: logging.exception(e) warnings.warn(f"failed to create JSON file {filename} for fit problem") def load_file(filename): with open(filename, "r") as fid: serialized: SerializedObject = json.loads(fid.read()) final_version, migrated = migrate(serialized) # print("final version: ", final_version) return deserialize(migrated) #### MIGRATIONS def validate_version(version: str, variable_name="from_version"): if version not in list(SCHEMA_VERSIONS): raise ValueError(f"must choose a valid {variable_name} from this list: {[s.value for s in SCHEMA_VERSIONS]}") def migrate( serialized: dict, from_version: Optional[SCHEMA_VERSIONS] = None, to_version: Optional[SCHEMA_VERSIONS] = SCHEMA ): """ Migrate a serialized object from one version to another By default, the `from_version` is determined by inspection of the serialized object. This is overriden by setting the `from_version` keyword argument to a member of `SCHEMA_VERSIONS` Also by default, the target version is the current schema, which can be overriden with the `to_version` keyword argument """ if from_version is None: from_version = serialized.get( "$schema", SCHEMA_VERSIONS.BUMPS_DRAFT_O1 ) # fall back to first version if not specified validate_version(from_version, "from_version") validate_version(to_version, "to_version") current_version = from_version while current_version != to_version: print(f"migrating {current_version}...") current_version, serialized = MIGRATIONS[current_version](serialized) return current_version, serialized def _migrate_draft_01_to_draft_02(serialized: dict): references = {} def rename_type(obj): if isinstance(obj, dict): t: str = obj.pop("type", obj.pop(TYPE_KEY, MISSING)) # print(f"moving type to __class__ for id {obj.get('id', 'no id')}") if t is not MISSING: obj[TYPE_KEY] = t for v in obj.values(): rename_type(v) elif isinstance(obj, list): for v in obj: rename_type(v) def build_references(obj): if isinstance(obj, dict): t: str = obj.get(TYPE_KEY, MISSING) obj_id: str = obj.get("id", MISSING) # if obj_id is not MISSING: # print(f"building reference for id {obj_id}") if obj_id is not MISSING and not t in [MISSING, REFERENCE_TYPE_NAME]: if not obj_id in references: references[obj_id] = deepcopy(obj) obj[TYPE_KEY] = REFERENCE_TYPE_NAME for k in list(obj.keys()): if k not in [TYPE_KEY, "id"]: del obj[k] for v in obj.values(): build_references(v) elif isinstance(obj, list): for v in obj: build_references(v) rename_type(serialized) build_references(serialized) migrated = { "$schema": SCHEMA_VERSIONS.BUMPS_DRAFT_02.value, "object": serialized, "references": references, } return SCHEMA_VERSIONS.BUMPS_DRAFT_02, migrated def _migrate_draft_02_to_draft_03(serialized: dict): # add migration code here def div_to_truediv(obj): # remove all 'div' operators and replace with 'truediv' if isinstance(obj, dict) and obj.get(TYPE_KEY, MISSING) == "bumps.parameter.Expression": if obj.get("op", MISSING) == "div": obj["op"] = "truediv" for v in obj.get("args", []): div_to_truediv(v) elif isinstance(obj, dict): for k, v in obj.items(): div_to_truediv(v) elif isinstance(obj, list): for v in obj: div_to_truediv(v) migrated = deepcopy(serialized) div_to_truediv(migrated) return SCHEMA_VERSIONS.BUMPS_DRAFT_03, migrated MIGRATIONS = { SCHEMA_VERSIONS.BUMPS_DRAFT_O1: _migrate_draft_01_to_draft_02, SCHEMA_VERSIONS.BUMPS_DRAFT_02: _migrate_draft_02_to_draft_03, }