working refactor

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,

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

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]):
         """

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(

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.")

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")

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")