kopia lustrzana https://github.com/carson-katri/geometry-script
381 wiersze
14 KiB
Python
381 wiersze
14 KiB
Python
import bpy
|
|
from bpy.types import NodeSocketStandard
|
|
import nodeitems_utils
|
|
import enum
|
|
from .state import State
|
|
from .static.sample_mode import SampleMode
|
|
import geometry_script
|
|
|
|
INT_MAX = 2147483647
|
|
INT_MIN = -INT_MAX -1
|
|
|
|
class SubtypeInt(enum.Enum):
|
|
NONE = 'None'
|
|
PERCENTAGE = 'Percentage'
|
|
FACTOR = 'Factor'
|
|
|
|
class SubtypeFloat(enum.Enum):
|
|
NONE = 'None'
|
|
PERCENTAGE = 'Percentage'
|
|
FACTOR = 'Factor'
|
|
ANGLE = 'Angle'
|
|
TIME = 'Time (Scene Relative)'
|
|
TIME_ABSOLUTE = 'Time (Absolute)'
|
|
DISTANCE = 'Distance'
|
|
|
|
class InputOptions:
|
|
min_value: int | float
|
|
max_value: int | float
|
|
bl_subtype_label: str
|
|
name: str
|
|
description: str
|
|
hide_in_modifier: bool
|
|
|
|
def __init__(
|
|
self,
|
|
min: int | float | None = None,
|
|
max: int | float | None = None,
|
|
subtype: SubtypeInt | SubtypeFloat | None = None,
|
|
name: str | None = None,
|
|
tooltip: str = '',
|
|
hide_in_modifier: bool = False
|
|
):
|
|
self.min_value = min
|
|
self.max_value = max
|
|
self.bl_subtype_label = subtype.value if subtype != None else None
|
|
self.name = name
|
|
self.description = tooltip
|
|
self.hide_in_modifier = hide_in_modifier
|
|
|
|
def process(self, node_input_type: str):
|
|
if node_input_type == 'INT':
|
|
if self.min_value != None and self.max_value == None:
|
|
self.max_value = INT_MAX
|
|
if self.max_value != None and self.min_value == None:
|
|
self.min_value = INT_MIN
|
|
if isinstance(self.min_value, float):
|
|
self.min_value = int(self.min_value)
|
|
if isinstance(self.max_value, float):
|
|
self.max_value = int(self.max_value)
|
|
if self.bl_subtype_label == None:
|
|
self.bl_subtype_label = SubtypeInt.NONE
|
|
elif node_input_type == 'VALUE':
|
|
if self.min_value != None and self.max_value == None:
|
|
self.max_value = float('inf')
|
|
if self.max_value != None and self.min_value == None:
|
|
self.min_value = float('-inf')
|
|
if isinstance(self.min_value, int):
|
|
self.min_value = float(self.min_value)
|
|
if isinstance(self.max_value, int):
|
|
self.max_value = float(self.max_value)
|
|
if self.bl_subtype_label == None:
|
|
self.bl_subtype_label = SubtypeFloat.NONE
|
|
|
|
def map_case_name(i):
|
|
return ('_' if not i.identifier[0].isalpha() else '') + i.identifier.replace(' ', '_').upper()
|
|
|
|
def socket_type_to_data_type(socket_type):
|
|
match socket_type:
|
|
case 'VALUE':
|
|
return 'FLOAT'
|
|
case 'VECTOR':
|
|
return 'FLOAT_VECTOR'
|
|
case 'COLOR':
|
|
return 'FLOAT_COLOR'
|
|
case _:
|
|
return socket_type
|
|
|
|
def socket_class_to_data_type(socket_class_name):
|
|
match socket_class_name:
|
|
case 'NodeSocketGeometry':
|
|
return 'GEOMETRY'
|
|
case 'NodeSocketFloat':
|
|
return 'FLOAT'
|
|
case _:
|
|
return socket_class_name
|
|
|
|
# The base class all exposed socket types conform to.
|
|
class _TypeMeta(type):
|
|
def __getitem__(self, args):
|
|
input_options = None
|
|
if isinstance(args, int) or isinstance(args, float):
|
|
input_options = InputOptions(min=args)
|
|
elif isinstance(args, tuple):
|
|
tuple_args = {}
|
|
if isinstance(args[0], int) or isinstance(args[0], float):
|
|
tuple_args['min'] = args[0]
|
|
if len(args) > 1 and (isinstance(args[1], int) or isinstance(args[1], float)):
|
|
tuple_args['max'] = args[1]
|
|
if len(tuple_args) > 0:
|
|
input_options = InputOptions(**tuple_args)
|
|
elif isinstance(args, slice):
|
|
slice_args = {}
|
|
if isinstance(args.start, int) or isinstance(args.start, float):
|
|
slice_args['min'] = args.start
|
|
if isinstance(args.stop, int) or isinstance(args.stop, float):
|
|
slice_args['max'] = args.stop
|
|
if len(slice_args) > 0:
|
|
input_options = InputOptions(**slice_args)
|
|
elif isinstance(args, InputOptions):
|
|
input_options = args
|
|
|
|
if input_options != None:
|
|
setattr(self, 'input_options', input_options)
|
|
|
|
return self
|
|
|
|
class Type(metaclass=_TypeMeta):
|
|
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 type(value) == int:
|
|
print("Making an integer node?")
|
|
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, reverse=False):
|
|
if other is None:
|
|
vector_or_value = self
|
|
else:
|
|
vector_or_value = (other, self) if reverse else (self, other)
|
|
|
|
if self._socket.type == 'VECTOR':
|
|
return geometry_script.vector_math(operation=operation, vector=vector_or_value)
|
|
else:
|
|
return geometry_script.math(operation=operation, value=vector_or_value)
|
|
|
|
def __add__(self, other):
|
|
return self._math(other, 'ADD')
|
|
|
|
def __radd__(self, other):
|
|
return self._math(other, 'ADD', True)
|
|
|
|
def __sub__(self, other):
|
|
return self._math(other, 'SUBTRACT')
|
|
|
|
def __rsub__(self, other):
|
|
return self._math(other, 'SUBTRACT', True)
|
|
|
|
def __mul__(self, other):
|
|
return self._math(other, 'MULTIPLY')
|
|
|
|
def __rmul__(self, other):
|
|
return self._math(other, 'MULTIPLY', True)
|
|
|
|
def __truediv__(self, other):
|
|
return self._math(other, 'DIVIDE')
|
|
|
|
def __rtruediv__(self, other):
|
|
return self._math(other, 'DIVIDE', True)
|
|
|
|
def __mod__(self, other):
|
|
return self._math(other, 'MODULO')
|
|
|
|
def __rmod__(self, other):
|
|
return self._math(other, 'MODULO', True)
|
|
|
|
def __floordiv__(self, other):
|
|
return self._math(other, 'DIVIDE')._math(None,'FLOOR')
|
|
|
|
def __rfloordiv__(self, other):
|
|
return self._math(other, 'DIVIDE',True)._math(None,'FLOOR')
|
|
|
|
def __pow__(self, other):
|
|
return self._math(other, 'POWER')
|
|
|
|
def __rpow__(self, other):
|
|
return self._math(other, 'POWER', True)
|
|
|
|
def __matmul__(self, other):
|
|
return self._math(other, 'DOT_PRODUCT')
|
|
|
|
def __rmatmul__(self, other):
|
|
return self._math(other, 'DOT_PRODUCT', True)
|
|
|
|
def __abs__(self):
|
|
return self._math(None,'ABSOLUTE')
|
|
|
|
def __neg__(self):
|
|
return self._math(-1, 'MULTIPLY')
|
|
|
|
def __pos__(self):
|
|
return self
|
|
|
|
def __round__(self):
|
|
return self._math(None,'ROUND')
|
|
|
|
def _compare(self, other, operation):
|
|
return geometry_script.compare(operation=operation, a=self, b=other)
|
|
|
|
def __eq__(self, other):
|
|
if self._socket.type == 'BOOLEAN':
|
|
return self._boolean_math(other, 'XNOR')
|
|
else:
|
|
return self._compare(other, 'EQUAL')
|
|
|
|
def __ne__(self, other):
|
|
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')
|
|
|
|
def __le__(self, other):
|
|
return self._compare(other, 'LESS_EQUAL')
|
|
|
|
def __gt__(self, other):
|
|
return self._compare(other, 'GREATER_THAN')
|
|
|
|
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, -1, -1) if self._socket.type == 'VECTOR' else -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.")
|
|
separate_node = State.current_node_tree.nodes.new('ShaderNodeSeparateXYZ')
|
|
State.current_node_tree.links.new(self._socket, separate_node.inputs[0])
|
|
return Type(separate_node.outputs[component])
|
|
@property
|
|
def x(self):
|
|
return self._get_xyz_component(0)
|
|
@property
|
|
def y(self):
|
|
return self._get_xyz_component(1)
|
|
@property
|
|
def z(self):
|
|
return self._get_xyz_component(2)
|
|
|
|
def capture(self, value, **kwargs):
|
|
data_type = socket_type_to_data_type(value._socket.type)
|
|
res = self.capture_attribute(data_type=data_type, value=value, **kwargs)
|
|
return res.geometry, res.attribute
|
|
def transfer(self, attribute, **kwargs):
|
|
data_type = socket_type_to_data_type(attribute._socket.type)
|
|
return self.transfer_attribute(data_type=data_type, attribute=attribute, **kwargs)
|
|
|
|
def __getitem__(self, subscript):
|
|
if self._socket.type == 'VECTOR' and isinstance(subscript, int):
|
|
return self._get_xyz_component(subscript)
|
|
if isinstance(subscript, tuple):
|
|
accessor = subscript[0]
|
|
args = subscript[1:]
|
|
else:
|
|
accessor = subscript
|
|
args = []
|
|
sample_mode = SampleMode.INDEX if len(args) < 1 else args[0]
|
|
domain = 'POINT' if len(args) < 2 else (args[1].value if isinstance(args[1], enum.Enum) else args[1])
|
|
sample_position = None
|
|
sampling_index = None
|
|
if isinstance(accessor, slice):
|
|
data_type = socket_type_to_data_type(accessor.start._socket.type)
|
|
value = accessor.start
|
|
match sample_mode:
|
|
case SampleMode.INDEX:
|
|
sampling_index = accessor.stop
|
|
case SampleMode.NEAREST_SURFACE:
|
|
sample_position = accessor.stop
|
|
case SampleMode.NEAREST:
|
|
sample_position = accessor.stop
|
|
if accessor.step is not None:
|
|
domain = accessor.step.value if isinstance(accessor.step, enum.Enum) else accessor.step
|
|
else:
|
|
data_type = socket_type_to_data_type(accessor._socket.type)
|
|
value = accessor
|
|
match sample_mode:
|
|
case SampleMode.INDEX:
|
|
return self.sample_index(
|
|
data_type=data_type,
|
|
domain=domain,
|
|
value=value,
|
|
index=sampling_index or geometry_script.index()
|
|
)
|
|
case SampleMode.NEAREST_SURFACE:
|
|
return self.sample_nearest_surface(
|
|
data_type=data_type,
|
|
value=value,
|
|
sample_position=sample_position or geometry_script.position()
|
|
)
|
|
case SampleMode.NEAREST:
|
|
return self.sample_index(
|
|
data_type=data_type,
|
|
value=value,
|
|
index=self.sample_nearest(domain=domain, sample_position=sample_position or geometry_script.position())
|
|
)
|
|
|
|
for standard_socket in list(filter(lambda x: 'NodeSocket' in x, dir(bpy.types))):
|
|
name = standard_socket.replace('NodeSocket', '')
|
|
if len(name) < 1:
|
|
continue
|
|
globals()[name] = type(name, (Type,), { 'socket_type': standard_socket, '__module__': Type.__module__ })
|
|
if name == 'Int':
|
|
class IntIterator:
|
|
def __init__(self, integer):
|
|
self.integer = integer
|
|
self.points = State.current_node_tree.nodes.new('GeometryNodePoints')
|
|
State.current_node_tree.links.new(self.integer._socket, self.points.inputs[0])
|
|
self.index = State.current_node_tree.nodes.new('GeometryNodeInputIndex')
|
|
self._did_iterate = False
|
|
def __next__(self):
|
|
if not self._did_iterate:
|
|
self._did_iterate = True
|
|
return Type(self.index.outputs[0]), Type(self.points.outputs[0])
|
|
else:
|
|
raise StopIteration()
|
|
globals()[name].__iter__ = lambda self: IntIterator(self) |