cleanup collections code

molecularnodes/blender/bpyd/collection.py

@@ -0,0 +1,23 @@
+import bpy
+from bpy.types import Collection
+
+
+def create_collection(
+    name: str = "NewCollection", parent: Collection | str | None = None
+) -> Collection:
+    if isinstance(parent, str):
+        try:
+            parent = bpy.data.collections[name]
+        except KeyError:
+            parent = bpy.data.collections.new(name)
+            bpy.context.scene.collection.children.linke(parent)
+    try:
+        coll = bpy.data.collections[name]
+    except KeyError:
+        coll = bpy.data.collections.new(name)
+        if parent is None:
+            bpy.context.scene.collection.children.link(coll)
+        else:
+            parent.children.link(coll)
+
+    return coll

molecularnodes/blender/coll.py

@@ -1,70 +1,34 @@
 import bpy
+from bpy.types import Collection
+from .bpyd.collection import create_collection
 
 
-def mn():
-    """Return the MolecularNodes Collection
+def mn() -> Collection:
+    "Return the 'MolecularNodes' collection, creating it first if required"
+    return create_collection("MolecularNodes")
 
-    The collection called 'MolecularNodes' inside the Blender scene is returned. If the 
-    collection does not exist first, it is created.
-    """
-    coll = bpy.data.collections.get('MolecularNodes')
-    if not coll:
-        coll = bpy.data.collections.new('MolecularNodes')
-        bpy.context.scene.collection.children.link(coll)
-    return coll
 
+def data() -> Collection:
+    "Return the MolecularNodes/data collection and disable it"
+    name = ".MN_data"
 
-def armature(name='MN_armature'):
-    coll = bpy.data.collections.get(name)
-    if not coll:
-        coll = bpy.data.collections.new(name)
-        mn().children.link(coll)
-    return coll
+    try:
+        return bpy.data.collections[name]
+    except KeyError:
+        collection = create_collection(name=name, parent=mn())
 
+        bpy.context.view_layer.layer_collection.children["MolecularNodes"].children[
+            collection.name
+        ].exclude = True
+        return collection
 
-def data(suffix=""):
-    """A collection for storing MN related data objects.
-    """
-    name = f"MN_data{suffix}"
 
-    collection = bpy.data.collections.get(name)
-    if not collection:
-        collection = bpy.data.collections.new(name)
-        mn().children.link(collection)
+def frames(name: str = "") -> Collection:
+    "Return a collection for storing the objects that are the frames of a trajectory"
+    return create_collection(f".data_{name}_frames", parent=data())
 
-        # disable the view of the data collection
-        bpy.context.view_layer.layer_collection.children['MolecularNodes'].children[name].exclude = True
-    return collection
 
-
-def frames(name="", parent=None, suffix="_frames"):
-    """Create a Collection for Frames of a Trajectory
-
-    Args:
-        name (str, optional): Name of the collection for the frames. Defaults to "".
-        parent (_type_, optional): A blender collection which will become the parent 
-        collection. Defaults to the MolecularNodes collection if None.
-    """
-    coll_frames = bpy.data.collections.new(name + suffix)
-    if not parent:
-        mn().children.link(coll_frames)
-    else:
-        parent.children.link(coll_frames)
-
-    return coll_frames
-
-
-def cellpack(name="", parent=None, fallback=False):
+def cellpack(name: str = "") -> Collection:
+    "Return a collection for storing the instances for a CellPack Ensemble"
     full_name = f"cellpack_{name}"
-    coll = bpy.data.collections.get(full_name)
-    if coll and fallback:
-        return coll
-
-    coll = bpy.data.collections.new(full_name)
-
-    if parent:
-        parent.children.link(coll)
-    else:
-        data().children.link(coll)
-
-    return coll
+    return create_collection(full_name, parent=data())

molecularnodes/blender/mesh.py

@@ -76,7 +76,12 @@ def evaluate_using_mesh(obj):
     return bob.evaluate().object
 
 
-def create_data_object(array, collection=None, name="DataObject", world_scale=0.01):
+def create_data_object(
+    array: np.ndarray,
+    name: str = "DataObject",
+    collection: str | bpy.types.Collection | None = None,
+    world_scale: float = 0.01,
+) -> bpy.types.Object:
     # still requires a unique call TODO: figure out why
     # I think this has to do with the bcif instancing extraction
     array = np.unique(array)

molecularnodes/entities/molecule/molecule.py

@@ -647,19 +647,13 @@ def att_sec_struct():
 
     coll_frames = None
     if frames:
-        coll_frames = bl.coll.frames(bob.name, parent=bl.coll.data())
+        coll_frames = bl.coll.frames(bob.name)
         for i, frame in enumerate(frames):
             frame = bl.mesh.create_object(
                 name=bob.name + "_frame_" + str(i),
                 collection=coll_frames,
                 vertices=frame.coord * world_scale,
-                # vertices=frame.coord * world_scale - centroid
             )
-            # TODO if update_attribute
-            # bl.mesh.store_named_attribute(attribute)
-
-    # this has started to throw errors for me. I'm not sure why.
-    # mol.mn['molcule_type'] = 'pdb'
 
     # add custom properties to the actual blender object, such as number of chains, biological assemblies etc
     # currently biological assemblies can be problematic to holding off on doing that

molecularnodes/entities/trajectory/dna.py

@@ -288,7 +288,7 @@ def load(top, traj, name="oxDNA", setup_nodes=True, world_scale=0.01):
     # create a collection to store all of the frame objects that are part of the trajectory
     # they will contain all of the possible attributes which can be interpolated betewen
     # frames such as position, base_vector, base_normal, velocity, angular_velocity
-    collection = coll.frames(f"{name}_frames", parent=coll.data())
+    collection = coll.frames(name)
     for i, frame in enumerate(trajectory):
         fill_n = int(np.ceil(np.log10(n_frames)))
         frame_name = f"{name}_frame_{str(i).zfill(fill_n)}"

tests/test_coll.py

@@ -4,9 +4,7 @@
 def test_coll():
     assert mn.blender.coll.mn().name == "MolecularNodes"
     assert mn.blender.coll.mn().name == "MolecularNodes"
-    assert mn.blender.coll.data().name == "MN_data"
+    assert mn.blender.coll.data().name == ".MN_data"
     assert mn.blender.coll.cellpack().name == "cellpack_"
-    assert mn.blender.coll.cellpack(fallback=True).name == "cellpack_"
-    assert mn.blender.coll.cellpack().name == "cellpack_.001"
-    assert mn.blender.coll.frames().name == "_frames"
-    assert mn.blender.coll.frames("4OZS").name == "4OZS_frames"
+    assert mn.blender.coll.frames().name == ".data__frames"
+    assert mn.blender.coll.frames("4OZS").name == ".data_4OZS_frames"