200 lines
7.1 KiB
Python
200 lines
7.1 KiB
Python
import bpy
|
|
import mathutils
|
|
|
|
from .utils import global_loc, calc_node
|
|
import gpu
|
|
from gpu_extras.presets import draw_circle_2d
|
|
from gpu_extras.batch import batch_for_shader
|
|
|
|
DRAW_COLOR = (1, 1, 1, 0.5)
|
|
DRAW_RADIUS = 10
|
|
|
|
|
|
def draw_callback(self):
|
|
if self.drag_mode:
|
|
if self.dragging_node:
|
|
node = self.dragging_node
|
|
loc = global_loc(node)
|
|
|
|
x1, y1 = loc[0] - DRAW_RADIUS, loc[1] + DRAW_RADIUS
|
|
x2, y2 = loc[0] + node.dimensions[0] + DRAW_RADIUS, loc[1] - node.dimensions[1] - DRAW_RADIUS
|
|
|
|
shader = gpu.shader.from_builtin('2D_SMOOTH_COLOR')
|
|
|
|
vertices = ((x1, y1), (x2, y1), (x2, y2), (x1, y2), (x1, y1))
|
|
vertex_colors = (DRAW_COLOR, DRAW_COLOR, DRAW_COLOR, DRAW_COLOR, DRAW_COLOR)
|
|
|
|
batch = batch_for_shader(shader, 'LINE_STRIP', {"pos": vertices, "color": vertex_colors})
|
|
|
|
shader.bind()
|
|
batch.draw(shader)
|
|
else:
|
|
draw_circle_2d(self.cursor_pos, DRAW_COLOR, self.radius)
|
|
|
|
|
|
class NodeRelaxBrush(bpy.types.Operator):
|
|
"""Relax Nodes"""
|
|
bl_idname = "node_relax.brush"
|
|
bl_label = "Relax Nodes"
|
|
|
|
bl_options = {"UNDO", "REGISTER"}
|
|
|
|
@classmethod
|
|
def poll(cls, context):
|
|
space = context.space_data
|
|
return space.type == 'NODE_EDITOR' and space.node_tree is not None
|
|
|
|
def update_cursor_pos(self, context, event):
|
|
self.cursor_prev_pos = self.cursor_pos
|
|
self.cursor_pos = mathutils.Vector(
|
|
context.region.view2d.region_to_view(event.mouse_region_x, event.mouse_region_y))
|
|
|
|
def update_radius(self, context, original_radius):
|
|
radiusM = context.region.view2d.region_to_view(original_radius, 0)
|
|
radius0 = context.region.view2d.region_to_view(0, 0)
|
|
self.radius = radiusM[0] - radius0[0]
|
|
|
|
def get_brush_influence(self, loc, size):
|
|
self.delta.x = self.cursor_pos.x - min(max(self.cursor_pos.x, loc.x), loc.x + size.x)
|
|
self.delta.y = self.cursor_pos.y - max(min(self.cursor_pos.y, loc.y), loc.y - size.y)
|
|
|
|
dist_sqr = self.delta.x * self.delta.x + self.delta.y * self.delta.y
|
|
|
|
return 1 - (dist_sqr / (self.radius * self.radius))
|
|
|
|
def main_operation(self, context):
|
|
nodes = self.tree.nodes
|
|
props = context.scene.NodeRelax_props
|
|
|
|
self.slide_vec = self.cursor_pos - self.cursor_prev_pos
|
|
context.area.tag_redraw()
|
|
|
|
if self.drag_mode:
|
|
if self.is_dragging:
|
|
if self.dragging_node:
|
|
self.dragging_node.location += self.slide_vec
|
|
else:
|
|
self.update_dragging_node(nodes)
|
|
else:
|
|
if self.lmb:
|
|
dist = mathutils.Vector((props.Distance, props.Distance))
|
|
for node in nodes:
|
|
if node.type == 'FRAME':
|
|
continue
|
|
# Brush
|
|
loc = global_loc(node)
|
|
size = node.dimensions
|
|
infl = self.get_brush_influence(loc, size)
|
|
if infl <= 0:
|
|
continue
|
|
|
|
# Calculate physics
|
|
calc_node(node, nodes, infl, self.slide_vec * props.SlidePower, props.RelaxPower,
|
|
props.CollisionPower, dist, False)
|
|
|
|
def update_dragging_node(self, nodes):
|
|
self.dragging_node = None
|
|
nearest = 0
|
|
for node in nodes:
|
|
if node.type == 'FRAME':
|
|
continue
|
|
loc = global_loc(node)
|
|
pos = mathutils.Vector((loc.x, loc.y))
|
|
pos.x += node.dimensions.x / 2
|
|
pos.y -= node.dimensions.y / 2
|
|
pos -= self.cursor_pos
|
|
dist = pos.x * pos.x + pos.y * pos.y # Squared length
|
|
if self.dragging_node is None or dist < nearest:
|
|
self.dragging_node = node
|
|
nearest = dist
|
|
|
|
def finish(self, context, props):
|
|
st = bpy.types.SpaceNodeEditor
|
|
st.draw_handler_remove(self.draw_handler, 'WINDOW')
|
|
props.IsRunning = False
|
|
|
|
def modal(self, context, event):
|
|
props = context.scene.NodeRelax_props
|
|
|
|
if event.type in {'MIDDLEMOUSE', 'WHEELUPMOUSE', 'WHEELDOWNMOUSE'}: # allow navigation shortcuts
|
|
return {'PASS_THROUGH'}
|
|
|
|
if event.type == 'LEFT_SHIFT' or event.type == 'RIGHT_SHIFT': # drag individual node
|
|
if event.value == 'PRESS':
|
|
self.drag_mode = True
|
|
if event.value == 'RELEASE':
|
|
self.drag_mode = False
|
|
self.is_dragging = False
|
|
self.update_dragging_node(self.tree.nodes)
|
|
context.area.tag_redraw()
|
|
|
|
if event.type == 'MOUSEMOVE':
|
|
self.update_cursor_pos(context, event)
|
|
self.update_radius(context, props.BrushSize)
|
|
self.main_operation(context)
|
|
|
|
if event.type == 'WHEELUPMOUSE' or event.type == 'WHEELDOWNMOUSE':
|
|
self.update_cursor_pos(context, event)
|
|
self.update_radius(context, props.BrushSize)
|
|
context.area.tag_redraw()
|
|
|
|
elif event.type == 'LEFTMOUSE':
|
|
if event.value == 'PRESS':
|
|
self.lmb = True
|
|
if self.drag_mode:
|
|
self.is_dragging = True
|
|
else:
|
|
self.update_cursor_pos(context, event)
|
|
self.cursor_prev_pos = self.cursor_pos # No sliding
|
|
self.main_operation(context)
|
|
if event.value == 'RELEASE':
|
|
self.lmb = False
|
|
if self.drag_mode:
|
|
self.is_dragging = False
|
|
|
|
elif event.type in {'RIGHTMOUSE', 'ESC'}:
|
|
self.finish(context, props)
|
|
context.area.tag_redraw()
|
|
return {'FINISHED'}
|
|
|
|
if event.type == "LEFT_BRACKET":
|
|
props.BrushSize -= 10
|
|
props.BrushSize = max(props.BrushSize, 10)
|
|
self.update_radius(context, props.BrushSize)
|
|
context.area.tag_redraw()
|
|
|
|
elif event.type == "RIGHT_BRACKET":
|
|
props.BrushSize += 10
|
|
props.BrushSize = min(props.BrushSize, 1000)
|
|
self.update_radius(context, props.BrushSize)
|
|
context.area.tag_redraw()
|
|
|
|
return {'RUNNING_MODAL'}
|
|
|
|
def invoke(self, context, event):
|
|
props = context.scene.NodeRelax_props
|
|
if props.IsRunning:
|
|
return {'CANCELLED'}
|
|
|
|
self.tree = context.space_data.edit_tree
|
|
# Initialize runtime state here (avoid __init__ for Blender RNA types)
|
|
self.radius = 100
|
|
self.lmb = False
|
|
self.delta = mathutils.Vector((0, 0))
|
|
self.cursor_pos = mathutils.Vector((0, 0))
|
|
self.cursor_prev_pos = mathutils.Vector((0, 0))
|
|
self.slide_vec = mathutils.Vector((0, 0))
|
|
self.drag_mode = False
|
|
self.is_dragging = False
|
|
self.dragging_node = None
|
|
context.window_manager.modal_handler_add(self)
|
|
st = bpy.types.SpaceNodeEditor
|
|
self.draw_handler = st.draw_handler_add(draw_callback, (self,), 'WINDOW', 'POST_VIEW')
|
|
|
|
props.IsRunning = True
|
|
self.update_cursor_pos(context, event)
|
|
self.update_radius(context, props.BrushSize)
|
|
|
|
context.area.tag_redraw()
|
|
return {'RUNNING_MODAL'}
|