From 66ade0a7395f0382ccde993ae8160a0c1aecf25c Mon Sep 17 00:00:00 2001 From: Carson Katri Date: Sat, 19 Nov 2022 15:00:26 -0500 Subject: [PATCH] Support prefixes for InputGroups --- api/node_mapper.py | 5 ++- api/static/input_group.py | 11 ++++- api/tree.py | 12 ++++-- api/types.py | 89 +++++++++++++++++++++++++++------------ 4 files changed, 83 insertions(+), 34 deletions(-) diff --git a/api/node_mapper.py b/api/node_mapper.py index b914ec1..891e55d 100644 --- a/api/node_mapper.py +++ b/api/node_mapper.py @@ -2,6 +2,7 @@ import bpy import bl_ui import itertools import enum +import re from .state import State from .types import * from .static.input_group import InputGroup @@ -299,6 +300,8 @@ def create_documentation(): def enum_namespace(k): return f"""class {k}: {newline.join(enums[k])}""" + def add_self_arg(x): + return re.sub('\(', '(self, ', x, 1) contents = f"""from typing import * import enum def tree(builder): @@ -327,7 +330,7 @@ class Type: y = Type() z = Type() def capture(self, attribute: Type, **kwargs) -> Callable[[], Type]: return transfer_attribute - {(newline + ' ').join(map(lambda x: x.replace('(', '(self, '), filter(lambda x: x.startswith('def'), symbols)))} + {(newline + ' ').join(map(add_self_arg, filter(lambda x: x.startswith('def'), symbols)))} {newline.join(map(type_symbol, Type.__subclasses__()))} {newline.join(map(enum_namespace, enums.keys()))} diff --git a/api/static/input_group.py b/api/static/input_group.py index 0232349..74a079f 100644 --- a/api/static/input_group.py +++ b/api/static/input_group.py @@ -1,4 +1,13 @@ -class InputGroup: +class _InputGroupMeta(type): + def __getitem__(cls, args): + if isinstance(args, str): + class PrefixedInputGroup(InputGroup): + prefix = args + PrefixedInputGroup.__annotations__ = cls.__annotations__ + return PrefixedInputGroup + return cls + +class InputGroup(metaclass=_InputGroupMeta): """ A group of inputs that will be expanded in the node tree. diff --git a/api/tree.py b/api/tree.py index c6f5bd6..dbdbeec 100644 --- a/api/tree.py +++ b/api/tree.py @@ -47,11 +47,12 @@ def tree(name): raise Exception(f"Type of tree input '{param.name}' is not a valid 'Type' subclass.") for param in signature.parameters.values(): if issubclass(param.annotation, InputGroup): + prefix = (param.annotation.prefix + "_") if hasattr(param.annotation, "prefix") else "" for group_param, annotation in param.annotation.__annotations__.items(): - inputs[group_param] = (annotation, inspect.Parameter.empty, param.name) + inputs[prefix + group_param] = (annotation, inspect.Parameter.empty, param.name, prefix) else: validate_param(param) - inputs[param.name] = (param.annotation, param.default, None) + inputs[param.name] = (param.annotation, param.default, None, None) # Create the input sockets and collect input values. for i, node_input in enumerate(node_group.inputs): @@ -72,7 +73,7 @@ def tree(name): 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() - setattr(builder_inputs[arg[1][2]], arg[0], arg[1][0](group_input_node.outputs[i])) + setattr(builder_inputs[arg[1][2]], arg[0].replace(arg[1][3], ''), arg[1][0](group_input_node.outputs[i])) else: builder_inputs[arg[0]] = arg[1][0](group_input_node.outputs[i]) @@ -132,7 +133,10 @@ def tree(name): group_outputs = [] for group_output in result._socket.node.outputs: group_outputs.append(Type(group_output)) - return tuple(group_outputs) + if len(group_outputs) == 1: + return group_outputs[0] + else: + return tuple(group_outputs) return group_reference if isinstance(name, str): return build_tree diff --git a/api/types.py b/api/types.py index 73320ad..667d1f1 100644 --- a/api/types.py +++ b/api/types.py @@ -2,6 +2,7 @@ import bpy from bpy.types import NodeSocketStandard import nodeitems_utils from .state import State +import geometry_script def map_case_name(i): return ('_' if not i.identifier[0].isalpha() else '') + i.identifier.replace(' ', '_').upper() @@ -16,7 +17,16 @@ def socket_type_to_data_type(socket_type): return socket_type # The base class all exposed socket types conform to. -class Type: +class _TypeMeta(type): + def __getitem__(self, args): + for s in filter(lambda x: isinstance(x, slice), args): + if (isinstance(s.start, float) or isinstance(s.start, int)) and (isinstance(s.stop, float) or isinstance(s.stop, int)): + print(f"minmax: ({s.start}, {s.stop})") + elif isinstance(s.start, str): + print(f"{s.start} = {s.stop}") + return self + +class Type(metaclass=_TypeMeta): socket_type: str def __init__(self, socket: bpy.types.NodeSocket = None, value = None): @@ -41,14 +51,10 @@ class Type: self.socket_type = type(socket).__name__ def _math(self, other, operation, reverse=False): - math_node = State.current_node_tree.nodes.new('ShaderNodeVectorMath' if self._socket.type == 'VECTOR' else 'ShaderNodeMath') - math_node.operation = operation - State.current_node_tree.links.new(self._socket, math_node.inputs[1 if reverse else 0]) - if issubclass(type(other), Type): - State.current_node_tree.links.new(other._socket, math_node.inputs[0 if reverse else 1]) + if self._socket.type == 'VECTOR': + return geometry_script.vector_math(operation=operation, vector=(other, self) if reverse else (self, other)) else: - math_node.inputs[0 if reverse else 1].default_value = other - return Type(math_node.outputs[0]) + return geometry_script.math(operation=operation, value=(other, self) if reverse else (self, other)) def __add__(self, other): return self._math(other, 'ADD') @@ -81,30 +87,19 @@ class Type: return self._math(other, 'MODULO', True) def _compare(self, other, operation): - compare_node = State.current_node_tree.nodes.new('FunctionNodeCompare') - compare_node.data_type = 'FLOAT' if self._socket.type == 'VALUE' else self._socket.type - compare_node.operation = operation - a = None - b = None - for node_input in compare_node.inputs: - if not node_input.enabled: - continue - elif a is None: - a = node_input - else: - b = node_input - State.current_node_tree.links.new(self._socket, a) - if issubclass(type(other), Type): - State.current_node_tree.links.new(other._socket, b) - else: - b.default_value = other - return Type(compare_node.outputs[0]) + return geometry_script.compare(operation=operation, a=self, b=other) def __eq__(self, other): - return self._compare(other, 'EQUAL') + if self._socket.type == 'BOOLEAN': + return self._boolean_math(other, 'XNOR') + else: + return self._compare(other, 'EQUAL') def __ne__(self, other): - return self._compare(other, 'NOT_EQUAL') + if self._socket.type == 'BOOLEAN': + return self._boolean_math(other, 'XOR') + else: + return self._compare(other, 'NOT_EQUAL') def __lt__(self, other): return self._compare(other, 'LESS_THAN') @@ -118,6 +113,44 @@ class Type: def __ge__(self, other): return self._compare(other, 'GREATER_EQUAL') + def _boolean_math(self, other, operation, reverse=False): + boolean_math_node = State.current_node_tree.nodes.new('FunctionNodeBooleanMath') + boolean_math_node.operation = operation + a = None + b = None + for node_input in boolean_math_node.inputs: + if not node_input.enabled: + continue + elif a is None: + a = node_input + else: + b = node_input + State.current_node_tree.links.new(self._socket, a) + if other is not None: + if issubclass(type(other), Type): + State.current_node_tree.links.new(other._socket, b) + else: + b.default_value = other + return Type(boolean_math_node.outputs[0]) + + def __and__(self, other): + return self._boolean_math(other, 'AND') + + def __rand__(self, other): + return self._boolean_math(other, 'AND', reverse=True) + + def __or__(self, other): + return self._boolean_math(other, 'OR') + + def __ror__(self, other): + return self._boolean_math(other, 'OR', reverse=True) + + def __invert__(self): + if self._socket.type == 'BOOLEAN': + return self._boolean_math(None, 'NOT') + else: + return self._math(-1, 'MULTIPLY') + def _get_xyz_component(self, component): if self._socket.type != 'VECTOR': raise Exception("`x`, `y`, `z` properties are not available on non-Vector types.")