From 58bc0b635a09e6ceb382e81b8bd9594c667e9fb2 Mon Sep 17 00:00:00 2001 From: Carson Katri Date: Sat, 12 Nov 2022 18:59:00 -0500 Subject: [PATCH] Add missing lib folder --- .gitignore | 1 - lib/__init__.py | 3 + lib/node_mapper.py | 208 +++++++++++++++++++++++++++++++++++++++++++++ lib/state.py | 3 + lib/tree.py | 107 +++++++++++++++++++++++ lib/types.py | 53 ++++++++++++ 6 files changed, 374 insertions(+), 1 deletion(-) create mode 100644 lib/__init__.py create mode 100644 lib/node_mapper.py create mode 100644 lib/state.py create mode 100644 lib/tree.py create mode 100644 lib/types.py diff --git a/.gitignore b/.gitignore index 98fd9f0..d594108 100644 --- a/.gitignore +++ b/.gitignore @@ -16,7 +16,6 @@ dist/ downloads/ eggs/ .eggs/ -lib/ lib64/ parts/ sdist/ diff --git a/lib/__init__.py b/lib/__init__.py new file mode 100644 index 0000000..e2a6537 --- /dev/null +++ b/lib/__init__.py @@ -0,0 +1,3 @@ +from .tree import * +from .types import * +from .node_mapper import * \ No newline at end of file diff --git a/lib/node_mapper.py b/lib/node_mapper.py new file mode 100644 index 0000000..c0f6673 --- /dev/null +++ b/lib/node_mapper.py @@ -0,0 +1,208 @@ +import bpy +from .state import State +from .types import Type + +class OutputsList(dict): + __getattr__ = dict.get + __setattr__ = dict.__setitem__ + __delattr__ = dict.__delitem__ + +def build_node(node_type): + def build(_primary_arg=None, **kwargs): + node = State.current_node_tree.nodes.new(node_type.__name__) + if _primary_arg is not None: + State.current_node_tree.links.new(_primary_arg._socket, node.inputs[0]) + for prop in node.bl_rna.properties: + argname = prop.name.lower().replace(' ', '_') + if argname in kwargs: + setattr(node, prop.identifier, kwargs[argname]) + for node_input in (node.inputs[1:] if _primary_arg is not None else node.inputs): + argname = node_input.name.lower().replace(' ', '_') + if argname in kwargs: + if node_input.is_multi_input and hasattr(kwargs[argname], '__iter__') and len(kwargs[argname]) > 0 and issubclass(type(next(iter(kwargs[argname]))), Type): + for x in kwargs[argname]: + State.current_node_tree.links.new(x._socket, node_input) + elif issubclass(type(kwargs[argname]), Type): + State.current_node_tree.links.new(kwargs[argname]._socket, node_input) + else: + try: + node_input.default_value = kwargs[argname] + except: + constant = Type(value=kwargs[argname]) + State.current_node_tree.links.new(constant._socket, node_input) + outputs = {} + for node_output in node.outputs: + outputs[node_output.name.lower().replace(' ', '_')] = Type(node_output) + if len(outputs) == 1: + return list(outputs.values())[0] + else: + return OutputsList(outputs) + return build + +documentation = {} +registered_nodes = set() +def register_node(node_type, category_path=None): + if node_type in registered_nodes: + return + snake_case_name = node_type.bl_rna.name.lower().replace(' ', '_') + globals()[snake_case_name] = build_node(node_type) + globals()[snake_case_name].bl_category_path = category_path + globals()[snake_case_name].bl_node_type = node_type + documentation[snake_case_name] = globals()[snake_case_name] + def build_node_method(node_type): + def build(self, *args, **kwargs): + return build_node(node_type)(self, *args, **kwargs) + return build + setattr(Type, snake_case_name, build_node_method(node_type)) + registered_nodes.add(node_type) +for category_name in list(filter(lambda x: x.startswith('NODE_MT_category_GEO_'), dir(bpy.types))): + category = getattr(bpy.types, category_name) + category_path = category.category.name.lower().replace(' ', '_') + for node in category.category.items(None): + node_type = getattr(bpy.types, node.nodetype) + register_node(node_type, category_path) +for node_type in bpy.types.GeometryNode.__subclasses__(): + register_node(node_type) + +def create_documentation(): + temp_node_group = bpy.data.node_groups.new('temp_node_group', 'GeometryNodeTree') + color_mappings = { + 'INT': '#598C5C', + 'FLOAT': '#A1A1A1', + 'BOOLEAN': '#CCA6D6', + 'GEOMETRY': '#00D6A3', + 'VALUE': '#A1A1A1', + 'VECTOR': '#6363C7', + 'MATERIAL': '#EB7582', + 'TEXTURE': '#9E4FA3', + 'COLLECTION': '#F5F5F5', + 'OBJECT': '#ED9E5C', + 'STRING': '#70B2FF', + 'RGBA': '#C7C729', + } + default_color = '#A1A1A1' + docstrings = [] + symbols = [] + for func in sorted(documentation.keys()): + method = documentation[func] + link = f"https://docs.blender.org/manual/en/latest/modeling/geometry_nodes/{method.bl_category_path}/{func}.html" + image = f"https://docs.blender.org/manual/en/latest/_images/node-types_{method.bl_node_type.__name__}" + node_instance = temp_node_group.nodes.new(method.bl_node_type.__name__) + props_inputs = {} + symbol_inputs = {} + parent_props = [prop.identifier for base in method.bl_node_type.__bases__ for prop in base.bl_rna.properties] + for prop in method.bl_node_type.bl_rna.properties: + if not prop.identifier in parent_props: + if prop.type == 'ENUM': + enum_items = 'Literal[' + ', '.join(map(lambda i: f"'{i.identifier}'", prop.enum_items)) + ']' + props_inputs[prop.identifier] = f"{enum_items}" + symbol_inputs[prop.identifier] = enum_items + else: + props_inputs[prop.identifier] = f"{prop.type.title()}" + symbol_inputs[prop.identifier] = prop.type.title() + primary_arg = None + for node_input in node_instance.inputs: + name = node_input.name.lower().replace(' ', '_') + typename = type(node_input).__name__.replace('NodeSocket', '') + if node_input.is_multi_input: + typename = f"List[{typename}]" + type_str = f"{typename}" + if name in props_inputs: + props_inputs[name] = props_inputs[name] + f' | {type_str}' + symbol_inputs[name] = symbol_inputs[name] + f' | {typename}' + else: + props_inputs[name] = type_str + symbol_inputs[name] = typename + if primary_arg is None: + primary_arg = (name, props_inputs[name]) + arg_docs = [] + symbol_args = [] + for name, value in props_inputs.items(): + arg_docs.append(f"{name}: {value}") + symbol_args.append(f"{name}: {symbol_inputs[name]}") + outputs = {} + symbol_outputs = {} + for node_output in node_instance.outputs: + output_name = node_output.name.lower().replace(' ', '_') + output_type = type(node_output).__name__.replace('NodeSocket', '') + outputs[output_name] = f"{output_type}" + symbol_outputs[output_name] = output_type + output_docs = [] + output_symbols = [] + for name, value in outputs.items(): + output_docs.append(f"{name}: {value}") + output_symbols.append(f"{name}: {symbol_outputs[name]}") + outputs_doc = f"{{ {', '.join(output_docs)} }}" if len(output_docs) > 1 else ''.join(output_docs) + arg_separator = ',\n ' + def primary_arg_docs(): + return f""" +

Chain Syntax

+ 
{primary_arg[0]}: {primary_arg[1]} = ...
+{primary_arg[0]}.{func}(...)
+ """ + docstrings.append(f""" +
+ {func} - {method.bl_node_type.bl_rna.name} +
+ +

Signature

+ 
{func}(
+  {arg_separator.join(arg_docs)}
+)
+

Result

+ 
{outputs_doc}
+ {primary_arg_docs() if primary_arg is not None else ""} +
+
+ """) + output_symbol_separator = '\n ' + symbol_return_type = f"_{func}_result" + if len(output_symbols) > 1: + symbols.append(f"""class {symbol_return_type}: + {output_symbol_separator.join(output_symbols)}""") + symbols.append(f"""def {func}({', '.join(symbol_args)}) -> {list(symbol_outputs.values())[0] if len(output_symbols) == 1 else symbol_return_type}: pass""") + bpy.data.node_groups.remove(temp_node_group) + html = f""" + + + + + +

Geometry Script

+

Nodes

+ {''.join(docstrings)} + + + """ + with open('documentation.html', 'w') as f: + f.write(html) + with open('geometry_script.py', 'w') as f: + newline = '\n' + def type_symbol(t): + return f"class {t.__name__}: pass" + f.write(f"""from typing import * +{newline.join(map(type_symbol, Type.__subclasses__()))} +{newline.join(symbols)}""") + +def create_docs(): + create_documentation() +bpy.app.timers.register(create_docs) \ No newline at end of file diff --git a/lib/state.py b/lib/state.py new file mode 100644 index 0000000..b61133f --- /dev/null +++ b/lib/state.py @@ -0,0 +1,3 @@ +# Tree generation state +class State: + current_node_tree = None \ No newline at end of file diff --git a/lib/tree.py b/lib/tree.py new file mode 100644 index 0000000..3db607c --- /dev/null +++ b/lib/tree.py @@ -0,0 +1,107 @@ +import bpy +import re +from inspect import getfullargspec +try: + import node_arrange as node_arrange +except: + pass +from .state import State +from .types import Type +from .node_mapper import * + +def _as_iterable(input): + try: + return iter(input) + except TypeError: + return {input} + +def tree(name): + tree_name = name + def build_tree(builder): + # 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') + # Clear the node group before building + for node in node_group.nodes: + node_group.nodes.remove(node) + for group_input in node_group.inputs: + node_group.inputs.remove(group_input) + for group_output in node_group.outputs: + node_group.outputs.remove(group_output) + + # Setup the group inputs + group_input_node = node_group.nodes.new('NodeGroupInput') + group_output_node = node_group.nodes.new('NodeGroupOutput') + + # Collect the inputs + argspec = getfullargspec(builder) + inputs = {} + for arg in argspec.args: + if not arg in argspec.annotations: + raise Exception(f"Tree input '{arg}' has no type specified. Please specify a valid NodeInput subclass.") + type_annotation = argspec.annotations[arg] + if not issubclass(type_annotation, Type): + raise Exception(f"Type of tree input '{arg}' is not a valid 'Type' subclass.") + inputs[arg] = type_annotation + + # Create the input sockets and collect input values. + builder_inputs = [] + for i, arg in enumerate(inputs.items()): + node_group.inputs.new(arg[1].socket_type, re.sub('([A-Z])', r' \1', arg[0]).title()) + builder_inputs.append(arg[1](group_input_node.outputs[i])) + + # Run the builder function + State.current_node_tree = node_group + outputs = builder(*builder_inputs) + + # Create the output sockets + 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.") + node_group.outputs.new(result.socket_type, 'Result') + link = node_group.links.new(result._socket, group_output_node.inputs[i]) + + # Attempt to run the "Node Arrange" add-on on the tree. + try: + for area in bpy.context.screen.areas: + for space in area.spaces: + if space.type == 'NODE_EDITOR': + space.node_tree = node_group + with bpy.context.temp_override(area=area, space=space, space_data=space): + ntree = node_group + ntree.nodes[0].select = True + ntree.nodes.active = ntree.nodes[0] + n_groups = [] + for i in ntree.nodes: + if i.type == 'GROUP': + n_groups.append(i) + + while n_groups: + j = n_groups.pop(0) + node_arrange.nodes_iterate(j.node_tree) + for i in j.node_tree.nodes: + if i.type == 'GROUP': + n_groups.append(i) + + node_arrange.nodes_iterate(ntree) + + # arrange nodes + this center nodes together + if bpy.context.scene.node_center: + node_arrange.nodes_center(ntree) + except: + pass + + # 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): + return group(node_tree=node_group, *args, **kwargs) + return group_reference + if isinstance(name, str): + return build_tree + else: + tree_name = name.__name__ + return build_tree(name) \ No newline at end of file diff --git a/lib/types.py b/lib/types.py new file mode 100644 index 0000000..581f7b4 --- /dev/null +++ b/lib/types.py @@ -0,0 +1,53 @@ +import bpy +from .state import State + +# The base class all exposed socket types conform to. +class Type: + socket_type: str + + def __init__(self, socket: bpy.types.NodeSocket = None, value = None): + if value is not None: + input_nodes = { + int: ('FunctionNodeInputInt', 'integer'), + bool: ('FunctionNodeInputBool', 'boolean'), + str: ('FunctionNodeInputString', 'string'), + tuple: ('FunctionNodeInputVector', 'vector'), + float: ('ShaderNodeValue', None), + } + if not type(value) in input_nodes: + raise Exception(f"'{value}' cannot be expressed as a node.") + input_node_info = input_nodes[type(value)] + value_node = State.current_node_tree.nodes.new(input_node_info[0]) + if input_node_info[1] is None: + value_node.outputs[0].default_value = value + else: + setattr(value_node, input_node_info[1], value) + socket = value_node.outputs[0] + self._socket = socket + self.socket_type = type(socket).__name__ + + def _math(self, other, operation): + math_node = State.current_node_tree.nodes.new('ShaderNodeVectorMath' if self._socket.type else 'ShaderNodeMath') + math_node.operation = operation + State.current_node_tree.links.new(self._socket, math_node.inputs[0]) + if issubclass(type(other), Type): + State.current_node_tree.links.new(other._socket, math_node.inputs[1]) + else: + math_node.inputs[1].default_value = other + return Type(math_node.outputs[0]) + + def __add__(self, other): + return self._math(other, 'ADD') + + def __sub__(self, other): + return self._math(other, 'SUBTRACT') + + def __mul__(self, other): + return self._math(other, 'SUBTRACT') + + def __truediv__(self, other): + return self._math(other, 'DIVIDE') + +for standard_socket in bpy.types.NodeSocketStandard.__subclasses__(): + name = standard_socket.__name__.replace('NodeSocket', '') + globals()[name] = type(name, (Type,), { 'socket_type': standard_socket.__name__ }) \ No newline at end of file