cleanup

molecularnodes/blender/obj.py

@@ -1,7 +1,7 @@
 from dataclasses import dataclass
-from typing import Union, List, Optional
-from typing import Type
+from typing import Union, List, Optional, Type
 from types import TracebackType
+from pathlib import Path
 import bpy
 import numpy as np
 
@@ -169,6 +169,13 @@ def create_object(
     return object
 
 
+def set_position(bob: bpy.types.Object, positions: np.ndarray) -> None:
+    "A stripped-back way to set the positions for higher performance."
+    attribute = bob.data.attributes["position"]
+    attribute.data.foreach_set("vector", positions.reshape(-1))
+    bob.data.vertices[0].co = bob.data.vertices[0].co
+
+
 def set_attribute(
     object: bpy.types.Object,
     name: str,
@@ -246,7 +253,7 @@ def set_attribute(
     # is the case For now we will set a single vert to it's own position, which triggers a
     # proper refresh of the object data.
     try:
-        object.data.vertices[0].co = object.data.certices[0].co
+        object.data.vertices[0].co = object.data.vertices[0].co
     except AttributeError:
         object.data.update()
 
@@ -300,7 +307,7 @@ def get_attribute(object: bpy.types.Object, name: str = "position", evaluate: bo
         return array
 
 
-def import_vdb(file: str, collection: bpy.types.Collection = None) -> bpy.types.Object:
+def import_vdb(file: Union[Path, str], collection: bpy.types.Collection = None) -> bpy.types.Object:
     """
     Imports a VDB file as a Blender volume object, in the MolecularNodes collection.
 

molecularnodes/color.py

@@ -1,8 +1,9 @@
 import random
 import colorsys
 import numpy as np
+import numpy.typing as npt
 from numpy.typing import NDArray
-from typing import List, Dict, Tuple, Any
+from typing import List, Dict, Tuple, Any, Iterable
 
 
 def random_rgb(seed: int = 6) -> NDArray[np.float64]:
@@ -18,7 +19,7 @@ def color_from_atomic_number(atomic_number: int) -> Tuple[int, int, int, int]:
 
 
 def colors_from_elements(
-    atomic_numbers: NDArray[np.int32],
+    atomic_numbers: Iterable,
 ) -> NDArray[np.float64]:
     colors = np.array([color_from_atomic_number(x) for x in atomic_numbers])
     return colors

molecularnodes/io/parse/density.py

@@ -2,21 +2,24 @@
 import os
 import bpy
 
+from typing import Union, Optional
+from pathlib import Path
+
 
 class Density(metaclass=ABCMeta):
     """
     Abstract base class for molecular density objects.
 
     """
 
-    def __init__(self, file_path):
-        self.file_path: str = None
+    def __init__(self, file_path: Union[str, Path]) -> None:
+        self.file_path = file_path
         self.grid = None
-        self.file_vdb: str = None
-        self.threshold: float = None
-        self.object: bpy.types.Object = None
+        self.file_vdb: Union[Path, str]
+        self.threshold: float = 1.0
+        self.object: Optional[bpy.types.Object] = None
 
-    def path_to_vdb(self, file: str, center: False, invert: False):
+    def path_to_vdb(self, file: Union[Path, str], center: False, invert: False) -> Path:
         """
         Convert a file path to a corresponding VDB file path.
 
@@ -37,4 +40,4 @@ def path_to_vdb(self, file: str, center: False, invert: False):
         name += "_invert" if invert else ""
         file_name = name + ".vdb"
         file_path = os.path.join(folder_path, file_name)
-        return file_path
+        return Path(file_path)

molecularnodes/io/parse/mda.py

@@ -1,27 +1,10 @@
 import bpy
 from bpy.app.handlers import persistent
 
-try:
-    import MDAnalysis as mda
-except ImportError:
-    HAS_mda = False
-    import types
-
-    class MockAtomGroup:
-        pass
-
-    class MockUniverse:
-        pass
-
-    mda = types.ModuleType("MDAnalysis")
-    mda.Universe = MockUniverse
-    mda.AtomGroup = MockAtomGroup
-    mda.core = types.ModuleType("core")
-    mda.topology = types.ModuleType("topology")
-
-else:
-    HAS_mda = True
+
+import MDAnalysis as mda
 import numpy as np
+from numpy.typing import NDArray
 import warnings
 import pickle
 from typing import Union, List, Dict
@@ -103,31 +86,29 @@ def __init__(self, ag: mda.AtomGroup, style: str = "vdw", world_scale: float = 0
         is_solvent : np.ndarray
             Whether the atoms in the atomgroup are solvent.
         """
-        if not HAS_mda:
-            raise ImportError("MDAnalysis is not installed.")
         self.ag = ag
         self.world_scale = world_scale
         self.style = style
 
     @property
     def n_atoms(self) -> int:
-        return self.ag.n_atoms
+        return int(self.ag.n_atoms)
 
     @property
     def style(self) -> str:
-        return self._style
+        return str(self._style)
 
     @style.setter
-    def style(self, style):
+    def style(self, style: str) -> None:
         self._style = style
 
     @staticmethod
-    def bool_selection(ag, selection) -> np.ndarray:
+    def bool_selection(ag: mda.AtomGroup, selection: NDArray[np.bool_]) -> NDArray[np.bool_]:
         return np.isin(ag.ix, ag.select_atoms(selection).ix).astype(bool)
 
     @property
-    def positions(self) -> np.ndarray:
-        return self.ag.positions * self.world_scale
+    def positions(self) -> NDArray[np.float32]:
+        return self.ag.positions * self.world_scale  # type: ignore
 
     @property
     def bonds(self) -> List[List[int]]:
@@ -159,7 +140,7 @@ def elements(self) -> List[str]:
                 ValueError,
             ):  # If 'x' is not in 'data.elements.keys()' or 'guess_atom_element(x)' fails
                 elements = ["X"] * self.ag.n_atoms
-        return elements
+        return elements  # type: ignore
 
     @property
     def atomic_number(self) -> np.ndarray:
@@ -419,8 +400,6 @@ def __init__(self, world_scale: float = 0.01, in_memory: bool = False):
             Whether the old import is used (default: False).
         """
         log = start_logging(logfile_name="mda")
-        if not HAS_mda:
-            raise ImportError("MDAnalysis is not installed.")
 
         # if the session already exists, load the existing session
         if hasattr(bpy.types.Scene, "mda_session"):

molecularnodes/io/parse/molecule.py

@@ -713,7 +713,7 @@ def att_sec_struct() -> NDArray[np.int32]:
             )
             if verbose:
                 print(f'Added {att["name"]} after {time.process_time() - start} s')
-        except ValueError:
+        except AttributeError:
             if verbose:
                 warnings.warn(f"Unable to add attribute: {att['name']}")
                 print(f'Failed adding {att["name"]} after {time.process_time() - start} s')

molecularnodes/io/parse/mrc.py

@@ -1,8 +1,10 @@
 import os
-
+import mrcfile
+import pyopenvdb as vdb
 import bpy
 import numpy as np
-
+from typing import Union
+from pathlib import Path
 from ...blender import coll, nodes, obj
 from .density import Density
 
@@ -15,13 +17,16 @@ class MRC(Density):
     that can be written as `.vdb` files and the imported into Blender as volumetric objects.
     """
 
-    def __init__(self, file_path, center=False, invert=False, overwrite=False):
-        super().__init__(self)
-        self.file_path = file_path
+    def __init__(
+        self, file_path: Union[str, Path], center: bool = False, invert: bool = False, overwrite: bool = False
+    ) -> None:
+        super().__init__(file_path=file_path)
         self.grid = self.map_to_grid(self.file_path, center=center)
         self.file_vdb = self.map_to_vdb(self.file_path, center=center, invert=invert, overwrite=overwrite)
 
-    def create_model(self, name="NewDensity", style="density_surface", setup_nodes=True) -> bpy.types.Object:
+    def create_model(
+        self, name: str = "NewDensity", style: str = "density_surface", setup_nodes: bool = True
+    ) -> bpy.types.Object:
         """
         Loads an MRC file into Blender as a volumetric object.
 
@@ -55,12 +60,12 @@ def create_model(self, name="NewDensity", style="density_surface", setup_nodes=T
 
     def map_to_vdb(
         self,
-        file: str,
+        file: Union[str, Path],
         invert: bool = False,
-        world_scale=0.01,
+        world_scale: float = 0.01,
         center: bool = False,
-        overwrite=False,
-    ) -> (str, float):
+        overwrite: bool = False,
+    ) -> Path:
         """
         Converts an MRC file to a .vdb file using pyopenvdb.
 
@@ -124,7 +129,7 @@ def map_to_vdb(
         # Return the path to the output file
         return file_path
 
-    def map_to_grid(self, file: str, invert: bool = False, center: bool = False):
+    def map_to_grid(self, file: Union[str, Path], invert: bool = False, center: bool = False) -> vdb.GridBase:
         """
         Reads an MRC file and converts it into a pyopenvdb FloatGrid object.
 
@@ -144,8 +149,6 @@ def map_to_grid(self, file: str, invert: bool = False, center: bool = False):
         pyopenvdb.FloatGrid
             A pyopenvdb FloatGrid object containing the density data.
         """
-        import mrcfile
-        import pyopenvdb as vdb
 
         volume = mrcfile.read(file)
 

pyproject.toml

@@ -27,6 +27,7 @@ syrupy = "*"
 quartodoc = "*"
 scipy = "*"
 mypy = "*"
+ruff = "*"
 
 
 [build-system]
@@ -38,6 +39,8 @@ build-backend = "poetry.core.masonry.api"
 [tool.mypy]
 strict = true
 ignore_missing_imports = true
+plugins = "numpy.typing.mypy_plugin"
+
 
 [tool.ruff]
 # Enable pycodestyle (`E`) and Pyflakes (`F`) codes by default.