Tree

Root class for GeoNodes and ShaderNodes trees.

The system manages a stack of Trees. When a Tree is created, it is placed at the top of the stack and becomes the current tree. The Tree is poped from the stack with the method pop.

Better use the context management syntax:

with GeoNodes("My Tree") as tree:
    pass

Important

Trees scripted with geonodes are kept distinct from manually created trees by putting the marker string 'GEONODES' in the description attribute. If you initialize a Tree with the name of an existing tree:

• it will be replaced if it is a tree scripted with geonodes
• it will be renamed if it is not the case

This avoids to accidentally delete a manually created tree.

Warning

This doesn't work with shaders embedded in a Material. It is why the argument replace_material is set to False by default when creating a Shader. Hence, an existing material is not replaced by a geonodes scripts except if you change the default value of replace_material argument to True.

# Create the material if it doesn't exist
# Do nothing if it already exists
with ShaderNodes("My Material"):
    pass

# Replace the existing material
with ShaderNodes("My Material", replace_material=True):
    pass

Note

prefix argument is added at the begining of the name.

# The two following modifiers have the same name

with GeoNodes("My Modifier", prefix="Utils"):
    pass

with GeoNodes("Utils My Modifier"):
    pass
````

The `prefix` is usefull for big projects with numerous Groups and when you want the Groups to be sorted by their name.
You can then use the special class `G` to call the groups as python functions:

``` python
# Prefix
util = "Util"
math = "Math"

# Util Change Shape
with GeoNodes("Change Shape", prefix=util):
    pass

# Util Rotate
with GeoNodes("Rotate", prefix=util):
    pass

# Math Multiply
with GeoNodes("Multiply", prefix=math)
    pass

# Math Divide
with GeoNodes("Divide", prefix=math)
    pass

# Call the groups
with GeoNodes("My Modifier"):
    # Call a math group with the special class G
    a = G(math).divide(...)

Source code in core/treeclass.py

class Tree:
    """ Root class for GeoNodes and ShaderNodes trees.

    The system manages a stack of Trees. When a Tree is created, it is placed at the top of the stack
    and becomes the current tree.
    The Tree is poped from the stack with the method pop.

    Better use the context management syntax:

    ``` python
    with GeoNodes("My Tree") as tree:
        pass
    ```

    !!! important

        Trees scripted with **geonodes** are kept distinct from manually created trees by putting the
        marker string _'GEONODES'_ in the description attribute. If you initialize a Tree with the
        name of an existing tree:

        - it will be replaced if it is a tree scripted with **geonodes**
        - it will be renamed if it is not the case

        This avoids to accidentally delete a manually created tree.

    !!! warning
        This doesn't work with shaders embedded in a Material. It is why the argument `replace_material` is set to False
        by default when creating a Shader. Hence, an existing material is not replaced by a geonodes scripts
        except if you change the default value of `replace_material` argument to True.

    ``` python
    # Create the material if it doesn't exist
    # Do nothing if it already exists
    with ShaderNodes("My Material"):
        pass

    # Replace the existing material
    with ShaderNodes("My Material", replace_material=True):
        pass
    ```

    !!! note
        `prefix` argument is added at the begining of the name.

    ``` python
    # The two following modifiers have the same name

    with GeoNodes("My Modifier", prefix="Utils"):
        pass

    with GeoNodes("Utils My Modifier"):
        pass
    ````

    The `prefix` is usefull for big projects with numerous Groups and when you want the Groups to be sorted by their name.
    You can then use the special class `G` to call the groups as python functions:

    ``` python
    # Prefix
    util = "Util"
    math = "Math"

    # Util Change Shape
    with GeoNodes("Change Shape", prefix=util):
        pass

    # Util Rotate
    with GeoNodes("Rotate", prefix=util):
        pass

    # Math Multiply
    with GeoNodes("Multiply", prefix=math)
        pass

    # Math Divide
    with GeoNodes("Divide", prefix=math)
        pass

    # Call the groups
    with GeoNodes("My Modifier"):
        # Call a math group with the special class G
        a = G(math).divide(...)
    ```
    """

    _total_nodes = 0
    _total_links = 0

    TREE_STACK = []

    def __init__(self,
            tree_name        : str, 
            tree_type        : str   = 'GeometryNodeTree',
            *,
            fake_user        : bool  = False, 
            is_group         : bool  = False,
            prefix           : str   = "",
            replace_material : bool  = False,
            color_tag        : str = 'NONE',
            ):

        """ 

        Parameters
        ----------
        tree_name : str
            tree name

        tree_type : {'GeometryNodeTree', 'ShaderNodeTree'}, default='GeometryNodeTree'
            tree type 

        fake_user : bool, default=False
            fake user flag

        is_group : bool, default=False
            Group or not

        prefix : str, default=""
            prefix to add at the begining of the tree name

        replace_material : bool, default=False
            replace the existing matertial if True

        color_tag : str, default='NONE'
            group or modifier color_tag attribute

        """
        self._btree = None

        # Loaded vs created
        self._loaded = False

        # Arrange when finished
        self._use_arrange = True

        # In / out nodes cache
        self._input_node  = None
        self._output_node = None

        # ---------------------------------------------------------------------------
        # Stacks
        # ---------------------------------------------------------------------------

        self._panels  = [ItemPath(None)]
        self._layouts = []

        # Input / output socket creation are performed by node with dynamic
        # sockets (either with items or with tree_node)
        # By default, these are input_node and output_node
        # This behavior can be overloaded using input and output stacks

        self._output_stack  = []
        self._input_stack   = []

        # Prefix
        if prefix is None:
            prefix =  ""

        prefix = str(prefix)
        if len(prefix):
            tree_name = f"{prefix} {tree_name}"
        self._prefix = prefix

        # ----------------------------------------------------------------------------------------------------
        # Shader Node tree is a property of the Material
        # ----------------------------------------------------------------------------------------------------

        self.ignore = False

        if tree_type is None:
            btree = tree_name
            if not isinstance(btree, bpy.types.NodeTree):
                raise NodeError(
                    f"Tree inititialization error: when 'tree_type' is None, 'tree_name' must be a tree, not <{tree_name}>."
                    )
            self._btree = btree
            self._loaded = True

        elif tree_type == 'ShaderNodeTree':

            self._material = None
            if not is_group:
                self._material = bpy.data.materials.get(tree_name)

                if self._material is None:
                    self._material = bpy.data.materials.new(tree_name)
                else:
                    if not replace_material:
                        self._material = bpy.data.materials.new("GN TEMP MATERIAL")
                        self.ignore = True

                self._material.use_nodes = True
                self._btree = self._material.node_tree
                self._material.use_fake_user = fake_user

        # ----------------------------------------------------------------------------------------------------
        # Get the tree nodes from node_groups
        # ----------------------------------------------------------------------------------------------------

        if self._btree is None:
            self._btree = utils.get_tree(tree_name, tree_type=tree_type, create=True)
            self._btree.use_fake_user = fake_user

        self._is_group = is_group

        # Managed lists
        self._kept_nodes = []
        self._nodes      = {} # Registered nodes: bpy.types.Node.name -> Node
        self._mod_vals   = {} # To restore modifiers values when necessary (menus)

        # ---------------------------------------------------------------------------
        # Clear / Keep the tree
        # ---------------------------------------------------------------------------

        # Store the current value of each input socket
        if tree_type == 'GeometryNodeTree' or self._is_group:
            for bnode in self._btree.nodes:
                if bnode.bl_idname != 'NodeGroupInput':
                    continue

                for name, mod in self.get_modifiers().items():
                    values = {}
                    for bsock in bnode.outputs:
                        values[(bsock.name, bsock.bl_idname)] = mod.get(bsock.identifier)
                    self._mod_vals[name] = values

                break

        # Clear tree
        if not self._loaded:
            self.clear(keep_nodes=False)

        # ---------------------------------------------------------------------------
        # Color tag
        # ---------------------------------------------------------------------------

        if hasattr(self._btree, 'color_tag'):
            utils.set_rna_attr(self._btree, 'color_tag', color_tag)

        # ---------------------------------------------------------------------------
        # Interface
        # ---------------------------------------------------------------------------

        self._interface = None
        if self._is_group or tree_type == 'GeometryNodeTree':
            self._interface = TreeInterface(self._btree)
            self._interface.clear(use_bin=True)

        # ---------------------------------------------------------------------------
        # Random generator for random colors
        # ---------------------------------------------------------------------------

        self._rng = np.random.default_rng(0)

        # Load the tree
        if self._loaded:
            self.load_btree()


    # ====================================================================================================
    # Stack of trees
    # ====================================================================================================

    @staticmethod
    def current_tree():
        """ > Get the Current Tree.

        Returns None if no tree is currently open

        Returns
        -------
        Tree
            current tree or None

        """
        if len(Tree.TREE_STACK):
            return Tree.TREE_STACK[-1]
        else:
            return None

    # ----------------------------------------------------------------------------------------------------
    # Push the current tree zone
    # ----------------------------------------------------------------------------------------------------

    def push(self):
        """ > Make this tree zone the current one

        !!! important
            This methods shouldn't be called directly, better use a **with** context block.

        ``` python
        with Tree("My Name"):
            pass
        ```
        """
        Tree.TREE_STACK.append(self)
        return self

    # ----------------------------------------------------------------------------------------------------
    # Pop from the stack
    # ----------------------------------------------------------------------------------------------------

    def pop(self, error=False):
        """ > Remove this tree from the stack

        !!! important
            This methods shouldn't be called directly, better use a **with** context block.

        Raises
        ------
        NodeError
            if this tree is not the current one


        ``` python
        with Tree("My Name"):
            pass
        ```
        """
        # ---------------------------------------------------------------------------
        # Ignore changes
        # ---------------------------------------------------------------------------

        if self.ignore:
            print(f"CHANGES IGNORED : Material '{self._btree.name}' already exists. Use 'replace_material=True' to overwite it.")

            assert self._btree.bl_idname == 'ShaderNodeTree'
            assert self._material is not None

            bpy.data.materials.remove(self._material)
            self._btree = None

            return

        if self._loaded:
            return

        # ---------------------------------------------------------------------------
        # Finalize
        # ---------------------------------------------------------------------------

        # Remove kept nodes
        for bnode in self._kept_nodes:
            self._btree.nodes.remove(bnode)

        # Adjust menu inputs
        if self._has_tree and not error:
            for node in self._nodes.values():
                if '_default_menu' in node.__slots__:
                    node._tree_is_completed(self._mod_vals)

            innode = self.input_node
            for bsock in innode._bnode.outputs:
                if bsock.type != 'MENU' or not bsock.is_linked:
                    continue
                isock = self._interface.by_identifier(bsock.identifier)
                enums = list(utils.get_enums(isock, 'default_value'))
                if not len(enums):
                    continue

                for mod in self.get_modifiers().values():

                    cur_val = mod.get(bsock.identifier)
                    if cur_val is not None and (cur_val < 2 or cur_val > len(enums) + 2):
                        mod[bsock.identifier] = enums.index(isock.default_value) + 2


        # Check dead ends
        warnings = 0
        if not error:
            warnings = self.check_warnings()

        # Remove from stack
        tree = Tree.TREE_STACK.pop()
        if tree != self:
            raise RuntimeError(f"Error in tree stack management")

        # Empty the interface bin
        if (not error) and self._interface is not None:
            self._interface.empty_bin()

        # Last node in error
        if error and len(self._nodes):
            for node in reversed(self._nodes.values()):
                if node._bnode.bl_idname in ['NodeGroupOutput', 'NodeGroupInput']:
                    continue
                utils.set_node_error(node._bnode)
                break

        # Arrange the nodes
        self.arrange()

        # Stats
        if self._btree.bl_idname == 'ShaderNodeTree' and self._material is not None:
            name = f"Material '{self._material.name}'"
        else:
            name = f"Tree '{self._btree.name}'"

        print(f"{name} built: {self._str_stats}, warnings: {warnings}")

        Tree._total_nodes += len(self._btree.nodes)
        Tree._total_links += len(self._btree.links)      

    # ----------------------------------------------------------------------------------------------------
    # Context management
    # ----------------------------------------------------------------------------------------------------

    def __enter__(self):
        self.push()
        return self

    def __exit__(self, type, exc_value, traceback):

        ok = exc_value is None or isinstance(exc_value, Break)

        self.pop(not ok)

        if isinstance(exc_value, Break):
            return True

    # ====================================================================================================
    # Remove node groups
    # ====================================================================================================

    @staticmethod
    def remove_groups(names=None, prefix=None, geonodes=True, shadernodes=True):
        """ > Remove Groups created by GeoNodes.

        !!! important
             This method can only remove groups created by **GeoNodes**.

        Parameters
        ----------
        names : str or list of strs, optional
            name of the node groups to remove (all if None)

        prefix : str, optional
            name prefix for the groups to delete

        geonodes : bool, default=True
            remove geometry nodes groups

        shadernodes : bool, default=True
            remove shader nodes groups


        Returns
        -------
        None
        """

        groups = []
        if prefix is None:
            prefix_ = ""
        else:
            prefix_ = prefix + ' '

        tree_types = []
        if geonodes:
            tree_types.append('GeometryNodeTree')
        if shadernodes:
            tree_types.append('ShaderNodeTree')

        for group in bpy.data.node_groups:
            if group.bl_idname not in tree_types:
                continue

            if not group.description.startswith('GEONODES'):
                continue

            if prefix is not None and group.name[:len(prefix_)] != prefix_:
                continue

            if names is None:
                pass
            elif isinstance(names, str):
                if group.name != prefix_ + names:
                    continue
            elif isinstance(names, list):
                ok = False
                for name in names:
                    if group.name == prefix_ + name:
                        ok = True
                        break
                if not ok:
                    continue
            else:
                raise Exception(f"Argument names '{names}' is not valid: None, str or list of strs expected")

            groups.append(group)

        print("Remove node groups")
        for group in groups:
            print(' -', group.name)
            bpy.data.node_groups.remove(group)

        print(f"{len(groups)} node groups removed")

    # ====================================================================================================
    # Check if the blender node is valid in the context
    #
    # Nodes specific to tools can't be inserted in modifiers
    # ====================================================================================================

    def check_node_validity(self, bnode):
        pass

    # ====================================================================================================
    # Counters
    # ====================================================================================================

    def __str__(self):
        return f"<Tree '{self._btree.name}' ({type(self).__name__}): {len(self._btree.nodes)} nodes and {len(self._btree.links)} links>"

    @staticmethod
    def _btree_repr(btree):
        s = ""

        sepa = "\n   - "

        s += "\nInputs:"
        in_node, out_node = None, None
        for node in btree.nodes:
            if node.bl_idname == 'NodeGroupInput':
                in_node = node
            if node.bl_idname == 'NodeGroupOutput':
                out_node = node

        if in_node is not None:
            a = []
            for sock in in_node.outputs:
                name = f"[{sock.name}]"
                if not sock.enabled:
                    name += "( d)"
                if sock.is_linked:
                    name += " -> " + sock.links[0].to_socket.name
                a.append(name)

            s += sepa + sepa.join(a)

        s += "\nOutputs:"
        if out_node is not None:
            a = []
            for sock in out_node.inputs:
                name = f"[{sock.name}]"
                if not sock.enabled:
                    name += "( d)"
                if sock.is_linked:
                    name = sock.links[0].from_socket.name + " -> " + name
                a.append(name)
            s += sepa + sepa.join(a) + "\n"

        return s + "\n"


    def __repr__(self):
        return str(self) + "\n" + Tree._btree_repr(self._btree)

    @classmethod
    def _reset_counters(cls):
        cls._total_nodes = 0
        cls._total_links = 0
        cls._total_time  = 0.

    @classmethod
    def _display_counter(cls, title):
        print(f"{title}: {Tree._total_nodes} nodes, {Tree._total_links} links in {Tree._total_time:.1f} s")

    @property
    def _str_stats(self):
        return f"{len(self._btree.nodes)} nodes, {len(self._btree.links)} links"

    # ====================================================================================================
    # Types of tree
    # ====================================================================================================

    @classmethod
    def is_geonodes(cls):
        """ > Current Tree is Geometry Nodes.

        Returns
        -------
        True if Tree is GeoNodes, False otherwise
        """
        return cls.current_tree()._btree.bl_idname == 'GeometryNodeTree'

    @classmethod
    def is_shader(cls):
        """ > Current Tree is Shader Nodes.

        Returns
        -------
        True if Tree is ShaderNodes, False otherwise
        """
        return cls.current_tree()._btree.bl_idname == 'ShaderNodeTree'

    @classmethod
    def is_compositor(cls):
        return cls.current_tree()._btree.bl_idname == 'CompositorNodeTree'

    # ====================================================================================================
    # Panel
    # ====================================================================================================

    def get_panel(self, sub_panel: str = ""):
        return self._panels[-1] + ItemPath.to_item_path(sub_panel)

    def push_panel(self, new_panel: str):
        self._panels.append(ItemPath.to_item_path(new_panel))

    def pop_panel(self):
        self._panels.pop()
        assert len(self._panels) >= 1, f"Error in pushing / poping tree panels."

    # ====================================================================================================
    # Nodes
    # ====================================================================================================

    # ----------------------------------------------------------------------------------------------------
    # Clear the Tree
    # ----------------------------------------------------------------------------------------------------

    def clear(self, keep_nodes: bool = True):
        """ Clear the content of the Tree.

        Remove all the nodes in the Tree.
        """
        self._nodes.clear()

        # Keep menu switch nodes to preserver input links
        # They will be deleted at the end

        if keep_nodes:
            del_nodes = []
            self._kept_nodes = []
            for bnode in self._btree.nodes:
                if bnode.bl_idname in ['GeometryNodeMenuSwitch', 'NodeGroupInput']:
                    self._kept_nodes.append(bnode)
                else:
                    del_nodes.append(bnode)

            for bnode in del_nodes:
                self._btree.nodes.remove(bnode)

        else:
            self._btree.links.clear()
            self._btree.nodes.clear()

    # ----------------------------------------------------------------------------------------------------
    # Register a new Node
    # ----------------------------------------------------------------------------------------------------

    def register_node(self, node):
        """ Register a new Node.

        This method is called by the Node class constructor. Never call directly.

        Registered Nodes are stored in the dictionary : _nodes = bpy.types.Node.name -> Node

        Parameters
        ----------
        node : Node
            the newly created Node to register


        Returns
        -------
        Node
            the passed argument

        """

        self._nodes[node._bnode.name] = node
        if len(self._layouts):
            self._layouts[-1].include_node(node)

        node._stack = NodeError.stack_lines()

        if self._input_node is None and node._bnode.bl_idname == 'NodeGroupInput':
            self._input_node = node

        if self._output_node is None and node._bnode.bl_idname == 'NodeGroupOutput':
            self._output_node = node

        return node

    # ----------------------------------------------------------------------------------------------------
    # Load the tree
    # ----------------------------------------------------------------------------------------------------

    def load_btree(self):
        from .nodeclass import Node

        self._use_arrange = False

        self.push()
        for bnode in self._btree.nodes:
            if bnode.bl_idname in []:
                pass
            else:
                Node(bnode)
        self.pop()

    # ====================================================================================================
    # Sockets creation
    # ====================================================================================================

    # ----------------------------------------------------------------------------------------------------
    # Set the default value of an input socket
    # ----------------------------------------------------------------------------------------------------

    def set_input_socket_default(self, socket, value=None):

        if value is None:
            return

        bsocket = utils.get_bsocket(socket)
        if bsocket is None:
            return

        try:
            bsocket.default_value = value
        except Exception:
            pass

        io_socket = self._interface.by_identifier(bsocket.identifier)
        if io_socket is None:
            return

        if not hasattr(io_socket, 'default_value'):
            return

        try:
            io_socket.default_value = value
        except Exception:
            return

    # ----------------------------------------------------------------------------------------------------
    # Create an input socket
    # ----------------------------------------------------------------------------------------------------

    def create_input_socket(self, bl_idname, name, panel="", **props):
        """ Create a new input socket.

        This is an **input socket** of the zone, hence an **output socket** of the input node.

        Parameters
        ----------
        bl_idname : str
            Socket bl_idname

        name : str
            Socket name

        panel : str, optional
            Panel to place the socket in

        **props : dict, optional
            Properties specific to interface socket


        Returns
        -------
        Socket
        """        
        input_node = self.get_input_node()
        socket = input_node.create_socket('OUTPUT', bl_idname, name, panel=panel, **props)
        def_val = props.get('default_value')
        if def_val is not None:
            try:
                utils.get_bsocket(socket).default_value = def_val
            except Exception as e:
                print(f"WARNING Tree.create_input_socket: {str(e)}")

        return socket

    # ====================================================================================================
    # Get the signature
    # ====================================================================================================

    def get_signature(self, include: list = None, exclude: list = [], enabled_only=True, with_sockets: bool = False):
        """ Build the closure signature of the tree.

        The closure signature is the tuple made of the output signature of the input node
        and the input signature of the output node

        Returns
        -------
        Signature
        """
        return Signature(
            self.input_node.get_signature(
                include         = include,
                exclude         = exclude,
                enabled_only    = enabled_only,
                with_sockets    = with_sockets).outputs,
            self.output_node.get_signature(
                enabled_only    = enabled_only,
                with_sockets    = with_sockets).inputs)

    # ====================================================================================================
    # Arranges nodes
    # ====================================================================================================

    def arrange(self):
        """ > Arrange the nodes in the editor.

        Try to arrange properly the nodes from left to right.

        This method is called when the Tree is poped from the stack.

        Returns
        -------
        None
        """
        if self._use_arrange:
            treearrange.arrange(self._btree)

    # ====================================================================================================
    # Node colors
    # ====================================================================================================

    @staticmethod
    def _get_color(value):

        if isinstance(value, str):
            if value == 'RANDOM':
                rng = Tree.current_tree()._rng
                value = tuple(rng.uniform(0., .7, 3))

            elif value in ['OP', 'OPERATION']:
                value = (.406, .541, .608)
            elif value == 'MACRO':
                value = (.261, .963, .409)
            elif value == 'AUTO_GEN':
                value = (.583, .229, .963)
            elif value == 'WARNING':
                value = (.949, .574, .119)
            else:
                return SysColor(value)

        return SysColor(value)

    # ====================================================================================================
    # Create a link
    # ====================================================================================================

    def link(self, out_socket, in_socket, handle_dynamic_sockets=False):
        """ > Create a link between two sockets.

        Parameters
        ----------
        out_socket : socket
            a node output socket

        in_socket : socket
            input socket from another node


        Returns
        -------
        Link
        """

        #if hasattr(out_socket, '_bsocket'):
        if '_bsocket' in dir(out_socket):
            out_socket = out_socket._bsocket

        #if hasattr(in_socket, '_bsocket'):
        if '_bsocket' in dir(in_socket):
            in_socket = in_socket._bsocket

        link = self._btree.links.new(out_socket, in_socket, handle_dynamic_sockets=handle_dynamic_sockets)

        utils.check_link(link)
        utils.check_zones(self._btree)

        return link

    # ====================================================================================================
    # Tree Input / Output
    # ====================================================================================================

    @property
    def _has_tree(self):
        if self.is_geonodes():
            return True
        elif self.is_shader():
            return self._is_group
        return False

    # ----------------------------------------------------------------------------------------------------
    # Default input node
    # ----------------------------------------------------------------------------------------------------

    @property
    def input_node(self):
        """ > Return a Group Input node

        If the node doesn't already exist, it is created.

        Returns
        -------
        Node
        """
        from .nodeclass import Node

        if self._input_node is None:
            return Node('NodeGroupInput')
        else:
            return self._input_node

    # ----------------------------------------------------------------------------------------------------
    # Default output node
    # ----------------------------------------------------------------------------------------------------

    @property
    def output_node(self):
        """ Returns a Group Output node

        If the node doesn't already exist, it is created.

        Returns
        -------
        Node
        """
        from .nodeclass import Node

        if self._output_node is None:
            return Node('NodeGroupOutput')
        else:
            return self._output_node

    # ----------------------------------------------------------------------------------------------------
    # Using stack
    # ----------------------------------------------------------------------------------------------------

    def get_input_node(self):
        if len(self._input_stack):
            return self._input_stack[-1]
        else:
            return self.input_node

    def get_output_node(self):
        if len(self._output_stack):
            return self._output_stack[-1]
        else:
            return self.output_node

    # ====================================================================================================
    # Dict of modifiers using this tree
    # ====================================================================================================

    def get_modifiers(self):
        return blender.get_geonodes_modifiers(self._btree)

    # ====================================================================================================
    # Check warnings
    # ====================================================================================================

    def check_warnings(self):

        from .nodezone import ZoneNode

        tab = "\n   | " 

        count = 0

        dead_ends = []
        unlinkeds = []

        for node in self._nodes.values():

            if node._bnode.bl_idname in (
                'NodeGroupInput', 'GeometryNodeWarning',
                'GeometryNodeGizmoDial', 'GeometryNodeGizmoLinear', 'GeometryNodeGizmoTransform',
                ):
                continue

            # ---------------------------------------------------------------------------
            # Dead ends
            # ---------------------------------------------------------------------------

            is_linked = False
            n = 0
            for bsock in node._bnode.outputs:
                if bsock.is_linked:
                    is_linked = True
                    break
                n += 1

            if not is_linked and n > 0:
                dead_ends.append((node))

            # ---------------------------------------------------------------------------
            # Unlinked zone output nodes
            # ---------------------------------------------------------------------------

            if isinstance(node, ZoneNode):
                unlinked = []
                for bsock in node._bnode.inputs:
                    if bsock.type == 'CUSTOM':
                        continue
                    if bsock.name in ["Skip"]:
                        continue
                    if not bsock.is_linked:
                        unlinked.append(bsock.name)

                if len(unlinked):
                    unlinkeds.append((node, unlinked))

        # ---------------------------------------------------------------------------
        # Display
        # ---------------------------------------------------------------------------

        if len(dead_ends):
            print()
            print("-"*100)
            print(f"The following node{' is' if len(dead_ends) == 1 else 's are'} not connected")

        for i, node in enumerate(dead_ends):
            #node._bnode.label = f"DEAD END {i + 1}"
            utils.set_node_warning(node._bnode)

            print()
            print(node)
            if node._stack is not None:
                print(tab + tab.join(node._stack[1:]))

        if len(unlinkeds):
            print()
            print("-"*100)
            print(f"The output node of the following zone{' is' if len(unlinkeds) == 1 else 's are'} missing input links.")
            print("The variables uses in the loop are not updated.")
            print()

            for i, (node, names) in enumerate(unlinkeds):
                node._bnode.label = f"UPDATE MISSING {i + 1}"
                utils.set_node_warning(node._bnode)

                print()
                print(node)
                #node._bnode.label = f"DEAD END {i + 1}"
                print("Unlinked sockets:", names)
                if node._stack is not None:
                    print(tab + tab.join(node._stack[1:]))

        # Total number of warnings

        return len(dead_ends) + len(unlinkeds)

    # ====================================================================================================
    # Create a method calling this Group
    # ====================================================================================================

    def add_method(self, 
        target_class    : type, 
        func_name       : str = None,
        self_attr       : str = None, 
        ret_class       : type = None, 
        **fixed):
        """ Add a method calling the Group.

        Parameters
        ----------
        target_class : type
            class to add the method to

        func_name : str, optional
            name of the method to create (snae case version of group name if None) default=None.

        self_attr : str, optional
            self name attribute name default=None.

        ret_class : type, optional
            class to use to transtype the output socket default=None.

        fixed : dict
            fixed values for sockets        

        """
        from .nodeclass import Group

        group_name = self._btree.name
        if self._prefix != "":
            group_name = group_name[len(self._prefix) + 1:]

        Group.add_method(group_name, target_class,
            func_name       = func_name,
            self_attr       = self_attr, 
            ret_class       = ret_class, 
            prefix          = self._prefix,
            **fixed)


    # ====================================================================================================
    # Dump content
    # ====================================================================================================

    def dump(self):
        for bnode in self._btree.nodes:
            print(f"Node '{bnode.name}' ({bnode.bl_idname})")
            print("Inputs")
            for bsock in bnode.inputs:
                print(f"   - {bsock.name:20s} : {bsock.bl_idname}")
            print("Outputs")
            for bsock in bnode.outputs:
                print(f"   - {bsock.name:20s} : {bsock.bl_idname}")
            print()

    def gen_node_headers(self, default_values=True):
        for node in self._btree.nodes:
            if default_values:
                print(f"def FUNCTION(self,", ", ".join([utils.snake_case(sock.name) + f"={sock.default_value}" for sock in node.inputs]))
            else:
                print(f"def FUNCTION(self,", ", ".join([utils.snake_case(sock.name) + "=None" for sock in node.inputs]))

            sockets = ", ".join([f"'{sock.name}': {utils.snake_case(sock.name)}" for sock in node.inputs])
            print(f"Node('{node.name}', " + "{" + sockets + "}")
            print()

input_node property

Return a Group Input node

If the node doesn't already exist, it is created.

Returns:

Type	Description
Node

output_node property

Returns a Group Output node

If the node doesn't already exist, it is created.

Returns:

Type	Description
Node

__init__(tree_name, tree_type='GeometryNodeTree', *, fake_user=False, is_group=False, prefix='', replace_material=False, color_tag='NONE')

Parameters:

Name	Type	Description	Default
tree_name	str	tree name	required
tree_type	(GeometryNodeTree, ShaderNodeTree)	tree type	'GeometryNodeTree'
fake_user	bool	fake user flag	False
is_group	bool	Group or not	False
prefix	str	prefix to add at the begining of the tree name	""
replace_material	bool	replace the existing matertial if True	False
color_tag	str	group or modifier color_tag attribute	'NONE'

Source code in core/treeclass.py

def __init__(self,
        tree_name        : str, 
        tree_type        : str   = 'GeometryNodeTree',
        *,
        fake_user        : bool  = False, 
        is_group         : bool  = False,
        prefix           : str   = "",
        replace_material : bool  = False,
        color_tag        : str = 'NONE',
        ):

    """ 

    Parameters
    ----------
    tree_name : str
        tree name

    tree_type : {'GeometryNodeTree', 'ShaderNodeTree'}, default='GeometryNodeTree'
        tree type 

    fake_user : bool, default=False
        fake user flag

    is_group : bool, default=False
        Group or not

    prefix : str, default=""
        prefix to add at the begining of the tree name

    replace_material : bool, default=False
        replace the existing matertial if True

    color_tag : str, default='NONE'
        group or modifier color_tag attribute

    """
    self._btree = None

    # Loaded vs created
    self._loaded = False

    # Arrange when finished
    self._use_arrange = True

    # In / out nodes cache
    self._input_node  = None
    self._output_node = None

    # ---------------------------------------------------------------------------
    # Stacks
    # ---------------------------------------------------------------------------

    self._panels  = [ItemPath(None)]
    self._layouts = []

    # Input / output socket creation are performed by node with dynamic
    # sockets (either with items or with tree_node)
    # By default, these are input_node and output_node
    # This behavior can be overloaded using input and output stacks

    self._output_stack  = []
    self._input_stack   = []

    # Prefix
    if prefix is None:
        prefix =  ""

    prefix = str(prefix)
    if len(prefix):
        tree_name = f"{prefix} {tree_name}"
    self._prefix = prefix

    # ----------------------------------------------------------------------------------------------------
    # Shader Node tree is a property of the Material
    # ----------------------------------------------------------------------------------------------------

    self.ignore = False

    if tree_type is None:
        btree = tree_name
        if not isinstance(btree, bpy.types.NodeTree):
            raise NodeError(
                f"Tree inititialization error: when 'tree_type' is None, 'tree_name' must be a tree, not <{tree_name}>."
                )
        self._btree = btree
        self._loaded = True

    elif tree_type == 'ShaderNodeTree':

        self._material = None
        if not is_group:
            self._material = bpy.data.materials.get(tree_name)

            if self._material is None:
                self._material = bpy.data.materials.new(tree_name)
            else:
                if not replace_material:
                    self._material = bpy.data.materials.new("GN TEMP MATERIAL")
                    self.ignore = True

            self._material.use_nodes = True
            self._btree = self._material.node_tree
            self._material.use_fake_user = fake_user

    # ----------------------------------------------------------------------------------------------------
    # Get the tree nodes from node_groups
    # ----------------------------------------------------------------------------------------------------

    if self._btree is None:
        self._btree = utils.get_tree(tree_name, tree_type=tree_type, create=True)
        self._btree.use_fake_user = fake_user

    self._is_group = is_group

    # Managed lists
    self._kept_nodes = []
    self._nodes      = {} # Registered nodes: bpy.types.Node.name -> Node
    self._mod_vals   = {} # To restore modifiers values when necessary (menus)

    # ---------------------------------------------------------------------------
    # Clear / Keep the tree
    # ---------------------------------------------------------------------------

    # Store the current value of each input socket
    if tree_type == 'GeometryNodeTree' or self._is_group:
        for bnode in self._btree.nodes:
            if bnode.bl_idname != 'NodeGroupInput':
                continue

            for name, mod in self.get_modifiers().items():
                values = {}
                for bsock in bnode.outputs:
                    values[(bsock.name, bsock.bl_idname)] = mod.get(bsock.identifier)
                self._mod_vals[name] = values

            break

    # Clear tree
    if not self._loaded:
        self.clear(keep_nodes=False)

    # ---------------------------------------------------------------------------
    # Color tag
    # ---------------------------------------------------------------------------

    if hasattr(self._btree, 'color_tag'):
        utils.set_rna_attr(self._btree, 'color_tag', color_tag)

    # ---------------------------------------------------------------------------
    # Interface
    # ---------------------------------------------------------------------------

    self._interface = None
    if self._is_group or tree_type == 'GeometryNodeTree':
        self._interface = TreeInterface(self._btree)
        self._interface.clear(use_bin=True)

    # ---------------------------------------------------------------------------
    # Random generator for random colors
    # ---------------------------------------------------------------------------

    self._rng = np.random.default_rng(0)

    # Load the tree
    if self._loaded:
        self.load_btree()

add_method(target_class, func_name=None, self_attr=None, ret_class=None, **fixed)

Add a method calling the Group.

Parameters:

Name	Type	Description	Default
target_class	type	class to add the method to	required
func_name	str	name of the method to create (snae case version of group name if None) default=None.	None
self_attr	str	self name attribute name default=None.	None
ret_class	type	class to use to transtype the output socket default=None.	None
fixed	dict	fixed values for sockets	{}

Source code in core/treeclass.py

def add_method(self, 
    target_class    : type, 
    func_name       : str = None,
    self_attr       : str = None, 
    ret_class       : type = None, 
    **fixed):
    """ Add a method calling the Group.

    Parameters
    ----------
    target_class : type
        class to add the method to

    func_name : str, optional
        name of the method to create (snae case version of group name if None) default=None.

    self_attr : str, optional
        self name attribute name default=None.

    ret_class : type, optional
        class to use to transtype the output socket default=None.

    fixed : dict
        fixed values for sockets        

    """
    from .nodeclass import Group

    group_name = self._btree.name
    if self._prefix != "":
        group_name = group_name[len(self._prefix) + 1:]

    Group.add_method(group_name, target_class,
        func_name       = func_name,
        self_attr       = self_attr, 
        ret_class       = ret_class, 
        prefix          = self._prefix,
        **fixed)

arrange()

Arrange the nodes in the editor.

Try to arrange properly the nodes from left to right.

This method is called when the Tree is poped from the stack.

Returns:

Type	Description
None

Source code in core/treeclass.py

def arrange(self):
    """ > Arrange the nodes in the editor.

    Try to arrange properly the nodes from left to right.

    This method is called when the Tree is poped from the stack.

    Returns
    -------
    None
    """
    if self._use_arrange:
        treearrange.arrange(self._btree)

clear(keep_nodes=True)

Clear the content of the Tree.

Remove all the nodes in the Tree.

Source code in core/treeclass.py

def clear(self, keep_nodes: bool = True):
    """ Clear the content of the Tree.

    Remove all the nodes in the Tree.
    """
    self._nodes.clear()

    # Keep menu switch nodes to preserver input links
    # They will be deleted at the end

    if keep_nodes:
        del_nodes = []
        self._kept_nodes = []
        for bnode in self._btree.nodes:
            if bnode.bl_idname in ['GeometryNodeMenuSwitch', 'NodeGroupInput']:
                self._kept_nodes.append(bnode)
            else:
                del_nodes.append(bnode)

        for bnode in del_nodes:
            self._btree.nodes.remove(bnode)

    else:
        self._btree.links.clear()
        self._btree.nodes.clear()

create_input_socket(bl_idname, name, panel='', **props)

Create a new input socket.

This is an input socket of the zone, hence an output socket of the input node.

Parameters:

Name	Type	Description	Default
bl_idname	str	Socket bl_idname	required
name	str	Socket name	required
panel	str	Panel to place the socket in	''
**props	dict	Properties specific to interface socket	{}

Returns:

Type	Description
Socket

Source code in core/treeclass.py

def create_input_socket(self, bl_idname, name, panel="", **props):
    """ Create a new input socket.

    This is an **input socket** of the zone, hence an **output socket** of the input node.

    Parameters
    ----------
    bl_idname : str
        Socket bl_idname

    name : str
        Socket name

    panel : str, optional
        Panel to place the socket in

    **props : dict, optional
        Properties specific to interface socket


    Returns
    -------
    Socket
    """        
    input_node = self.get_input_node()
    socket = input_node.create_socket('OUTPUT', bl_idname, name, panel=panel, **props)
    def_val = props.get('default_value')
    if def_val is not None:
        try:
            utils.get_bsocket(socket).default_value = def_val
        except Exception as e:
            print(f"WARNING Tree.create_input_socket: {str(e)}")

    return socket

current_tree() staticmethod

Get the Current Tree.

Returns None if no tree is currently open

Returns:

Type	Description
Tree	current tree or None

Source code in core/treeclass.py

@staticmethod
def current_tree():
    """ > Get the Current Tree.

    Returns None if no tree is currently open

    Returns
    -------
    Tree
        current tree or None

    """
    if len(Tree.TREE_STACK):
        return Tree.TREE_STACK[-1]
    else:
        return None

get_signature(include=None, exclude=[], enabled_only=True, with_sockets=False)

Build the closure signature of the tree.

The closure signature is the tuple made of the output signature of the input node and the input signature of the output node

Returns:

Type	Description
Signature

Source code in core/treeclass.py

def get_signature(self, include: list = None, exclude: list = [], enabled_only=True, with_sockets: bool = False):
    """ Build the closure signature of the tree.

    The closure signature is the tuple made of the output signature of the input node
    and the input signature of the output node

    Returns
    -------
    Signature
    """
    return Signature(
        self.input_node.get_signature(
            include         = include,
            exclude         = exclude,
            enabled_only    = enabled_only,
            with_sockets    = with_sockets).outputs,
        self.output_node.get_signature(
            enabled_only    = enabled_only,
            with_sockets    = with_sockets).inputs)

is_geonodes() classmethod

Current Tree is Geometry Nodes.

Returns:

Type	Description
True if Tree is GeoNodes, False otherwise

Source code in core/treeclass.py

@classmethod
def is_geonodes(cls):
    """ > Current Tree is Geometry Nodes.

    Returns
    -------
    True if Tree is GeoNodes, False otherwise
    """
    return cls.current_tree()._btree.bl_idname == 'GeometryNodeTree'

is_shader() classmethod

Current Tree is Shader Nodes.

Returns:

Type	Description
True if Tree is ShaderNodes, False otherwise

Source code in core/treeclass.py

@classmethod
def is_shader(cls):
    """ > Current Tree is Shader Nodes.

    Returns
    -------
    True if Tree is ShaderNodes, False otherwise
    """
    return cls.current_tree()._btree.bl_idname == 'ShaderNodeTree'

link(out_socket, in_socket, handle_dynamic_sockets=False)

Create a link between two sockets.

Parameters:

Name	Type	Description	Default
out_socket	socket	a node output socket	required
in_socket	socket	input socket from another node	required

Returns:

Type	Description
Link

Source code in core/treeclass.py

def link(self, out_socket, in_socket, handle_dynamic_sockets=False):
    """ > Create a link between two sockets.

    Parameters
    ----------
    out_socket : socket
        a node output socket

    in_socket : socket
        input socket from another node


    Returns
    -------
    Link
    """

    #if hasattr(out_socket, '_bsocket'):
    if '_bsocket' in dir(out_socket):
        out_socket = out_socket._bsocket

    #if hasattr(in_socket, '_bsocket'):
    if '_bsocket' in dir(in_socket):
        in_socket = in_socket._bsocket

    link = self._btree.links.new(out_socket, in_socket, handle_dynamic_sockets=handle_dynamic_sockets)

    utils.check_link(link)
    utils.check_zones(self._btree)

    return link

pop(error=False)

Remove this tree from the stack

Important

This methods shouldn't be called directly, better use a with context block.

Raises:

Type	Description
NodeError	if this tree is not the current one
``` python
with Tree("My Name"):	pass
```

Source code in core/treeclass.py

def pop(self, error=False):
    """ > Remove this tree from the stack

    !!! important
        This methods shouldn't be called directly, better use a **with** context block.

    Raises
    ------
    NodeError
        if this tree is not the current one


    ``` python
    with Tree("My Name"):
        pass
    ```
    """
    # ---------------------------------------------------------------------------
    # Ignore changes
    # ---------------------------------------------------------------------------

    if self.ignore:
        print(f"CHANGES IGNORED : Material '{self._btree.name}' already exists. Use 'replace_material=True' to overwite it.")

        assert self._btree.bl_idname == 'ShaderNodeTree'
        assert self._material is not None

        bpy.data.materials.remove(self._material)
        self._btree = None

        return

    if self._loaded:
        return

    # ---------------------------------------------------------------------------
    # Finalize
    # ---------------------------------------------------------------------------

    # Remove kept nodes
    for bnode in self._kept_nodes:
        self._btree.nodes.remove(bnode)

    # Adjust menu inputs
    if self._has_tree and not error:
        for node in self._nodes.values():
            if '_default_menu' in node.__slots__:
                node._tree_is_completed(self._mod_vals)

        innode = self.input_node
        for bsock in innode._bnode.outputs:
            if bsock.type != 'MENU' or not bsock.is_linked:
                continue
            isock = self._interface.by_identifier(bsock.identifier)
            enums = list(utils.get_enums(isock, 'default_value'))
            if not len(enums):
                continue

            for mod in self.get_modifiers().values():

                cur_val = mod.get(bsock.identifier)
                if cur_val is not None and (cur_val < 2 or cur_val > len(enums) + 2):
                    mod[bsock.identifier] = enums.index(isock.default_value) + 2


    # Check dead ends
    warnings = 0
    if not error:
        warnings = self.check_warnings()

    # Remove from stack
    tree = Tree.TREE_STACK.pop()
    if tree != self:
        raise RuntimeError(f"Error in tree stack management")

    # Empty the interface bin
    if (not error) and self._interface is not None:
        self._interface.empty_bin()

    # Last node in error
    if error and len(self._nodes):
        for node in reversed(self._nodes.values()):
            if node._bnode.bl_idname in ['NodeGroupOutput', 'NodeGroupInput']:
                continue
            utils.set_node_error(node._bnode)
            break

    # Arrange the nodes
    self.arrange()

    # Stats
    if self._btree.bl_idname == 'ShaderNodeTree' and self._material is not None:
        name = f"Material '{self._material.name}'"
    else:
        name = f"Tree '{self._btree.name}'"

    print(f"{name} built: {self._str_stats}, warnings: {warnings}")

    Tree._total_nodes += len(self._btree.nodes)
    Tree._total_links += len(self._btree.links)      

push()

Make this tree zone the current one

Important

This methods shouldn't be called directly, better use a with context block.

with Tree("My Name"):
    pass

Source code in core/treeclass.py

def push(self):
    """ > Make this tree zone the current one

    !!! important
        This methods shouldn't be called directly, better use a **with** context block.

    ``` python
    with Tree("My Name"):
        pass
    ```
    """
    Tree.TREE_STACK.append(self)
    return self

register_node(node)

Register a new Node.

This method is called by the Node class constructor. Never call directly.

Registered Nodes are stored in the dictionary : _nodes = bpy.types.Node.name -> Node

Parameters:

Name	Type	Description	Default
node	Node	the newly created Node to register	required

Returns:

Type	Description
Node	the passed argument

Source code in core/treeclass.py

def register_node(self, node):
    """ Register a new Node.

    This method is called by the Node class constructor. Never call directly.

    Registered Nodes are stored in the dictionary : _nodes = bpy.types.Node.name -> Node

    Parameters
    ----------
    node : Node
        the newly created Node to register


    Returns
    -------
    Node
        the passed argument

    """

    self._nodes[node._bnode.name] = node
    if len(self._layouts):
        self._layouts[-1].include_node(node)

    node._stack = NodeError.stack_lines()

    if self._input_node is None and node._bnode.bl_idname == 'NodeGroupInput':
        self._input_node = node

    if self._output_node is None and node._bnode.bl_idname == 'NodeGroupOutput':
        self._output_node = node

    return node

remove_groups(names=None, prefix=None, geonodes=True, shadernodes=True) staticmethod

Remove Groups created by GeoNodes.

Important

This method can only remove groups created by GeoNodes.

Parameters:

Name	Type	Description	Default
names	str or list of strs	name of the node groups to remove (all if None)	None
prefix	str	name prefix for the groups to delete	None
geonodes	bool	remove geometry nodes groups	True
shadernodes	bool	remove shader nodes groups	True

Returns:

Type	Description
None

Source code in core/treeclass.py

@staticmethod
def remove_groups(names=None, prefix=None, geonodes=True, shadernodes=True):
    """ > Remove Groups created by GeoNodes.

    !!! important
         This method can only remove groups created by **GeoNodes**.

    Parameters
    ----------
    names : str or list of strs, optional
        name of the node groups to remove (all if None)

    prefix : str, optional
        name prefix for the groups to delete

    geonodes : bool, default=True
        remove geometry nodes groups

    shadernodes : bool, default=True
        remove shader nodes groups


    Returns
    -------
    None
    """

    groups = []
    if prefix is None:
        prefix_ = ""
    else:
        prefix_ = prefix + ' '

    tree_types = []
    if geonodes:
        tree_types.append('GeometryNodeTree')
    if shadernodes:
        tree_types.append('ShaderNodeTree')

    for group in bpy.data.node_groups:
        if group.bl_idname not in tree_types:
            continue

        if not group.description.startswith('GEONODES'):
            continue

        if prefix is not None and group.name[:len(prefix_)] != prefix_:
            continue

        if names is None:
            pass
        elif isinstance(names, str):
            if group.name != prefix_ + names:
                continue
        elif isinstance(names, list):
            ok = False
            for name in names:
                if group.name == prefix_ + name:
                    ok = True
                    break
            if not ok:
                continue
        else:
            raise Exception(f"Argument names '{names}' is not valid: None, str or list of strs expected")

        groups.append(group)

    print("Remove node groups")
    for group in groups:
        print(' -', group.name)
        bpy.data.node_groups.remove(group)

    print(f"{len(groups)} node groups removed")