WIP: Modifier Interface Helper

databpy/__init__.py

@@ -1,44 +1,46 @@
-from .object import (
-    ObjectTracker,
-    BlenderObjectBase,
-    BlenderObjectAttribute,
-    BlenderObject,
-    BOB,
-    create_object,
-    create_bob,
-    create_mesh_object,
-    create_curves_object,
-    create_pointcloud_object,
-    LinkedObjectError,
-    bdo,
-)
-from .vdb import import_vdb
 from . import nodes
-from .nodes import utils
 from .addon import register, unregister
-from .utils import centre, lerp
-from .collection import create_collection
 from .array import AttributeArray
 from .attribute import (
-    named_attribute,
-    store_named_attribute,
-    remove_named_attribute,
-    list_attributes,
-    evaluate_object,
     Attribute,
+    AttributeDomain,
+    AttributeDomains,
+    AttributeMismatchError,
     AttributeType,
     AttributeTypeInfo,
     AttributeTypes,
-    AttributeDomains,
-    AttributeDomain,
     NamedAttributeError,
-    AttributeMismatchError,
+    evaluate_object,
+    list_attributes,
+    named_attribute,
+    remove_named_attribute,
+    store_named_attribute,
+)
+from .collection import create_collection
+from .modifier import NodesModifierInterface
+from .nodes import utils
+from .object import (
+    BOB,
+    BlenderObject,
+    BlenderObjectAttribute,
+    BlenderObjectBase,
+    LinkedObjectError,
+    ObjectTracker,
+    bdo,
+    create_bob,
+    create_curves_object,
+    create_mesh_object,
+    create_object,
+    create_pointcloud_object,
 )
+from .utils import centre, lerp
+from .vdb import import_vdb
 
 __all__ = [
     "ObjectTracker",
     "BlenderObjectBase",
     "BlenderObjectAttribute",
+    "NodesModifierInterface",
     "BlenderObject",
     "BOB",
     "create_object",

databpy/modifier.py

@@ -0,0 +1,128 @@
+import bpy
+
+from .object import BlenderObjectBase
+
+POSSIBLE_TYPES = int | float | bool
+
+
+def _trigger_mesh_update(obj: bpy.types.Object) -> None:
+    data = obj.data
+    if data is None or not isinstance(data, bpy.types.Mesh):
+        raise RuntimeError("Object data is None")
+
+    data.update()
+    # we can manually trigger an update of the scene if it is a mesh object
+    # by just re-writing the coordinates for a single vertex which properly
+    # triggers a refresh of the 3D scene
+    try:
+        vert = data.vertices[0]  # type: ignore
+        vert.co = list(vert.co)  # type: ignore
+    except Exception as e:
+        print(e)
+
+
+class NodesModifierInterface(BlenderObjectBase):
+    def __init__(self, modifier: bpy.types.NodesModifier):
+        assert isinstance(modifier, bpy.types.NodesModifier)
+        super().__init__(modifier.id_data)  # type: ignore
+        self._modifier_name = modifier.name
+
+    @property
+    def name(self) -> str:
+        return self.modifier.name
+
+    @name.setter
+    def name(self, value: str) -> None:
+        self.modifier.name = value
+        self._modifier_name = value
+
+    @property
+    def modifier(self) -> bpy.types.NodesModifier:
+        mod = self.object.modifiers[self._modifier_name]
+        assert isinstance(mod, bpy.types.NodesModifier)
+        return mod
+
+    @property
+    def tree(self) -> bpy.types.NodeTree:
+        tree = self.modifier.node_group
+        if tree is None:
+            raise RuntimeError("Tree not found")
+        return tree
+
+    @tree.setter
+    def tree(self, value: bpy.types.NodeTree) -> None:
+        assert isinstance(value, bpy.types.NodeTree)
+        self.modifier.node_group = value
+
+    @property
+    def tree_interface(self) -> bpy.types.NodeTreeInterface:
+        interface = self.tree.interface
+        assert interface is not None
+        return interface
+
+    @property
+    def input_sockets(self) -> list[bpy.types.NodeTreeInterfaceSocket]:
+        return [
+            item
+            for item in self.tree_interface.items_tree
+            if isinstance(item, bpy.types.NodeTreeInterfaceSocket)
+            and item.in_out == "INPUT"
+            and hasattr(item, "default_value")
+        ]
+
+    def get_id_from_name(self, name: str) -> str:
+        for item in self.input_sockets:
+            if item.name == name:
+                return item.identifier
+
+        raise ValueError(f"Input socket not found: {name=}")
+
+    def get_value(self, name: str) -> POSSIBLE_TYPES:
+        return self.modifier[self.get_id_from_name(name)]
+
+    def set_value(self, name: str, value: POSSIBLE_TYPES) -> None:
+        self.modifier[self.get_id_from_name(name)] = value
+        _trigger_mesh_update(self.object)
+
+    def _key_attr_name(self, name: str) -> str:
+        return "{}_attribute_name".format(self.get_id_from_name(name))
+
+    def get_attr_name(self, name: str) -> str:
+        return self.modifier[self._key_attr_name(name)]
+
+    def set_attr_name(self, name: str, value: str) -> None:
+        self.modifier[self._key_attr_name(name)] = value
+
+    def _key_use_att(self, name: str) -> str:
+        return "{}_use_attribute".format(self.get_id_from_name(name))
+
+    def get_attr_use(self, name: str):
+        return self.modifier[self._key_use_att(name)]
+
+    def set_attr_use(self, name: str, value: bool) -> None:
+        self.modifier[self._key_use_att(name)] = value
+
+    def list_inputs(self) -> list[str]:
+        """Return list of inputs names that can accept and return values"""
+        input_names: list[str] = []
+        for item in self.input_sockets:
+            try:
+                self.modifier[item.identifier]
+            except KeyError:
+                continue
+            input_names.append(item.name)
+
+        return input_names
+
+    def _ipython_key_completions_(self) -> list[str]:
+        return self.list_inputs()
+
+    def __getitem__(self, name: str) -> POSSIBLE_TYPES:
+        return self.get_value(name)
+
+    def __setitem__(self, name: str, value: POSSIBLE_TYPES) -> None:
+        self.set_value(name, value)
+
+    # might be able to get tab completion in Blender using this?
+    def __dir__(self) -> list[str]:
+        return self.list_inputs()

tests/test_modifier.py

@@ -0,0 +1,48 @@
+import bpy
+import pytest
+
+import databpy as db
+
+
+def test_modifier():
+    cube = bpy.data.objects["Cube"]
+    modifier = cube.modifiers.new("GeometryNodes", "NODES")
+    assert isinstance(modifier, bpy.types.NodesModifier)
+    tree = db.nodes.new_tree()
+    tree.interface.new_socket(  # type: ignore
+        name="Count", in_out="INPUT", socket_type="NodeSocketInt"
+    ).default_value = 3
+    mod = db.NodesModifierInterface(modifier)
+    with pytest.raises(RuntimeError):
+        mod.tree
+    mod.tree = tree
+    assert mod.name == "GeometryNodes"
+    mod.name = "New Name"
+    assert mod.name == "New Name"
+
+    assert mod["Count"] == 3
+    mod["Count"] = 4
+    assert mod["Count"] == 4
+
+    assert mod.list_inputs() == ["Count"]
+    assert mod.list_inputs() == mod._ipython_key_completions_()
+    assert mod.list_inputs() == dir(mod)
+
+    assert mod._key_use_att("Count") == "Socket_2_use_attribute"
+    assert mod._key_attr_name("Count") == "Socket_2_attribute_name"
+
+    assert mod.get_attr_name("Count") == ""
+    mod.set_attr_name("Count", "is_peptide")
+    assert mod.get_attr_name("Count") == "is_peptide"
+
+    assert not mod.get_attr_use("Count")
+    mod.set_attr_use("Count", True)
+    assert mod.get_attr_use("Count")
+
+    with pytest.raises(ValueError):
+        mod.get_id_from_name("non_existant_input")
+
+    wrong_mod = cube.modifiers.new(name="test", type="NORMAL_EDIT")
+
+    with pytest.raises(AssertionError):
+        db.NodesModifierInterface(wrong_mod)  # type: ignore