init

mednodes/custom_nodes/__init__.py

@@ -0,0 +1 @@
+from .node_menu import CustomNodes

mednodes/custom_nodes/node_menu.py

@@ -0,0 +1,131 @@
+import bpy
+from pathlib import Path
+from .nodes import AddCustomNode
+
+
+class CustomNodes():
+    def __init__(
+        self,
+        menu_items,
+        nodes_file,
+        root_label='CustomNodes',
+        root_icon=88
+    ) -> None:
+        if not Path(nodes_file).is_file():
+            raise FileNotFoundError(Path(nodes_file).resolve().as_posix())
+
+        self.menu_items = menu_items
+        self.nodes_file = Path(nodes_file).resolve().as_posix()
+
+        self.root_label = root_label
+        self.root_icon = root_icon
+
+        menu_classes = []
+        self._create_menu_class(
+            items=menu_items,
+            label=root_label,
+            menu_classes=menu_classes
+        )
+        self.menu_classes = menu_classes
+
+        idname = f"CUSTOMNODES_MT_NODES__{root_label.replace(' ', '').upper()}"
+
+        def add_node_menu(self, context):
+            if ('GeometryNodeTree' == bpy.context.area.spaces[0].tree_type):
+                layout = self.layout
+                layout.separator()
+                layout.menu(idname, icon_value=root_icon)
+
+        self.add_node_menu = add_node_menu
+
+    def _create_menu_class(self, menu_classes, items, label='CustomNodes', icon=0, idname_namespace=None):
+        nodes_file = self.nodes_file
+        idname_namespace = idname_namespace or "CUSTOMNODES_MT_NODES_"
+        idname = f"{idname_namespace}_{label.replace(' ', '').upper()}"
+
+        # create submenu class if item is menu.
+        for item in items:
+            item_items = item.get('items') if item != 'separator' else None
+            if item_items:
+                menu_class = self._create_menu_class(
+                    menu_classes=menu_classes,
+                    items=item_items,
+                    label=item.get('label') or 'CustomNodes',
+                    icon=item.get('icon') or 0,
+                    idname_namespace=idname
+                )
+                item['menu_class'] = menu_class
+
+        # create menu class
+        class Menu(bpy.types.Menu):
+            bl_idname = idname
+            bl_label = label
+
+            def draw(self, context):
+                layout = self.layout
+
+                for item in items:
+                    # print(item)
+                    if item == "separator":
+                        layout.separator()
+                    elif item.get('menu_class'):
+                        layout.menu(
+                            item.get('menu_class').bl_idname,
+                            icon_value=item.get('icon') or 0
+                        )
+                    else:
+                        op = layout.operator(
+                            'customnodes.add_custom_node', text=item.get('label'))
+                        op.nodes_file = nodes_file
+                        op.node_type = item['node_type']
+                        op.node_label = item.get('label') or ""
+                        op.node_description = item.get(
+                            'node_description') or "Add Custom Node."
+                        op.node_driver = item.get('node_driver') or ""
+
+        menu_classes.append(Menu)
+        return Menu
+
+    def _find_item(self, found_items, menu_items, node_type: str):
+
+        for item in menu_items:
+            if item.get("node_type") == node_type:
+                found_items.append(item)
+
+            item_items = item.get('items') if item != 'separator' else None
+            if item_items:
+                self._find_item(found_items, item_items, node_type)
+
+    def find_item(self, node_type: str):
+        found_items = []
+        self._find_item(found_items, self.menu_items, node_type)
+        return found_items[0] if len(found_items) > 0 else None
+
+    def add_node(self, node_tree, node_type: str):
+        item = self.find_item(node_type)
+        if item:
+            op = AddCustomNode()
+            op.nodes_file = self.nodes_file
+            op.node_type = item['node_type']
+            op.node_label = item.get('label') or ""
+            op.node_description = item.get(
+                'node_description') or "Add Custom Node."
+            op.node_driver = item.get('node_driver') or ""
+            return op.add_node(node_tree)
+        else:
+            raise RuntimeError("No custom node type found.")
+
+    def register(self):
+        for cls in self.menu_classes:
+            bpy.utils.register_class(cls)
+
+        bpy.types.NODE_MT_add.append(self.add_node_menu)
+
+    def unregister(self):
+        try:
+            for cls in reversed(self.menu_classes):
+                bpy.utils.unregister_class(cls)
+
+            bpy.types.NODE_MT_add.remove(self.add_node_menu)
+        except RuntimeError:
+            pass

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"}