diff --git a/molecularnodes/bpyd/object.py b/molecularnodes/bpyd/object.py index e6f5cd26..92514343 100644 --- a/molecularnodes/bpyd/object.py +++ b/molecularnodes/bpyd/object.py @@ -95,7 +95,7 @@ def __init__(self, obj: Object | None): """ if not isinstance(obj, Object): raise ValueError(f"{obj} must be a Blender object of type Object") - self._object = obj + self._object_name = obj.name @property def object(self) -> Object | None: @@ -107,30 +107,8 @@ def object(self) -> Object | None: Object | None The Blender object, or None if not found. """ - # If we don't have connection to an object, attempt to re-stablish to a new - # object in the scene with the same UUID. This helps if duplicating / deleting - # objects in the scene, but sometimes Blender just loses reference to the object - # we are working with because we are manually setting the data on the mesh, - # which can wreak havoc on the object database. To protect against this, - # if we have a broken link we just attempt to find a new suitable object for it - try: - # if the connection is broken then trying to the name will raise a connection - # error. If we are loading from a saved session then the object_ref will be - # None and get an AttributeError - self._object.name - return self._object - except (ReferenceError, AttributeError): - for obj in bpy.data.objects: - if obj.mn.uuid == self.uuid: - print( - Warning( - f"Lost connection to object: {self._object}, now connected to {obj}" - ) - ) - self._object = obj - return obj - return None + return bpy.data.objects.get(self._object_name) @object.setter def object(self, value: Object) -> None: @@ -142,7 +120,7 @@ def object(self, value: Object) -> None: value : Object The Blender object to set. """ - self._object = value + self._object_name = value.name def store_named_attribute( self, diff --git a/molecularnodes/entities/entity.py b/molecularnodes/entities/entity.py index 0cd3f3de..3ffc2144 100644 --- a/molecularnodes/entities/entity.py +++ b/molecularnodes/entities/entity.py @@ -1,6 +1,7 @@ from abc import ABCMeta import bpy from uuid import uuid1 +from bpy.types import Object from ..bpyd import ( BlenderObject, ) @@ -13,8 +14,27 @@ class MolecularEntity( def __init__(self) -> None: self.uuid: str = str(uuid1()) self.type: str = "" - self._object: bpy.types.Object | None @property def bob(self) -> BlenderObject: return BlenderObject(self.object) + + @property + def object(self) -> Object: + try: + return bpy.data.objects[self._object_name] + except KeyError: + # if we can't find a refernce to the object via a name, then we look via the + # unique uuids that were assigned to the object and the entity to match up + for obj in bpy.data.objects: + if obj.mn.uuid == self.uuid: + self._object_name = obj.name + return obj + + @object.setter + def object(self, value) -> None: + if not isinstance(value, Object): + raise ValueError( + f"Can only set object to be of type bpy.types.Object, not {type(value)=}" + ) + self._object_name = value.name diff --git a/molecularnodes/entities/molecule/molecule.py b/molecularnodes/entities/molecule/molecule.py index 79999f34..039c37f2 100644 --- a/molecularnodes/entities/molecule/molecule.py +++ b/molecularnodes/entities/molecule/molecule.py @@ -4,6 +4,7 @@ from abc import ABCMeta from pathlib import Path from typing import Optional, Tuple, Union +from bpy.types import Collection import biotite.structure as struc import bpy @@ -71,11 +72,21 @@ def __init__(self, file_path: Union[str, Path, io.BytesIO]): self._parse_filepath(file_path=file_path) self.file: str self.array: np.ndarray - self.frames: bpy.types.Collection | None = None - self.frames_name: str = "" + self._frames_collection_name: str = "" bpy.context.scene.MNSession.molecules[self.uuid] = self + @property + def frames(self) -> Collection: + return bpy.data.collections.get(self._frames_collection_name) + + @frames.setter + def frames(self, value) -> None: + if value is None: + self._frames_collection_name = None + else: + self._frames_collection_name = value.name + @classmethod def _read(self, file_path: Union[Path, io.BytesIO]): """ diff --git a/molecularnodes/entities/trajectory/trajectory.py b/molecularnodes/entities/trajectory/trajectory.py index bad87105..a728ec64 100644 --- a/molecularnodes/entities/trajectory/trajectory.py +++ b/molecularnodes/entities/trajectory/trajectory.py @@ -21,7 +21,7 @@ def __init__(self, universe: mda.Universe, world_scale: float = 0.01): self.calculations: Dict[str, Callable] = {} self.world_scale = world_scale self.frame_mapping: npt.NDArray[np.in64] | None = None - bpy.context.scene.MNSession.trajectories[self.uuid] = self + bpy.context.scene.MNSession.entities[self.uuid] = self def selection_from_ui(self, ui_item: TrajectorySelectionItem) -> Selection: self.selections[ui_item.name] = Selection( diff --git a/molecularnodes/session.py b/molecularnodes/session.py index b9718070..eab2148b 100644 --- a/molecularnodes/session.py +++ b/molecularnodes/session.py @@ -13,34 +13,6 @@ from .entities.trajectory.trajectory import Trajectory -def trim(dictionary: dict): - to_pop = [] - for name, item in dictionary.items(): - # currently there are problems with pickling the functions so we have to just - # clean up any calculations that are created on saving. Could potentially convert - # it to a string and back but that is likely a job for better implementations - if hasattr(item, "calculations"): - item.calculations = {} - try: - item.object = None - if hasattr(item, "frames"): - if isinstance(item.frames, bpy.types.Collection): - item.frames_name = item.frames.name - item.frames = None - - except ReferenceError as e: - to_pop.append(name) - print( - Warning( - f"Object reference for {item} broken, removing this item from the session: `{e}`" - ) - ) - - for name in to_pop: - dictionary.pop(name) - return dictionary - - def path_relative_to_blend_wd(filepath: str | Path) -> Path: "Get the path of something, relative to the working directory of the current .blend file" blend_working_directory = bpy.path.abspath("//") @@ -60,124 +32,130 @@ def make_paths_relative(trajectories: Dict[str, Trajectory]) -> None: traj.save_filepaths_on_object() -def trim_root_folder(filename): - "Remove one of the prefix folders from a filepath" - return os.sep.join(filename.split(os.sep)[1:]) +def find_matching_object(uuid): + for obj in bpy.data.objects: + if obj.mn.uuid == uuid: + return obj + + return None class MNSession: def __init__(self) -> None: - self.molecules: Dict[str, Molecule] = {} - self.trajectories: Dict[str, Trajectory] = {} - self.ensembles: Dict[str, Ensemble] = {} - - def items(self): - "Return UUID and item for all molecules, trajectories and ensembles being tracked." - return ( - list(self.molecules.items()) - + list(self.trajectories.items()) - + list(self.ensembles.items()) - ) - - def get_object(self, uuid: str) -> bpy.types.Object | None: - """ - Try and get an object from Blender's object database that matches the uuid given. + self.entities: Dict[str, Molecule | Trajectory | Ensemble] = {} - If nothing is be found to match, return None. - """ - for obj in bpy.data.objects: - try: - if obj.mn.uuid == uuid: - return obj - except Exception as e: - print(e) + @property + def molecules(self) -> dict: + return { + key: mol for key, mol in self.entities.items() if isinstance(mol, Molecule) + } - return None + @property + def trajectories(self) -> dict: + return { + key: traj + for key, traj in self.entities.items() + if isinstance(traj, Trajectory) + } - def remove(self, uuid: str) -> None: - "Remove the item from the list." - self.molecules.pop(uuid, None) - self.trajectories.pop(uuid, None) - self.ensembles.pop(uuid, None) + @property + def ensembles(self) -> dict: + return { + key: ens for key, ens in self.entities.items() if isinstance(ens, Ensemble) + } def get(self, uuid: str) -> Union[Molecule, Trajectory, Ensemble]: - for id, item in self.items(): - if item.uuid == uuid: - return item + return self.entities.get(uuid) - return None + def __repr__(self) -> str: + return f"MNSession with {len(self.molecules)} molecules, {len(self.trajectories)} trajectories and {len(self.ensembles)} ensembles." - @property - def n_items(self) -> int: - "The number of items being tracked by this session." - length = 0 + def __len__(self) -> int: + return len(self.entities) - for dic in [self.molecules, self.trajectories, self.ensembles]: - length += len(dic) - return length + def trim(self) -> None: + to_pop = [] + for name, item in self.entities.items(): + # currently there are problems with pickling the functions so we have to just + # clean up any calculations that are created on saving. Could potentially convert + # it to a string and back but that is likely a job for better implementations + if hasattr(item, "calculations"): + item.calculations = {} - def __repr__(self) -> str: - return f"MNSession with {len(self.molecules)} molecules, {len(self.trajectories)} trajectories and {len(self.ensembles)} ensembles." + if item.object is None: + to_pop.append(name) + + for name in to_pop: + self.entities.pop(name) def pickle(self, filepath) -> None: - pickle_path = self.stashpath(filepath) + path = Path(filepath) + self.trim() + if len(self) == 0: + return None make_paths_relative(self.trajectories) - self.molecules = trim(self.molecules) - self.trajectories = trim(self.trajectories) - self.ensembles = trim(self.ensembles) # don't save anything if there is nothing to save - if self.n_items == 0: + if len(self) == 0: + # if we aren't saving anything, remove the currently existing session file + # so that it isn't reloaded when we load the save with old session information + if path.exists() and path.suffix == ".MNSession": + os.remove(filepath) return None - with open(pickle_path, "wb") as f: + with open(filepath, "wb") as f: pk.dump(self, f) - print(f"Saved session to: {pickle_path}") + print(f"Saved MNSession to: {filepath}") def load(self, filepath) -> None: - pickle_path = self.stashpath(filepath) - if not os.path.exists(pickle_path): - raise FileNotFoundError(f"MNSession file `{pickle_path}` not found") - with open(pickle_path, "rb") as f: - session = pk.load(f) - - for uuid, item in session.items(): - item.object = bpy.data.objects[item.name] - if hasattr(item, "frames") and hasattr(item, "frames_name"): - item.frames = bpy.data.collections[item.frames_name] - - for uuid, mol in session.molecules.items(): - self.molecules[uuid] = mol + "Load all of the entities from a previously saved MNSession" + path = Path(filepath) - for uuid, uni in session.trajectories.items(): - self.trajectories[uuid] = uni + if not path.exists(): + raise FileNotFoundError(f"MNSession file `{path}` not found") - for uuid, ens in session.ensembles.items(): - self.ensembles[uuid] = ens + with open(path, "rb") as f: + loaded_session = pk.load(f) + current_session = bpy.context.scene.MNSession - print(f"Loaded a MNSession from: {pickle_path}") + # merge the loaded session with current session, handling if they used the old + # structure of separating entities into different categories + if hasattr(loaded_session, "entities"): + current_session.entites + loaded_session.entities + else: + items = [] + for attr in ["molecules", "trajectories", "ensembles"]: + try: + items.append((getattr(loaded_session, attr))) + except AttributeError: + pass + print(f"{items=}") + for key, item in items: + current_session.entities[key] = item + + print(f"Loaded a MNSession from: {filepath}") def stashpath(self, filepath) -> str: return f"{filepath}.MNSession" def clear(self) -> None: """Remove references to all molecules, trajectories and ensembles.""" - self.molecules.clear() - self.trajectories.clear() - self.ensembles.clear() + self.entities.clear() def get_session(context: Context | None = None) -> MNSession: - if not context: - context = bpy.context - return context.scene.MNSession + if isinstance(context, Context): + return context.scene.MNSession + else: + return bpy.context.scene.MNSession @persistent def _pickle(filepath) -> None: - get_session().pickle(filepath) + session = get_session() + session.pickle(session.stashpath(filepath)) @persistent @@ -210,7 +188,8 @@ def _load(filepath: str, printing: str = "quiet") -> None: if filepath == "": return None try: - get_session().load(filepath) + session = get_session() + session.load(session.stashpath(filepath)) except FileNotFoundError: if printing == "verbose": print("No MNSession found to load for this .blend file.") diff --git a/molecularnodes/ui/panel.py b/molecularnodes/ui/panel.py index a04db68f..0a4e5f0a 100644 --- a/molecularnodes/ui/panel.py +++ b/molecularnodes/ui/panel.py @@ -237,8 +237,6 @@ def item_ui(layout, item): def panel_session(layout, context): session = get_session(context) - # if session.n_items > 0: - # return None row = layout.row() row.label(text="Loaded items in the session") # row.operator("mn.session_reload") diff --git a/tests/test_trajectory.py b/tests/test_trajectory.py index b1320dbd..8b533c90 100644 --- a/tests/test_trajectory.py +++ b/tests/test_trajectory.py @@ -200,7 +200,8 @@ def test_save_persistance( assert os.path.exists(session.stashpath(filepath)) bpy.ops.wm.open_mainfile(filepath=filepath) - traj = mn.session.get_session().trajectories[uuid] + traj = mn.session.get_session().get(uuid) + assert traj is not None verts_frame_0 = traj.named_attribute("position") bpy.context.scene.frame_set(4) verts_frame_4 = traj.named_attribute("position")