socket_editor.pyw

import tkinter as tk
import tkinter.messagebox
import inspect
import sys
import os

from blender_asset_tracer import blendfile
import inspect
import pathlib
from dataclasses import dataclass, field


def lookup(object, mode = 0):
    if mode == 0:
        attributes = [attribute for attribute in inspect.getmembers(object) if not attribute[0].startswith('_')]
        print(*attributes, sep='\n')
    else:
        attributes = [attribute for attribute in dir(object) if not attribute.startswith('_')]
        print(*attributes, sep='\n')


files = sys.argv[1:]
if len(files) == 0:
    print("Drop a .blend file.")
    input("\nPress any key to exit.")
    quit()

blend_files = [file for file in files if file.lower().endswith(".blend")]

if len(blend_files) == 0:
    print("No .blend file was dropped. Drop a .blend file.")
    input("\nPress any key to exit.")
    quit()

class ScrollableFrame(tk.Frame):
    def __init__(self, parent, **kwargs):
        super().__init__(parent, **kwargs)

        self.canvas = tk.Canvas(self, highlightthickness = 0)
        self.scrollbar = tk.Scrollbar(self, orient="vertical", command = self.canvas.yview, width = 20)
        self.scrollable_frame = tk.Frame(self.canvas)

        self.w = self.canvas.create_window((0, 0), window = self.scrollable_frame, anchor = tk.NW)
        self.canvas.configure(yscrollcommand = self.scrollbar.set)

        self.columnconfigure(0, weight = 2)
        self.rowconfigure(0, weight = 2)

        self.canvas.grid(row = 0, column = 0, sticky = tk.NSEW)
        self.scrollbar.grid(row = 0, column = 1, sticky = tk.NSEW)

        self.bind("<Configure>", self.config_item)
        self.canvas.bind("<Configure>", self.config_canvas)
        self.canvas.bind_all("<MouseWheel>", self._on_mousewheel)

    def config_canvas(self, event = None):
        self.canvas.configure(scrollregion = self.canvas.bbox("all"))

    def config_item(self, event):
        self.canvas.itemconfigure(self.w, width = event.width - 25)

    def _on_mousewheel(self, event):
        self.canvas.yview_scroll(int(-1*(event.delta/120)), "units")

    def update(self):
        self.scrollbar.update()
        self.config_canvas()
        super().update()


class Application(tk.Frame):
    def __init__(self, parent, blend_file_path):
        super().__init__(parent)
        
        self.parent = parent
        self.parent.withdraw()

        self.blend_file_path = blend_file_path

        self.node_tree = tk.StringVar(self)
        self.get_node_groups()

        if len(self.node_trees) == 0:
            tkinter.messagebox.showinfo("Node groups not found", f"{self.blend_file_path} does not have any node groups.")
            self.parent.destroy()
        else:
            self.parent.deiconify()
            self.initUI()

    def initUI(self):

        self.pack(fill=tk.BOTH, expand=True)

        self.node_tree.set(self.node_trees[0])

        buttons = tk.Frame(self)

        save_button = tk.Button(buttons, text = 'Save' ,command = self.set_sockets, font = (14))
        save_button.pack(fill = tk.BOTH, padx=5, pady=5, side = tk.LEFT, expand=True)

        name_as_identifier = tk.Button(buttons, text = 'Name As Identifier' , command = self.set_name_as_identifier, font = (14))
        name_as_identifier.pack(fill = tk.BOTH, padx=5, pady=5, side = tk.LEFT, expand=True)

        option = tk.OptionMenu(buttons, self.node_tree, *self.node_trees)
        option.config(font = (14))
        option.pack(fill = tk.BOTH, padx=5, pady=5, side = tk.LEFT, expand=True)

        buttons.pack(fill = tk.BOTH)

        self.scrollable_frame = ScrollableFrame(self)
        self.scrollable_frame.pack(fill = tk.BOTH, expand=True, padx=5, pady=5)
        self.sockets_frame = self.scrollable_frame.scrollable_frame

        self.node_tree.trace("w", self.update_socket_grid)

        self.load_socket_grid(self.node_trees[0])

        self.centerWindow()

    def set_name_as_identifier(self):

        old_names = self.sockets_frame.grid_slaves(column= 1)
        old_names.reverse()
        old_names.pop(0)
        old_names = [name.get() for name in old_names]
        

        new_identifiers = self.sockets_frame.grid_slaves(column= 4)
        new_identifiers.reverse()
        new_identifiers.pop(0)

        for name, identifier in zip(old_names, new_identifiers):
            identifier.delete(0, tk.END)
            identifier.insert(0, name)
    
    def update_socket_grid(self, *args):
        slaves = self.sockets_frame.grid_slaves()
        if slaves:
            for slave in list(slaves):
                slave.destroy()

        self.load_socket_grid(self.node_tree.get())

        self.scrollable_frame.update()     

    def load_socket_grid(self, node_tree_name):

        self.sockets = self.get_sockets(node_tree_name)

        for i in range(5):
            self.sockets_frame.columnconfigure(i, weight = 1)
        self.sockets_frame.rowconfigure(0, weight = 1)

        type = tk.Label(self.sockets_frame, text="Type", font = (14))
        type.grid(row = 0, column = 0, sticky = tk.NSEW)

        old_name = tk.Label(self.sockets_frame, text = "Name", font = (14))
        old_name.grid(row = 0, column = 1, sticky = tk.NSEW)

        new_name = tk.Label(self.sockets_frame, text="New Name", font = (14))
        new_name.grid(row = 0, column = 2, sticky = tk.NSEW)

        old_identifier = tk.Label(self.sockets_frame, text="Identifier", font = (14))
        old_identifier.grid(row = 0, column = 3, sticky = tk.NSEW)

        new_identifier = tk.Label(self.sockets_frame, text="New Identifier", font = (14))
        new_identifier.grid(row = 0, column = 4, sticky = tk.NSEW)

        for row_index, socket in enumerate(self.sockets, start = 1):

            self.sockets_frame.rowconfigure(row_index, weight = 1)

            if socket.is_output:
                current_type = 'Output'
            else:
                current_type = 'Input'

            type = tk.Label(
                self.sockets_frame, 
                text = current_type, 
                background = self.row_coloring(row_index), 
                font = (14))
            type.grid(row = row_index, column = 0, sticky=tk.NSEW)

            old_name = tk.Entry(
                self.sockets_frame,
                font = (14),
                borderwidth = 5, 
                relief = tk.FLAT
                )
            old_name.insert(0, socket.name.decode("utf-8"))
            old_name.configure(state = "readonly", readonlybackground = self.row_coloring(row_index))
            old_name.grid(row = row_index, column = 1, sticky=tk.NSEW)

            new_name = tk.Entry(
                self.sockets_frame, 
                background = self.row_coloring(row_index), 
                font = (14),
                borderwidth = 1)
            new_name.grid(row = row_index, column = 2, sticky=tk.NSEW)

            old_identifier = tk.Entry(
                self.sockets_frame,
                font = (14),
                borderwidth = 5, 
                relief = tk.FLAT
                )
            old_identifier.insert(0, socket.identifier.decode("utf-8"))
            old_identifier.configure(state = "readonly", readonlybackground = self.row_coloring(row_index))
            old_identifier.grid(row = row_index, column = 3, sticky=tk.NSEW)

            new_identifier = tk.Entry(
                self.sockets_frame, 
                background = self.row_coloring(row_index), 
                font = (14),
                borderwidth = 1)
            new_identifier.grid(row = row_index, column = 4, sticky=tk.NSEW)

    def row_coloring(self, index):
        if index % 2 == 0:
            return "white"
        else:
            return "lightgray"

    def get_node_groups(self):

        with blendfile.open_cached(pathlib.Path(self.blend_file_path)) as blend:

            node_tree_data_blocks = blend.find_blocks_from_code(b'NT')

            self.node_trees = []
            for block in node_tree_data_blocks:
                self.node_trees.append(block.id_name[2:].decode())

    def get_sockets(self, node_tree_name):

        node_tree_name = node_tree_name.encode()

        with blendfile.open_cached(pathlib.Path(self.blend_file_path)) as blend:

            node_trees = blend.find_blocks_from_code(b'NT')

            for node_tree in node_trees:
                a = node_tree.id_name[2:]
                if node_tree.id_name[2:] == node_tree_name:

                    def collect_from_bat_list(from_block , name):
                        address = list(from_block.get_recursive_iter(name))[0][1]
                        items = []
                        while address != 0:
                            block = blend.block_from_addr[address]
                            items.append(block)
                            address = block.get(b'next')
                        return items

                    @dataclass
                    class Socket:
                        name: str
                        identifier: str
                        is_output: bool
                        data: field(default_factory=list)

                    sockets = []

                    group_inputs = collect_from_bat_list(node_tree, b'inputs')
                    group_outputs = collect_from_bat_list(node_tree, b'outputs')
                    nodes = collect_from_bat_list(node_tree, b'nodes')

                    def get_socket_list(node, socket_type):
                        node_sockets = collect_from_bat_list(node, socket_type)
                        for node_socket in node_sockets.copy():
                            if node_socket.get(b'identifier') == b'__extend__':
                                node_sockets.remove(node_socket)
                        return node_sockets

                    group_input_nodes = [node for node in nodes if node.get(b'idname') == b'NodeGroupInput']
                    group_output_nodes = [node for node in nodes if node.get(b'idname') == b'NodeGroupOutput']

                    input_socket_lists = [get_socket_list(node, b'outputs') for node in group_input_nodes]
                    output_socket_lists = [get_socket_list(node, b'inputs') for node in group_output_nodes]


                    for i, input in enumerate(group_inputs):
                        socket = Socket(input.get(b'name'), input.get(b'identifier'), False, [input.addr_old])

                        for input_socket_list in input_socket_lists:
                            socket.data.append(input_socket_list[i].addr_old)                       

                        sockets.append(socket)

                    for i, output in enumerate(group_outputs):
                        socket = Socket(output.get(b'name'), output.get(b'identifier'), True, [output.addr_old])

                        for output_socket_list in output_socket_lists:
                            socket.data.append(output_socket_list[i].addr_old)                       

                        sockets.append(socket)

                    def print_sockets():
                        for socket in sockets:
                            print(
                                "name : ",
                                socket.name,
                                "\nidentifier : ",
                                socket.identifier,
                                "\nis_output : ",
                                socket.is_output,
                                "\ndata : ",
                                socket.data,
                                sep=""
                            )
                            print()

                    return sockets

    def set_sockets(self):

        # print("--------------------")

        new_names = self.sockets_frame.grid_slaves(column= 2)
        new_names.reverse()
        new_names.pop(0)
        # print(*[name.grid_info() for name in new_names if name.widgetName == 'entry'], sep="\n")
        new_names = [name.get() for name in new_names if name.widgetName == 'entry']
        # print(*new_names, sep="\n")
        
        # print()

        new_identifiers = self.sockets_frame.grid_slaves(column= 4)
        new_identifiers.reverse()
        new_identifiers.pop(0)
        # print(*[identifier.grid_info() for identifier in new_identifiers if identifier.widgetName == 'entry'], sep="\n")
        new_identifiers = [identifier.get() for identifier in new_identifiers if identifier.widgetName == 'entry']
        # print(*new_identifiers, sep="\n")
        

        with blendfile.open_cached(pathlib.Path(self.blend_file_path), mode = "r+b") as blend:

            for index, socket in enumerate(self.sockets):

                if new_names[index] != "" or new_identifiers[index] != "":
                    socket_data_blocks = []
                    for address in socket.data:
                        socket_data_blocks.append(blend.block_from_addr[address])

                    for block in socket_data_blocks:
                        if new_names[index] != "":
                            block.set(b'name', new_names[index].encode())
                        if new_identifiers[index] != "":
                            block.set(b'identifier', new_identifiers[index].encode())

        self.update_socket_grid()

    def centerWindow(self):

        sw = self.master.winfo_screenwidth()
        sh = self.master.winfo_screenheight()

        w = sw/2
        h = sh/2

        x = (sw - w)/2
        y = (sh - h)/2
        self.master.geometry('%dx%d+%d+%d' % (w, h, x, y))

for blend_file in blend_files:
    blend_file_file_name = os.path.splitext(os.path.basename(blend_file))[0]

    root = tk.Tk()
    root.title("Socket Identifier Editor - " + blend_file_file_name)

    app = Application(root, blend_file)

    root.mainloop()