diff --git a/LoopStructural/modelling/core/_model_serializer.py b/LoopStructural/modelling/core/_model_serializer.py new file mode 100644 index 00000000..cbf4d112 --- /dev/null +++ b/LoopStructural/modelling/core/_model_serializer.py @@ -0,0 +1,240 @@ +"""Recipe (JSON) and pickle serialization logic for GeologicalModel. + +Extracted from GeologicalModel to separate serialization concerns from +feature-container orchestration (see API.md). GeologicalModel's +``@public_api``-decorated methods (``to_dict``, ``to_recipe_dict``, +``from_recipe_dict``, ``to_recipe_json``, ``from_recipe_json``, +``save_recipe``, ``load_recipe``, ``to_file``, ``from_file``) stay defined +directly on the class -- their ``__qualname__`` is part of the CI-checked +stable API surface (``tests/unit/test_public_api_contract.py``) -- and just +delegate to the staticmethods here. +""" + +import json +import pathlib + +import pandas as pd + +from ...geometry import BoundingBox +from ...utils import LoopValueError, getLogger +from ..features import GeologicalFeature, StructuralFrame, UnconformityFeature +from ..features.fault import FaultSegment +from .stratigraphic_column import StratigraphicColumn + +logger = getLogger(__name__) + + +class ModelSerializer: + @staticmethod + def feature_recipe_kind(feature): + if isinstance(feature, GeologicalFeature): + return "foliation" + if isinstance(feature, StructuralFrame): + return "structural_frame" + if isinstance(feature, UnconformityFeature): + return "unconformity" + if isinstance(feature, FaultSegment): + return "fault" + return feature.__class__.__name__.lower() + + @staticmethod + def to_dict(model): + result = {} + result["model"] = {} + result["model"]["features"] = [f.name for f in model.features] + result['model']['bounding_box'] = model.bounding_box.to_dict() + result["model"]["stratigraphic_column"] = model.stratigraphic_column + return result + + @staticmethod + def to_recipe_dict(model, data_reference=None): + recipe = { + "schema": "LoopStructural.GeologicalModelRecipe", + "version": 1, + "model": { + "bounding_box": model.bounding_box.to_dict(), + "stratigraphic_column": model.stratigraphic_column.to_dict(), + "data_source": None, + "features": [], + }, + } + for feature in model.features: + feature_entry = { + "name": feature.name, + "kind": ModelSerializer.feature_recipe_kind(feature), + "faults": [ + fault.name + for fault in getattr(feature, "faults", []) + if getattr(fault, "name", None) + ], + "regions": [], + } + recipe["model"]["features"].append(feature_entry) + if data_reference is not None: + recipe["model"]["data_source"] = { + "kind": "reference", + "path": str(pathlib.Path(data_reference)), + } + elif not model.data.empty: + recipe["model"]["data_source"] = { + "kind": "inline", + "dataframe": model.data.to_dict(orient="split"), + } + return recipe + + @staticmethod + def from_recipe_dict(cls, recipe): + if not isinstance(recipe, dict): + raise TypeError("recipe must be a dictionary") + + model_data = recipe.get("model", recipe) + bounding_box = model_data.get("bounding_box") + if isinstance(bounding_box, dict): + bounding_box = BoundingBox.from_dict(bounding_box) + if not isinstance(bounding_box, BoundingBox): + raise TypeError("recipe must include a bounding_box dictionary") + + model = cls(bounding_box) + + data_source = model_data.get("data_source") + if isinstance(data_source, dict): + kind = data_source.get("kind") + if kind == "reference": + model.data = pd.read_csv(pathlib.Path(data_source["path"])) + elif kind == "inline": + dataframe = data_source.get("dataframe") + if not isinstance(dataframe, dict): + raise TypeError("inline data_source must include a dataframe dictionary") + model.data = pd.DataFrame(**dataframe) + elif kind is not None: + raise ValueError(f"Unsupported data_source kind: {kind}") + elif isinstance(data_source, str): + model.data = pd.read_csv(pathlib.Path(data_source)) + elif data_source is not None: + raise TypeError("data_source must be a dictionary, string path, or None") + + stratigraphic_column = model_data.get("stratigraphic_column") + if isinstance(stratigraphic_column, dict): + model.stratigraphic_column = StratigraphicColumn.from_dict(stratigraphic_column) + elif stratigraphic_column is not None: + raise TypeError("stratigraphic_column must be a dictionary or None") + + features = model_data.get("features", []) + if features is None: + features = [] + if not isinstance(features, list): + raise TypeError("features must be a list") + + feature_map = {} + for feature_entry in features: + if not isinstance(feature_entry, dict): + raise TypeError("each feature entry must be a dictionary") + feature_name = feature_entry.get("name") + if not isinstance(feature_name, str): + raise TypeError("each feature entry must include a string name") + feature_data = model.data.loc[model.data["feature_name"] == feature_name].copy() + if feature_data.empty: + feature_data = None + feature = model.create_and_add_foliation(feature_name, data=feature_data) + if feature is None: + raise ValueError(f"Could not recreate feature '{feature_name}' from recipe") + feature_map[feature_name] = feature + + for feature_entry in features: + feature_name = feature_entry.get("name") + fault_names = feature_entry.get("faults", []) + if not isinstance(fault_names, list): + raise TypeError("faults must be a list") + if fault_names: + feature = feature_map[feature_name] + feature.faults = [feature_map[name] for name in fault_names if name in feature_map] + + return model + + @staticmethod + def to_recipe_json(model, data_reference=None, indent=2): + recipe = ModelSerializer.to_recipe_dict(model, data_reference=data_reference) + return json.dumps(recipe, indent=indent) + + @staticmethod + def from_recipe_json(cls, json_str): + if not isinstance(json_str, str): + raise TypeError("json_str must be a string") + try: + recipe = json.loads(json_str) + except json.JSONDecodeError as e: + raise TypeError(f"json_str is not valid JSON: {e}") + return ModelSerializer.from_recipe_dict(cls, recipe) + + @staticmethod + def save_recipe(model, filename, data_reference=None): + filename = pathlib.Path(filename) + recipe = ModelSerializer.to_recipe_dict(model, data_reference=data_reference) + with open(filename, "w") as f: + json.dump(recipe, f, indent=2) + logger.info(f"Recipe saved to {filename}") + + @staticmethod + def load_recipe(cls, filename): + filename = pathlib.Path(filename) + if not filename.exists(): + raise FileNotFoundError(f"Recipe file not found: {filename}") + with open(filename, "r") as f: + recipe = json.load(f) + logger.info(f"Recipe loaded from {filename}") + return ModelSerializer.from_recipe_dict(cls, recipe) + + @staticmethod + def to_file(model, file): + try: + import dill as pickle + except ImportError: + logger.error("Cannot write to file, dill not installed \n" "pip install dill") + return + try: + logger.info(f"Writing GeologicalModel to: {file}") + with open(file, "wb") as handle: + pickle.dump(model, handle) + except pickle.PicklingError: + logger.error("Error saving file") + + @staticmethod + def from_file(cls, file, allow_pickle: bool = True): + if not allow_pickle: + raise LoopValueError( + "Pickle-based loading is disabled (allow_pickle=False). " + f"Refusing to unpickle '{file}' because deserialising untrusted " + "pickle/dill data can execute arbitrary code. If you generated " + "this file yourself and trust its contents, call " + "GeologicalModel.from_file(file, allow_pickle=True). Otherwise, " + "use the JSON-based GeologicalModel.from_recipe_dict " + "(paired with GeologicalModel.to_recipe_dict) as a safe " + "alternative serialisation format." + ) + logger.warning( + f"Loading GeologicalModel from '{file}' using dill/pickle. " + "Only load model files from trusted sources: deserialising a " + "pickle file can execute arbitrary code. Pass allow_pickle=False " + "to refuse pickle-based loading, or use " + "GeologicalModel.from_recipe_dict for untrusted/JSON-based input." + ) + try: + import dill as pickle + except ImportError: + logger.error("Cannot import from file, dill not installed") + return None + path = pathlib.Path(file) + if not path.is_file(): + raise LoopValueError(f"Cannot load model, file does not exist: {file}") + try: + with open(path, "rb") as f: + model = pickle.load(f) + except Exception as e: + logger.error(f"Failed to load model from {file}: {e}") + raise LoopValueError(f"Failed to load model from {file}: {e}") from e + if isinstance(model, cls): + logger.info("GeologicalModel initialised from file") + return model + else: + logger.error(f"{file} does not contain a geological model") + return None diff --git a/LoopStructural/modelling/core/geological_model.py b/LoopStructural/modelling/core/geological_model.py index 74a9458f..8511e76a 100644 --- a/LoopStructural/modelling/core/geological_model.py +++ b/LoopStructural/modelling/core/geological_model.py @@ -3,7 +3,6 @@ """ from __future__ import annotations -import json import pathlib import warnings @@ -44,6 +43,7 @@ convert_feature_to_structural_frame as _convert_feature_to_structural_frame, ) from ._feature_registry import FeatureBuilderRegistry +from ._model_serializer import ModelSerializer from .stratigraphic_column import StratigraphicColumn logger = getLogger(__name__) @@ -154,25 +154,7 @@ def to_dict(self): json : str json string of the geological model """ - json = {} - json["model"] = {} - json["model"]["features"] = [f.name for f in self.features] - json['model']['bounding_box'] = self.bounding_box.to_dict() - json["model"]["stratigraphic_column"] = self.stratigraphic_column - # json["features"] = [f.to_json() for f in self.features] - return json - - @staticmethod - def _feature_recipe_kind(feature): - if isinstance(feature, GeologicalFeature): - return "foliation" - if isinstance(feature, StructuralFrame): - return "structural_frame" - if isinstance(feature, UnconformityFeature): - return "unconformity" - if isinstance(feature, FaultSegment): - return "fault" - return feature.__class__.__name__.lower() + return ModelSerializer.to_dict(self) @public_api(tier="provisional") def to_recipe_dict(self, data_reference=None): @@ -182,110 +164,13 @@ def to_recipe_dict(self, data_reference=None): box, stratigraphic column, and either inline model data or a file reference to it. """ - recipe = { - "schema": "LoopStructural.GeologicalModelRecipe", - "version": 1, - "model": { - "bounding_box": self.bounding_box.to_dict(), - "stratigraphic_column": self.stratigraphic_column.to_dict(), - "data_source": None, - "features": [], - }, - } - for feature in self.features: - feature_entry = { - "name": feature.name, - "kind": self._feature_recipe_kind(feature), - "faults": [ - fault.name - for fault in getattr(feature, "faults", []) - if getattr(fault, "name", None) - ], - "regions": [], - } - recipe["model"]["features"].append(feature_entry) - if data_reference is not None: - recipe["model"]["data_source"] = { - "kind": "reference", - "path": str(pathlib.Path(data_reference)), - } - elif not self.data.empty: - recipe["model"]["data_source"] = { - "kind": "inline", - "dataframe": self.data.to_dict(orient="split"), - } - return recipe + return ModelSerializer.to_recipe_dict(self, data_reference=data_reference) @classmethod @public_api(tier="provisional") def from_recipe_dict(cls, recipe): """Rebuild a geological model from a recipe dictionary.""" - if not isinstance(recipe, dict): - raise TypeError("recipe must be a dictionary") - - model_data = recipe.get("model", recipe) - bounding_box = model_data.get("bounding_box") - if isinstance(bounding_box, dict): - bounding_box = BoundingBox.from_dict(bounding_box) - if not isinstance(bounding_box, BoundingBox): - raise TypeError("recipe must include a bounding_box dictionary") - - model = cls(bounding_box) - - data_source = model_data.get("data_source") - if isinstance(data_source, dict): - kind = data_source.get("kind") - if kind == "reference": - model.data = pd.read_csv(pathlib.Path(data_source["path"])) - elif kind == "inline": - dataframe = data_source.get("dataframe") - if not isinstance(dataframe, dict): - raise TypeError("inline data_source must include a dataframe dictionary") - model.data = pd.DataFrame(**dataframe) - elif kind is not None: - raise ValueError(f"Unsupported data_source kind: {kind}") - elif isinstance(data_source, str): - model.data = pd.read_csv(pathlib.Path(data_source)) - elif data_source is not None: - raise TypeError("data_source must be a dictionary, string path, or None") - - stratigraphic_column = model_data.get("stratigraphic_column") - if isinstance(stratigraphic_column, dict): - model.stratigraphic_column = StratigraphicColumn.from_dict(stratigraphic_column) - elif stratigraphic_column is not None: - raise TypeError("stratigraphic_column must be a dictionary or None") - - features = model_data.get("features", []) - if features is None: - features = [] - if not isinstance(features, list): - raise TypeError("features must be a list") - - feature_map = {} - for feature_entry in features: - if not isinstance(feature_entry, dict): - raise TypeError("each feature entry must be a dictionary") - feature_name = feature_entry.get("name") - if not isinstance(feature_name, str): - raise TypeError("each feature entry must include a string name") - feature_data = model.data.loc[model.data["feature_name"] == feature_name].copy() - if feature_data.empty: - feature_data = None - feature = model.create_and_add_foliation(feature_name, data=feature_data) - if feature is None: - raise ValueError(f"Could not recreate feature '{feature_name}' from recipe") - feature_map[feature_name] = feature - - for feature_entry in features: - feature_name = feature_entry.get("name") - fault_names = feature_entry.get("faults", []) - if not isinstance(fault_names, list): - raise TypeError("faults must be a list") - if fault_names: - feature = feature_map[feature_name] - feature.faults = [feature_map[name] for name in fault_names if name in feature_map] - - return model + return ModelSerializer.from_recipe_dict(cls, recipe) @public_api(tier="provisional") def to_recipe_json(self, data_reference=None, indent=2): @@ -304,8 +189,7 @@ def to_recipe_json(self, data_reference=None, indent=2): str JSON-formatted recipe string. """ - recipe = self.to_recipe_dict(data_reference=data_reference) - return json.dumps(recipe, indent=indent) + return ModelSerializer.to_recipe_json(self, data_reference=data_reference, indent=indent) @classmethod @public_api(tier="provisional") @@ -327,13 +211,7 @@ def from_recipe_json(cls, json_str): TypeError If json_str is not a string or does not parse as valid JSON. """ - if not isinstance(json_str, str): - raise TypeError("json_str must be a string") - try: - recipe = json.loads(json_str) - except json.JSONDecodeError as e: - raise TypeError(f"json_str is not valid JSON: {e}") - return cls.from_recipe_dict(recipe) + return ModelSerializer.from_recipe_json(cls, json_str) @public_api(tier="provisional") def save_recipe(self, filename, data_reference=None): @@ -347,11 +225,7 @@ def save_recipe(self, filename, data_reference=None): Path to an external CSV file to reference instead of embedding data inline in the JSON. If None, data is embedded. """ - filename = pathlib.Path(filename) - recipe = self.to_recipe_dict(data_reference=data_reference) - with open(filename, "w") as f: - json.dump(recipe, f, indent=2) - logger.info(f"Recipe saved to {filename}") + ModelSerializer.save_recipe(self, filename, data_reference=data_reference) @classmethod @public_api(tier="provisional") @@ -368,13 +242,7 @@ def load_recipe(cls, filename): GeologicalModel The reconstructed geological model. """ - filename = pathlib.Path(filename) - if not filename.exists(): - raise FileNotFoundError(f"Recipe file not found: {filename}") - with open(filename, "r") as f: - recipe = json.load(f) - logger.info(f"Recipe loaded from {filename}") - return cls.from_recipe_dict(recipe) + return ModelSerializer.load_recipe(cls, filename) def __str__(self): return f"GeologicalModel with {len(self.features)} features" @@ -530,44 +398,7 @@ def from_file(cls, file, allow_pickle: bool = True): GeologicalModel the geological model object """ - if not allow_pickle: - raise LoopValueError( - "Pickle-based loading is disabled (allow_pickle=False). " - f"Refusing to unpickle '{file}' because deserialising untrusted " - "pickle/dill data can execute arbitrary code. If you generated " - "this file yourself and trust its contents, call " - "GeologicalModel.from_file(file, allow_pickle=True). Otherwise, " - "use the JSON-based GeologicalModel.from_recipe_dict " - "(paired with GeologicalModel.to_recipe_dict) as a safe " - "alternative serialisation format." - ) - logger.warning( - f"Loading GeologicalModel from '{file}' using dill/pickle. " - "Only load model files from trusted sources: deserialising a " - "pickle file can execute arbitrary code. Pass allow_pickle=False " - "to refuse pickle-based loading, or use " - "GeologicalModel.from_recipe_dict for untrusted/JSON-based input." - ) - try: - import dill as pickle - except ImportError: - logger.error("Cannot import from file, dill not installed") - return None - path = pathlib.Path(file) - if not path.is_file(): - raise LoopValueError(f"Cannot load model, file does not exist: {file}") - try: - with open(path, "rb") as f: - model = pickle.load(f) - except Exception as e: - logger.error(f"Failed to load model from {file}: {e}") - raise LoopValueError(f"Failed to load model from {file}: {e}") from e - if isinstance(model, GeologicalModel): - logger.info("GeologicalModel initialised from file") - return model - else: - logger.error(f"{file} does not contain a geological model") - return None + return ModelSerializer.from_file(cls, file, allow_pickle=allow_pickle) def __getitem__(self, feature_name): """Accessor for feature in features using feature_name_index @@ -681,17 +512,7 @@ def to_file(self, file): file : string path to file location """ - try: - import dill as pickle - except ImportError: - logger.error("Cannot write to file, dill not installed \n" "pip install dill") - return - try: - logger.info(f"Writing GeologicalModel to: {file}") - with open(file, "wb") as handle: - pickle.dump(self, handle) - except pickle.PicklingError: - logger.error("Error saving file") + ModelSerializer.to_file(self, file) def _add_feature(self, feature, index: int | None = None): """