init

mednodes/custom_nodes/nodes.py

@@ -0,0 +1,164 @@
+import bpy
+
+
+def add_driver(target_prop, var_sources, expression):
+
+    driver = target_prop.driver_add("default_value")
+    is_vector = isinstance(driver, list)
+    drivers = driver if is_vector else [driver]
+
+    for i, driver in enumerate(drivers):
+        for j, var_source in enumerate(var_sources):
+
+            source = var_source['source']
+            prop = var_source['prop']
+
+            var = driver.driver.variables.new()
+            var.name = f"var{j}"
+
+            var.targets[0].id_type = source.id_type
+            var.targets[0].id = source
+            var.targets[0].data_path = f'["{prop}"][{i}]'\
+                if is_vector else f'["{prop}"]'
+
+        driver.driver.expression = expression
+
+
+def add_direct_driver(target, target_prop, source, source_prop):
+    target_prop = target.inputs.get(target_prop)
+    drivers = [
+        {
+            "source": source,
+            "prop": source_prop
+        }
+    ]
+    expression = "var0"
+    add_driver(target_prop, drivers, expression)
+
+
+class AddCustomNode():
+
+    nodes_file: bpy.props.StringProperty(
+        name="nodes_file",
+        subtype='FILE_PATH',
+        default=""
+    )  # type: ignore
+
+    node_type: bpy.props.StringProperty(
+        name='node_type',
+        description='',
+        default='',
+        subtype='NONE',
+        maxlen=0
+    )  # type: ignore
+
+    node_label: bpy.props.StringProperty(
+        name='node_label',
+        default=''
+    )  # type: ignore
+
+    node_driver: bpy.props.StringProperty(
+        name='node_driver',
+        default=''
+    )  # type: ignore
+
+    node_description: bpy.props.StringProperty(
+        name="node_description",
+        description="",
+        default="Add custom node group.",
+        subtype="NONE"
+    )  # type: ignore
+
+    node_link: bpy.props.BoolProperty(
+        name='node_link',
+        default=True
+    )  # type: ignore
+
+    @classmethod
+    def description(cls, context, properties):
+        return properties.node_description
+
+    def assign_node_tree(self, node):
+        node.node_tree = bpy.data.node_groups[self.node_type]
+
+        node.width = 200.0
+        node.label = self.node_label or self.node_type
+        node.name = self.node_type
+
+        if self.node_driver:
+            source_prop = self.node_driver.split(":")[0]
+            target_prop = self.node_driver.split(":")[1]
+            add_direct_driver(
+                target=node,
+                target_prop=target_prop,
+                source=bpy.context.active_object,
+                source_prop=source_prop
+            )
+
+        return node
+
+    def append_node_tree(self, node_type):
+        # try to get node from current file if exists
+        node_tree = bpy.data.node_groups.get(node_type)
+        # if not exists, get it from asset file.
+        if not node_tree:
+            bpy.ops.wm.append(
+                'EXEC_DEFAULT',
+                directory=f"{self.nodes_file}/NodeTree",
+                filename=node_type,
+                link=False,
+                use_recursive=True
+            )
+
+        node_tree = bpy.data.node_groups.get(node_type)
+        if node_tree:
+            # self.recursive_append_material(node_tree)
+            return node_tree
+        else:
+            raise RuntimeError('No custom node found')
+
+    def recursive_append_material(self, node_tree):
+        for child in node_tree.nodes:
+            material_socket = child.inputs.get('Material')
+            if material_socket:
+                print(material_socket.default_value)
+            try:
+                self.recursive_append_material(child.node_tree)
+            except:
+                ...
+
+    def add_node(self, node_tree):
+        self.append_node_tree(self.node_type)
+
+        node = node_tree.new("GeometryNodeGroup")
+
+        self.assign_node_tree(node)
+
+        return node
+
+
+class CUSTOMNODES_OT_Add_Custom_Node(bpy.types.Operator, AddCustomNode):
+    bl_idname = "customnodes.add_custom_node"
+    bl_label = "Add Custom Node"
+    bl_options = {"REGISTER", "UNDO"}
+
+    def execute(self, context):
+        self.append_node_tree(self.node_type)
+
+        # intended to be called upon button press in the node tree
+        prev_context = bpy.context.area.type
+        bpy.context.area.type = 'NODE_EDITOR'
+        # actually invoke the operator to add a node to the current node tree
+        # use_transform=True ensures it appears where the user's mouse is and is currently
+        # being moved so the user can place it where they wish
+        bpy.ops.node.add_node(
+            'INVOKE_DEFAULT',
+            type='GeometryNodeGroup',
+            use_transform=True
+        )
+        bpy.context.area.type = prev_context
+        node = bpy.context.active_node
+
+        self.assign_node_tree(node)
+
+        return {"FINISHED"}