2022-11-12 23:59:00 +00:00
import bpy
2022-11-14 23:31:26 +00:00
import inspect
2022-11-12 23:59:00 +00:00
try :
import node_arrange as node_arrange
except :
pass
from . state import State
2022-11-13 19:38:31 +00:00
from . types import *
2022-11-12 23:59:00 +00:00
from . node_mapper import *
2022-11-29 02:09:27 +00:00
from . static . attribute import *
2023-02-12 16:43:16 +00:00
from . static . curve import *
2022-11-20 23:47:37 +00:00
from . static . expression import *
2022-11-29 02:09:27 +00:00
from . static . input_group import *
2023-11-11 19:43:12 +00:00
from . static . repeat import *
2022-12-01 12:37:35 +00:00
from . static . sample_mode import *
2022-12-01 14:07:07 +00:00
from . static . simulation import *
2022-12-01 12:37:35 +00:00
from . arrange import _arrange
2022-11-12 23:59:00 +00:00
2023-11-15 23:10:37 +00:00
IS_BLENDER_4 = bpy . app . version [ 0 ] > = 4
2022-11-13 19:38:31 +00:00
def _as_iterable ( x ) :
2022-12-01 12:37:35 +00:00
if isinstance ( x , Type ) :
return [ x , ]
2022-11-12 23:59:00 +00:00
try :
2022-11-13 19:38:31 +00:00
return iter ( x )
2022-11-12 23:59:00 +00:00
except TypeError :
2022-11-13 19:38:31 +00:00
return [ x , ]
2022-11-12 23:59:00 +00:00
2023-11-15 23:10:37 +00:00
def get_node_inputs ( x ) :
if IS_BLENDER_4 :
return [ i for i in x . interface . items_tree if i . item_type == ' SOCKET ' and i . in_out == ' INPUT ' ]
else :
return x . inputs
def get_node_outputs ( x ) :
if IS_BLENDER_4 :
return [ i for i in x . interface . items_tree if i . item_type == ' SOCKET ' and i . in_out == ' OUTPUT ' ]
else :
return x . outputs
2023-11-02 03:09:12 +00:00
2022-11-12 23:59:00 +00:00
def tree ( name ) :
tree_name = name
def build_tree ( builder ) :
2022-11-14 23:31:26 +00:00
signature = inspect . signature ( builder )
2022-11-12 23:59:00 +00:00
# Locate or create the node group
node_group = None
if tree_name in bpy . data . node_groups :
node_group = bpy . data . node_groups [ tree_name ]
else :
node_group = bpy . data . node_groups . new ( tree_name , ' GeometryNodeTree ' )
2023-11-02 03:09:12 +00:00
2023-11-15 23:10:37 +00:00
if IS_BLENDER_4 :
node_group . is_modifier = True
2023-11-02 03:09:12 +00:00
2022-11-12 23:59:00 +00:00
# Clear the node group before building
for node in node_group . nodes :
node_group . nodes . remove ( node )
2023-11-02 03:09:12 +00:00
node_inputs = get_node_inputs ( node_group )
input_count = sum ( map ( lambda p : len ( p . annotation . __annotations__ ) if issubclass ( p . annotation , InputGroup ) else 1 , list ( signature . parameters . values ( ) ) ) )
for node_input in node_inputs [ input_count : ] :
2023-11-15 23:10:37 +00:00
if IS_BLENDER_4 :
node_group . interface . remove ( node_input )
else :
node_group . inputs . remove ( node_input )
2023-11-02 03:09:12 +00:00
for group_output in get_node_outputs ( node_group ) :
2023-11-15 23:10:37 +00:00
if IS_BLENDER_4 :
node_group . interface . remove ( group_output )
else :
node_group . outputs . remove ( group_output )
2022-11-12 23:59:00 +00:00
# Setup the group inputs
group_input_node = node_group . nodes . new ( ' NodeGroupInput ' )
group_output_node = node_group . nodes . new ( ' NodeGroupOutput ' )
# Collect the inputs
inputs = { }
2022-11-17 23:13:07 +00:00
def validate_param ( param ) :
2022-11-14 23:31:26 +00:00
if param . annotation == inspect . Parameter . empty :
raise Exception ( f " Tree input ' { param . name } ' has no type specified. Please annotate with a valid node input type. " )
if not issubclass ( param . annotation , Type ) :
raise Exception ( f " Type of tree input ' { param . name } ' is not a valid ' Type ' subclass. " )
2022-11-17 23:13:07 +00:00
for param in signature . parameters . values ( ) :
if issubclass ( param . annotation , InputGroup ) :
2022-11-20 23:57:16 +00:00
instance = param . annotation ( )
2022-11-19 20:00:26 +00:00
prefix = ( param . annotation . prefix + " _ " ) if hasattr ( param . annotation , " prefix " ) else " "
2022-11-17 23:13:07 +00:00
for group_param , annotation in param . annotation . __annotations__ . items ( ) :
2022-11-20 23:57:16 +00:00
default = getattr ( instance , group_param , None )
inputs [ prefix + group_param ] = ( annotation , inspect . Parameter . empty if default is None else default , param . name , prefix )
2022-11-17 23:13:07 +00:00
else :
validate_param ( param )
2022-11-19 20:00:26 +00:00
inputs [ param . name ] = ( param . annotation , param . default , None , None )
2022-11-12 23:59:00 +00:00
# Create the input sockets and collect input values.
2023-11-02 03:09:12 +00:00
node_inputs = get_node_inputs ( node_group )
for i , node_input in enumerate ( node_inputs ) :
2022-11-14 23:31:26 +00:00
if node_input . bl_socket_idname != list ( inputs . values ( ) ) [ i ] [ 0 ] . socket_type :
2023-11-02 03:09:12 +00:00
for ni in node_inputs :
2023-11-15 23:10:37 +00:00
if IS_BLENDER_4 :
node_group . interface . remove ( ni )
else :
node_group . inputs . remove ( ni )
2022-11-14 23:31:26 +00:00
break
2022-11-17 23:13:07 +00:00
builder_inputs = { }
2023-11-02 03:09:12 +00:00
node_inputs = get_node_inputs ( node_group )
2022-11-12 23:59:00 +00:00
for i , arg in enumerate ( inputs . items ( ) ) :
2023-12-02 04:35:27 +00:00
if hasattr ( arg [ 1 ] [ 0 ] , ' input_options ' ) and arg [ 1 ] [ 0 ] . input_options . name != None :
input_name = arg [ 1 ] [ 0 ] . input_options . name
else :
input_name = arg [ 0 ] . replace ( ' _ ' , ' ' ) . title ( )
2023-11-02 03:09:12 +00:00
if len ( node_inputs ) > i :
node_inputs [ i ] . name = input_name
node_input = node_inputs [ i ]
2022-11-14 23:31:26 +00:00
else :
2023-11-15 23:10:37 +00:00
if IS_BLENDER_4 :
node_input = node_group . interface . new_socket ( socket_type = arg [ 1 ] [ 0 ] . socket_type , name = input_name , in_out = ' INPUT ' )
else :
node_input = node_group . inputs . new ( arg [ 1 ] [ 0 ] . socket_type , input_name )
2022-11-14 23:31:26 +00:00
if arg [ 1 ] [ 1 ] != inspect . Parameter . empty :
node_input . default_value = arg [ 1 ] [ 1 ]
2023-12-02 04:35:27 +00:00
if hasattr ( arg [ 1 ] [ 0 ] , ' input_options ' ) :
input_options = arg [ 1 ] [ 0 ] . input_options
input_options . process ( node_input . type )
node_input . min_value = input_options . min_value
node_input . max_value = input_options . max_value
node_input . bl_subtype_label = input_options . bl_subtype_label # DOES NOT WORK. Do e nee to change a UI property like this: ui_property(object, "property-name", expand=False, text="New Label")
node_input . description = input_options . description
node_input . hide_in_modifier = input_options . hide_in_modifier
else :
# reset all options to defaults ???????
pass
# input_options = InputOptions()
# input_options.process(node_input.type)
# node_input.min_value = input_options.min_value
# node_input.max_value = input_options.max_value
# node_input.bl_subtype_label = input_options.bl_subtype_label
# node_input.description = input_options.description
# node_input.hide_in_modifier = input_options.hide_in_modifier
2022-11-17 23:13:07 +00:00
if arg [ 1 ] [ 2 ] is not None :
if arg [ 1 ] [ 2 ] not in builder_inputs :
builder_inputs [ arg [ 1 ] [ 2 ] ] = signature . parameters [ arg [ 1 ] [ 2 ] ] . annotation ( )
2022-11-19 20:00:26 +00:00
setattr ( builder_inputs [ arg [ 1 ] [ 2 ] ] , arg [ 0 ] . replace ( arg [ 1 ] [ 3 ] , ' ' ) , arg [ 1 ] [ 0 ] ( group_input_node . outputs [ i ] ) )
2022-11-17 23:13:07 +00:00
else :
builder_inputs [ arg [ 0 ] ] = arg [ 1 ] [ 0 ] ( group_input_node . outputs [ i ] )
2022-11-12 23:59:00 +00:00
# Run the builder function
State . current_node_tree = node_group
2022-11-15 03:32:32 +00:00
if inspect . isgeneratorfunction ( builder ) :
2022-11-17 23:13:07 +00:00
generated_outputs = [ * builder ( * * builder_inputs ) ]
2022-11-15 03:32:32 +00:00
if all ( map ( lambda x : issubclass ( type ( x ) , Type ) and x . _socket . type == ' GEOMETRY ' , generated_outputs ) ) :
outputs = join_geometry ( geometry = generated_outputs )
else :
outputs = generated_outputs
else :
2022-11-17 23:13:07 +00:00
outputs = builder ( * * builder_inputs )
2022-11-12 23:59:00 +00:00
# Create the output sockets
2022-12-09 20:51:38 +00:00
if isinstance ( outputs , dict ) :
# Use a dict to name each return value
for i , ( k , v ) in enumerate ( outputs . items ( ) ) :
if not issubclass ( type ( v ) , Type ) :
v = Type ( value = v )
2023-11-15 23:10:37 +00:00
if IS_BLENDER_4 :
node_group . interface . new_socket ( socket_type = v . socket_type , name = k , in_out = ' OUTPUT ' )
else :
node_group . outputs . new ( v . socket_type , k )
2022-12-09 20:51:38 +00:00
node_group . links . new ( v . _socket , group_output_node . inputs [ i ] )
else :
for i , result in enumerate ( _as_iterable ( outputs ) ) :
if not issubclass ( type ( result ) , Type ) :
result = Type ( value = result )
# raise Exception(f"Return value '{result}' is not a valid 'Type' subclass.")
2023-11-15 23:10:37 +00:00
if IS_BLENDER_4 :
node_group . interface . new_socket ( socket_type = result . socket_type , name = ' Result ' , in_out = ' OUTPUT ' )
else :
node_group . outputs . new ( result . socket_type , ' Result ' )
2022-12-09 20:51:38 +00:00
node_group . links . new ( result . _socket , group_output_node . inputs [ i ] )
2022-11-12 23:59:00 +00:00
2022-12-01 12:37:35 +00:00
_arrange ( node_group )
2022-11-12 23:59:00 +00:00
# Return a function that creates a NodeGroup node in the tree.
# This lets @trees be used in other @trees via simple function calls.
def group_reference ( * args , * * kwargs ) :
2023-11-16 01:43:29 +00:00
if IS_BLENDER_4 :
result = geometrynodegroup ( node_tree = node_group , * args , * * kwargs )
else :
result = group ( node_tree = node_group , * args , * * kwargs )
2022-11-19 13:44:40 +00:00
group_outputs = [ ]
for group_output in result . _socket . node . outputs :
group_outputs . append ( Type ( group_output ) )
2022-11-19 20:00:26 +00:00
if len ( group_outputs ) == 1 :
return group_outputs [ 0 ]
else :
return tuple ( group_outputs )
2022-11-12 23:59:00 +00:00
return group_reference
if isinstance ( name , str ) :
return build_tree
else :
tree_name = name . __name__
return build_tree ( name )