more fixing

molecularnodes/blender/bones.py

@@ -105,7 +105,7 @@ class MN_MT_Add_Armature(bpy.types.Operator):
     bl_label = "Add Armature"
     bl_description = "Automatically add armature for each amino acid of the structure   "
 
-    def execute(self, context):
+    def execute(self, context: bpy.types.Context):
         object = context.active_object
         add_bones(bpy.data.objects[object.name], name=object.name)
 

molecularnodes/color.py

@@ -2,7 +2,7 @@
 import colorsys
 import numpy as np
 from numpy.typing import NDArray
-from typing import List, Dict, Tuple
+from typing import List, Dict, Tuple, Any
 
 
 def random_rgb(seed: int = 6) -> NDArray[np.float64]:
@@ -45,7 +45,7 @@ def equidistant_colors(
 
 
 def color_chains_equidistant(
-    chain_ids: NDArray[np.character],
+    chain_ids: NDArray[Any],
 ) -> NDArray[np.float32]:
     color_dict = equidistant_colors(chain_ids)
     chain_colors = np.array([color_dict[x] for x in chain_ids])

molecularnodes/io/cellpack.py

@@ -36,10 +36,10 @@ class MN_OT_Import_Cell_Pack(bpy.types.Operator):
     bl_options = {"REGISTER"}
 
     @classmethod
-    def poll(cls, context):
+    def poll(cls, context: bpy.types.Context):
         return True
 
-    def execute(self, context):
+    def execute(self, context: bpy.types.Context):
         s = context.scene
         load(
             file_path=s.mol_import_cell_pack_path,

molecularnodes/io/density.py

@@ -64,10 +64,10 @@ class MN_OT_Import_Map(bpy.types.Operator):
     bl_options = {"REGISTER"}
 
     @classmethod
-    def poll(cls, context):
+    def poll(cls, context: bpy.types.Context):
         return True
 
-    def execute(self, context):
+    def execute(self, context: bpy.types.Context):
         scene = context.scene
         load(
             file_path=scene.MN_import_density,

molecularnodes/io/dna.py

@@ -1,5 +1,8 @@
 import bpy
 import numpy as np
+from numpy.typing import NDArray
+from typing import Union, Tuple, Set
+from pathlib import Path
 
 from .. import color
 from ..blender import coll, nodes, obj
@@ -24,7 +27,7 @@
 )
 
 
-def base_to_int(bases: np.array) -> np.array:
+def base_to_int(bases: NDArray[np.character]) -> NDArray[np.int32]:
     """
     Convert an array of DNA bases to their corresponding MN integer values.
 
@@ -46,20 +49,20 @@ def base_to_int(bases: np.array) -> np.array:
     return ints
 
 
-def is_new_topology(filepath):
+def is_new_topology(filepath: Union[str, Path]) -> bool:
     with open(filepath) as f:
         firstline = f.readline()
 
     return "5 -> 3" in firstline
 
 
-def read_topology_new(filepath):
+def read_topology_new(filepath: Union[str, Path]) -> NDArray[np.int32]:
     with open(filepath, "r") as file:
         contents = file.read()
 
     lines = np.array(contents.split("\n"))
 
-    def read_seq_line(line):
+    def read_seq_line(line: str) -> NDArray[np.character]:
         sequence = line.split(" ")[0]
         return np.array([c for c in sequence])
 
@@ -85,7 +88,7 @@ def read_seq_line(line):
     return np.vstack(strands)
 
 
-def read_topology_old(filepath):
+def read_topology_old(filepath: Union[str, Path]) -> NDArray[np.int32]:
     """
     Read the topology from a file and convert it to a numpy array.
 
@@ -131,7 +134,7 @@ def read_topology_old(filepath):
     return array_int
 
 
-def read_trajectory(filepath):
+def read_trajectory(filepath: Union[str, Path]) -> NDArray[np.int32]:
     """
     Read an oxDNA trajectory file and return an array of frames.
 
@@ -175,7 +178,7 @@ def read_trajectory(filepath):
     return np.stack(frames)
 
 
-def set_attributes_to_dna_mol(mol, frame, scale_dna=0.1):
+def set_attributes_to_dna_mol(mol: bpy.types.Object, frame: NDArray[np.float64], scale_dna: float = 0.1) -> None:
     attributes = ("base_vector", "base_normal", "velocity", "angular_velocity")
     for i, att in enumerate(attributes):
         col_idx = np.array([3, 4, 5]) + i * 3
@@ -192,7 +195,7 @@ def set_attributes_to_dna_mol(mol, frame, scale_dna=0.1):
         obj.set_attribute(mol, att, data, type="FLOAT_VECTOR")
 
 
-def toplogy_to_bond_idx_pairs(topology: np.ndarray):
+def toplogy_to_bond_idx_pairs(topology: NDArray[np.int32]) -> NDArray[np.int32]:
     """
     Convert the given topology array into pairs of indices representing each distinct bond.
 
@@ -233,7 +236,13 @@ def toplogy_to_bond_idx_pairs(topology: np.ndarray):
     return np.sort(bond_idxs, axis=1)
 
 
-def load(top, traj, name="oxDNA", setup_nodes=True, world_scale=0.01):
+def load(
+    top: Union[str, Path],
+    traj: Union[str, Path],
+    name: str = "oxDNA",
+    setup_nodes: bool = True,
+    world_scale: float = 0.01,
+) -> Tuple[bpy.types.Object, bpy.types.Collection]:
     # the scale of the oxDNA files seems to be based on nanometres rather than angstrongs
     # like most structural biology files, so currently adjusting the world_scale to
     # compensate
@@ -293,13 +302,13 @@ def load(top, traj, name="oxDNA", setup_nodes=True, world_scale=0.01):
     return mol, collection
 
 
-class MN_OT_Import_OxDNA_Trajectory(bpy.types.Operator):
+class MN_OT_Import_OxDNA_Trajectory(bpy.types.Operator):  # type: ignore
     bl_idname = "mn.import_oxdna"
     bl_label = "Load"
     bl_description = "Will import the given file and toplogy."
     bl_options = {"REGISTER"}
 
-    def execute(self, context):
+    def execute(self, context: bpy.types.Context) -> Set[str]:
         s = context.scene
         load(
             top=s.MN_import_oxdna_topology,
@@ -309,7 +318,7 @@ def execute(self, context):
         return {"FINISHED"}
 
 
-def panel(layout, scene):
+def panel(layout: bpy.types.UIList, scene: bpy.types.Scene) -> bpy.types.UIList:
     layout.label(text="Load oxDNA File", icon="FILE_TICK")
     layout.separator()
     row = layout.row()

molecularnodes/io/local.py

@@ -66,7 +66,7 @@ class MN_OT_Import_Protein_Local(bpy.types.Operator):
     bl_description = "Open a local structure file"
     bl_options = {"REGISTER", "UNDO"}
 
-    def execute(self, context):
+    def execute(self, context: bpy.types.Context):
         scene = context.scene
         file_path = scene.MN_import_local_path
 

molecularnodes/io/md.py

@@ -107,10 +107,10 @@ class MN_OT_Import_Protein_MD(bpy.types.Operator):
     bl_options = {"REGISTER", "UNDO"}
 
     @classmethod
-    def poll(cls, context):
+    def poll(cls, context: bpy.types.Context):
         return True
 
-    def execute(self, context):
+    def execute(self, context: bpy.types.Context):
         scene = context.scene
         if not pkg.is_current("MDAnalysis"):
             self.report(
@@ -189,7 +189,7 @@ class TrajectorySelection_OT_NewItem(bpy.types.Operator):
     bl_idname = "trajectory_selection_list.new_item"
     bl_label = "+"
 
-    def execute(self, context):
+    def execute(self, context: bpy.types.Context):
         context.scene.trajectory_selection_list.add()
         return {"FINISHED"}
 
@@ -199,10 +199,10 @@ class TrajectorySelection_OT_DeleteIem(bpy.types.Operator):
     bl_label = "-"
 
     @classmethod
-    def poll(cls, context):
+    def poll(cls, context: bpy.types.Context):
         return context.scene.trajectory_selection_list
 
-    def execute(self, context):
+    def execute(self, context: bpy.types.Context):
         my_list = context.scene.trajectory_selection_list
         index = context.scene.list_index
 

molecularnodes/io/parse/cellpack.py

@@ -1,5 +1,6 @@
 from pathlib import Path
-
+from typing import Union
+from pathlib import Path
 import numpy as np
 import bpy
 
@@ -12,7 +13,7 @@
 
 
 class CellPack(Ensemble):
-    def __init__(self, file_path):
+    def __init__(self, file_path: Union[str, Path]) -> None:
         super().__init__(file_path)
         self.file_type = self._file_type()
         self.data = self._read(self.file_path)
@@ -21,12 +22,8 @@ def __init__(self, file_path):
         self.chain_ids = self.data.chain_ids
 
     def create_model(
-        self,
-        name="CellPack",
-        node_setup: bool = True,
-        world_scale: float = 0.01,
-        fraction: float = 1.0,
-    ):
+        self, name: str = "StarFileObject", node_setup: bool = True, world_scale: float = 0.01
+    ) -> bpy.types.Object:
         self.data_object = self._create_data_object(name=f"{name}")
         self._create_object_instances(name=name, node_setup=node_setup)
 

molecularnodes/io/parse/ensemble.py

@@ -1,12 +1,14 @@
 import bpy
 from abc import ABCMeta
 import numpy as np
+from typing import Union, Optional, Set, List, Dict
+from pathlib import Path
 from ... import blender as bl
 import warnings
 
 
 class Ensemble(metaclass=ABCMeta):
-    def __init__(self, file_path):
+    def __init__(self, file_path: Union[str, Path]) -> None:
         """
         Initialize an Ensemble object.
 
@@ -17,20 +19,15 @@ def __init__(self, file_path):
 
         """
         self.type: str = "ensemble"
-        self.file_path: str = file_path
+        self.file_path: str = str(file_path)
         self.object: bpy.types.Object = None
         self.instances: bpy.types.Collection = None
         self.frames: bpy.types.Collection = None
 
     @classmethod
     def create_model(
-        cls,
-        name: str = "NewEnsemble",
-        node_setup: bool = True,
-        world_scale: float = 0.01,
-        fraction: float = 1.0,
-        simplify=False,
-    ):
+        self, name: str = "EnsembleObject", node_setup: bool = True, world_scale: float = 0.01
+    ) -> bpy.types.Object:
         """
         Create a 3D model in the of the ensemble.
 

molecularnodes/io/parse/molecule.py

@@ -1,5 +1,6 @@
 from abc import ABCMeta
 from typing import Optional, Any, Tuple, Union, List, Dict
+import biotite.database
 import biotite.structure
 from numpy.typing import NDArray
 from pathlib import Path
@@ -54,8 +55,8 @@ class Molecule(metaclass=ABCMeta):
     """
 
     def __init__(self, file_path: str) -> None:
-        self.file_path: Optional[Union[Path, str]] = None
-        self.file = None
+        self.file_path: Union[Path, str]
+        self.file: biotite.file.TextFile
         self.object: Optional[bpy.types.Object] = None
         self.frames: Optional[bpy.types.Collection] = None
         self.array: Optional[np.ndarray] = None
@@ -440,7 +441,7 @@ def att_res_name() -> NDArray[np.int32]:
         mol["ligands"] = np.unique(other_res)
         return np.array(res_nums)
 
-    def att_chain_id():
+    def att_chain_id() -> NDArray[np.int32]:
         return np.unique(array.chain_id, return_inverse=True)[1]
 
     def att_entity_id():