more refactoring

molecularnodes/blender/attribute.py

@@ -1,11 +1,9 @@
 from dataclasses import dataclass
 from enum import Enum
 from typing import Type
-from functools import reduce
 
 import bpy
 import numpy as np
-from numpy.ma.core import prod
 
 
 @dataclass
@@ -121,6 +119,9 @@ def guess_atype_from_array(array: np.ndarray) -> AttributeType:
     return AttributeTypes.FLOAT.value
 
 class Attribute:
+    """
+    Wrapper around a Blender attribute to provide a more convenient interface with numpy arrays
+    """
     def __init__(self, attribute: bpy.types.Attribute):
         self.attribute = attribute
         self.n_attr = len(attribute.data)
@@ -158,6 +159,17 @@ def dtype(self) -> Type:
     def n_values(self) -> int:
         return np.prod(self.shape, dtype=int)
 
+    def from_array(self, array: np.ndarray) -> None:
+        """
+        Set the attribute data from a numpy array
+        """
+        if array.shape != self.shape:
+            raise ValueError(
+                f"Array shape {array.shape} does not match attribute shape {self.shape}"
+            )
+
+        self.attribute.data.foreach_set(self.value_name, array.reshape(-1))
+
 
     def as_array(self) -> np.ndarray:
             """

molecularnodes/blender/mesh.py

@@ -10,7 +10,7 @@
     AttributeTypes,
     guess_atype_from_array,
 )
-from .object import ObjectTracker, create_object
+from .object import ObjectTracker, create_object, evaluate_object
 
 
 def centre(array: np.ndarray):
@@ -96,7 +96,7 @@ def named_attribute(
     obj: bpy.types.Object, name="position", evaluate=False
 ) -> np.ndarray:
     """
-    Get the attribute data from the object.
+    Get the named attribute data from the object, optionally evaluating modifiers first.
 
     Parameters:
         object (bpy.types.Object): The Blender object.
@@ -107,20 +107,17 @@ def named_attribute(
     """
     if evaluate:
         obj = evaluate_object(obj)
-    attribute_names = obj.data.attributes.keys()
     verbose = False
-    if name not in attribute_names:
+    try:
+        attr = Attribute(obj.data.attributes[name])
+    except KeyError:
+        message = f"The selected attribute '{name}' does not exist on the mesh."
         if verbose:
-            raise AttributeError(
-                f"The selected attribute '{name}' does not exist on the mesh. \
-                Possible attributes are: {attribute_names=}"
-            )
-        else:
-            raise AttributeError(
-                f"The selected attribute '{name}' does not exist on the mesh."
-            )
+            message += f"Possible attributes are: {obj.data.attributes.keys()}"
+
+        raise AttributeError(message)
 
-    return Attribute(obj.data.attributes[name]).as_array()
+    return attr.as_array()
 
 
 def import_vdb(file: str, collection: bpy.types.Collection = None) -> bpy.types.Object:
@@ -153,10 +150,6 @@ def import_vdb(file: str, collection: bpy.types.Collection = None) -> bpy.types.
     return obj
 
 
-def evaluate_object(obj):
-    "Return an object which has the modifiers evaluated."
-    obj.update_tag()
-    return obj.evaluated_get(bpy.context.evaluated_depsgraph_get())
 
 
 def evaluate_using_mesh(obj):

molecularnodes/blender/object.py

@@ -63,6 +63,11 @@ def latest(self):
         return self.new_objects()[-1]
 
 
+def evaluate_object(obj):
+    "Return an object which has the modifiers evaluated."
+    obj.update_tag()
+    return obj.evaluated_get(bpy.context.evaluated_depsgraph_get())
+
 def create_object(
     vertices: np.ndarray = [],
     edges: np.ndarray = [],