Skip to content
Snippets Groups Projects
node_wrangler.py 210 KiB
Newer Older
# ##### BEGIN GPL LICENSE BLOCK #####
#
#  This program is free software; you can redistribute it and/or
#  modify it under the terms of the GNU General Public License
#  as published by the Free Software Foundation; either version 2
#  of the License, or (at your option) any later version.
#
#  This program is distributed in the hope that it will be useful,
#  but WITHOUT ANY WARRANTY; without even the implied warranty of
#  MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
#  GNU General Public License for more details.
#
#  You should have received a copy of the GNU General Public License
#  along with this program; if not, write to the Free Software Foundation,
#  Inc., 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301, USA.
#
# ##### END GPL LICENSE BLOCK #####

bl_info = {
    "name": "Node Wrangler",
florianfelix's avatar
florianfelix committed
    "author": "Bartek Skorupa, Greg Zaal, Sebastian Koenig, Christian Brinkmann, Florian Meyer",
Greg Zaal's avatar
Greg Zaal committed
    "version": (3, 37),
    "blender": (2, 83, 0),
    "location": "Node Editor Toolbar or Shift-W",
    "description": "Various tools to enhance and speed up node-based workflow",
    "warning": "",
    "doc_url": "{BLENDER_MANUAL_URL}/addons/node/node_wrangler.html",
from bpy.types import Operator, Panel, Menu
Campbell Barton's avatar
Campbell Barton committed
from bpy.props import (
    FloatProperty,
    EnumProperty,
    BoolProperty,
    IntProperty,
    StringProperty,
    FloatVectorProperty,
    CollectionProperty,
)
from bpy_extras.io_utils import ImportHelper, ExportHelper
from gpu_extras.batch import batch_for_shader
Bartek Skorupa's avatar
Bartek Skorupa committed
from mathutils import Vector
from math import cos, sin, pi, hypot
from glob import glob
from itertools import chain
from collections import namedtuple

#################
# rl_outputs:
# list of outputs of Input Render Layer
# with attributes determinig if pass is used,
# and MultiLayer EXR outputs names and corresponding render engines
#
# rl_outputs entry = (render_pass, rl_output_name, exr_output_name, in_eevee, in_cycles)
RL_entry = namedtuple('RL_Entry', ['render_pass', 'output_name', 'exr_output_name', 'in_eevee', 'in_cycles'])
    RL_entry('use_pass_ambient_occlusion', 'AO', 'AO', True, True),
    RL_entry('use_pass_combined', 'Image', 'Combined', True, True),
    RL_entry('use_pass_diffuse_color', 'Diffuse Color', 'DiffCol', False, True),
    RL_entry('use_pass_diffuse_direct', 'Diffuse Direct', 'DiffDir', False, True),
    RL_entry('use_pass_diffuse_indirect', 'Diffuse Indirect', 'DiffInd', False, True),
    RL_entry('use_pass_emit', 'Emit', 'Emit', False, True),
    RL_entry('use_pass_environment', 'Environment', 'Env', False, False),
    RL_entry('use_pass_glossy_color', 'Glossy Color', 'GlossCol', False, True),
    RL_entry('use_pass_glossy_direct', 'Glossy Direct', 'GlossDir', False, True),
    RL_entry('use_pass_glossy_indirect', 'Glossy Indirect', 'GlossInd', False, True),
    RL_entry('use_pass_indirect', 'Indirect', 'Indirect', False, False),
    RL_entry('use_pass_material_index', 'IndexMA', 'IndexMA', False, True),
    RL_entry('use_pass_mist', 'Mist', 'Mist', True, True),
    RL_entry('use_pass_normal', 'Normal', 'Normal', True, True),
    RL_entry('use_pass_object_index', 'IndexOB', 'IndexOB', False, True),
    RL_entry('use_pass_shadow', 'Shadow', 'Shadow', False, True),
    RL_entry('use_pass_subsurface_color', 'Subsurface Color', 'SubsurfaceCol', True, True),
    RL_entry('use_pass_subsurface_direct', 'Subsurface Direct', 'SubsurfaceDir', True, True),
    RL_entry('use_pass_subsurface_indirect', 'Subsurface Indirect', 'SubsurfaceInd', False, True),
    RL_entry('use_pass_transmission_color', 'Transmission Color', 'TransCol', False, True),
    RL_entry('use_pass_transmission_direct', 'Transmission Direct', 'TransDir', False, True),
    RL_entry('use_pass_transmission_indirect', 'Transmission Indirect', 'TransInd', False, True),
    RL_entry('use_pass_uv', 'UV', 'UV', True, True),
    RL_entry('use_pass_vector', 'Speed', 'Vector', False, True),
    RL_entry('use_pass_z', 'Z', 'Depth', True, True),
    )

# shader nodes
# (rna_type.identifier, type, rna_type.name)
# Keeping mixed case to avoid having to translate entries when adding new nodes in operators.
# Keeping things in alphabetical orde so we don't need to sort later.
shaders_input_nodes_props = (
    ('ShaderNodeAmbientOcclusion', 'AMBIENT_OCCLUSION', 'Ambient Occlusion'),
    ('ShaderNodeAttribute', 'ATTRIBUTE', 'Attribute'),
    ('ShaderNodeBevel', 'BEVEL', 'Bevel'),
    ('ShaderNodeCameraData', 'CAMERA', 'Camera Data'),
    ('ShaderNodeFresnel', 'FRESNEL', 'Fresnel'),
    ('ShaderNodeNewGeometry', 'NEW_GEOMETRY', 'Geometry'),
    ('ShaderNodeHairInfo', 'HAIR_INFO', 'Hair Info'),
    ('ShaderNodeLayerWeight', 'LAYER_WEIGHT', 'Layer Weight'),
    ('ShaderNodeLightPath', 'LIGHT_PATH', 'Light Path'),
    ('ShaderNodeObjectInfo', 'OBJECT_INFO', 'Object Info'),
    ('ShaderNodeParticleInfo', 'PARTICLE_INFO', 'Particle Info'),
    ('ShaderNodeRGB', 'RGB', 'RGB'),
    ('ShaderNodeTangent', 'TANGENT', 'Tangent'),
    ('ShaderNodeTexCoord', 'TEX_COORD', 'Texture Coordinate'),
    ('ShaderNodeUVMap', 'UVMAP', 'UV Map'),
    ('ShaderNodeValue', 'VALUE', 'Value'),
    ('ShaderNodeVertexColor', 'VERTEX_COLOR', 'Vertex Color'),
    ('ShaderNodeVolumeInfo', 'VOLUME_INFO', 'Volume Info'),
    ('ShaderNodeWireframe', 'WIREFRAME', 'Wireframe'),

)
# (rna_type.identifier, type, rna_type.name)
# Keeping mixed case to avoid having to translate entries when adding new nodes in operators.
# Keeping things in alphabetical orde so we don't need to sort later.
shaders_output_nodes_props = (
    ('ShaderNodeOutputAOV', 'OUTPUT_AOV', 'AOV Output'),
    ('ShaderNodeOutputLight', 'OUTPUT_LIGHT', 'Light Output'),
    ('ShaderNodeOutputMaterial', 'OUTPUT_MATERIAL', 'Material Output'),
    ('ShaderNodeOutputWorld', 'OUTPUT_WORLD', 'World Output'),
)
# (rna_type.identifier, type, rna_type.name)
# Keeping mixed case to avoid having to translate entries when adding new nodes in operators.
# Keeping things in alphabetical orde so we don't need to sort later.
shaders_shader_nodes_props = (
    ('ShaderNodeAddShader', 'ADD_SHADER', 'Add Shader'),
    ('ShaderNodeBsdfAnisotropic', 'BSDF_ANISOTROPIC', 'Anisotropic BSDF'),
    ('ShaderNodeBsdfDiffuse', 'BSDF_DIFFUSE', 'Diffuse BSDF'),
    ('ShaderNodeEmission', 'EMISSION', 'Emission'),
    ('ShaderNodeBsdfGlass', 'BSDF_GLASS', 'Glass BSDF'),
    ('ShaderNodeBsdfGlossy', 'BSDF_GLOSSY', 'Glossy BSDF'),
    ('ShaderNodeBsdfHair', 'BSDF_HAIR', 'Hair BSDF'),
    ('ShaderNodeHoldout', 'HOLDOUT', 'Holdout'),
    ('ShaderNodeMixShader', 'MIX_SHADER', 'Mix Shader'),
    ('ShaderNodeBsdfPrincipled', 'BSDF_PRINCIPLED', 'Principled BSDF'),
    ('ShaderNodeBsdfHairPrincipled', 'BSDF_HAIR_PRINCIPLED', 'Principled Hair BSDF'),
    ('ShaderNodeVolumePrincipled', 'PRINCIPLED_VOLUME', 'Principled Volume'),
    ('ShaderNodeBsdfRefraction', 'BSDF_REFRACTION', 'Refraction BSDF'),
    ('ShaderNodeSubsurfaceScattering', 'SUBSURFACE_SCATTERING', 'Subsurface Scattering'),
    ('ShaderNodeBsdfToon', 'BSDF_TOON', 'Toon BSDF'),
    ('ShaderNodeBsdfTranslucent', 'BSDF_TRANSLUCENT', 'Translucent BSDF'),
    ('ShaderNodeBsdfTransparent', 'BSDF_TRANSPARENT', 'Transparent BSDF'),
    ('ShaderNodeBsdfVelvet', 'BSDF_VELVET', 'Velvet BSDF'),
    ('ShaderNodeBackground', 'BACKGROUND', 'Background'),
    ('ShaderNodeVolumeAbsorption', 'VOLUME_ABSORPTION', 'Volume Absorption'),
    ('ShaderNodeVolumeScatter', 'VOLUME_SCATTER', 'Volume Scatter'),
)
# (rna_type.identifier, type, rna_type.name)
# Keeping things in alphabetical orde so we don't need to sort later.
# Keeping mixed case to avoid having to translate entries when adding new nodes in operators.
shaders_texture_nodes_props = (
    ('ShaderNodeTexBrick', 'TEX_BRICK', 'Brick Texture'),
    ('ShaderNodeTexChecker', 'TEX_CHECKER', 'Checker Texture'),
    ('ShaderNodeTexEnvironment', 'TEX_ENVIRONMENT', 'Environment Texture'),
    ('ShaderNodeTexGradient', 'TEX_GRADIENT', 'Gradient Texture'),
    ('ShaderNodeTexIES', 'TEX_IES', 'IES Texture'),
    ('ShaderNodeTexImage', 'TEX_IMAGE', 'Image Texture'),
    ('ShaderNodeTexMagic', 'TEX_MAGIC', 'Magic Texture'),
    ('ShaderNodeTexMusgrave', 'TEX_MUSGRAVE', 'Musgrave Texture'),
    ('ShaderNodeTexNoise', 'TEX_NOISE', 'Noise Texture'),
    ('ShaderNodeTexPointDensity', 'TEX_POINTDENSITY', 'Point Density'),
    ('ShaderNodeTexSky', 'TEX_SKY', 'Sky Texture'),
    ('ShaderNodeTexVoronoi', 'TEX_VORONOI', 'Voronoi Texture'),
    ('ShaderNodeTexWave', 'TEX_WAVE', 'Wave Texture'),
    ('ShaderNodeTexWhiteNoise', 'TEX_WHITE_NOISE', 'White Noise'),
)
# (rna_type.identifier, type, rna_type.name)
# Keeping mixed case to avoid having to translate entries when adding new nodes in operators.
# Keeping things in alphabetical orde so we don't need to sort later.
shaders_color_nodes_props = (
    ('ShaderNodeBrightContrast', 'BRIGHTCONTRAST', 'Bright Contrast'),
    ('ShaderNodeGamma', 'GAMMA', 'Gamma'),
    ('ShaderNodeHueSaturation', 'HUE_SAT', 'Hue/Saturation'),
    ('ShaderNodeInvert', 'INVERT', 'Invert'),
    ('ShaderNodeLightFalloff', 'LIGHT_FALLOFF', 'Light Falloff'),
    ('ShaderNodeMixRGB', 'MIX_RGB', 'MixRGB'),
    ('ShaderNodeRGBCurve', 'CURVE_RGB', 'RGB Curves'),
)
# (rna_type.identifier, type, rna_type.name)
# Keeping mixed case to avoid having to translate entries when adding new nodes in operators.
# Keeping things in alphabetical orde so we don't need to sort later.
shaders_vector_nodes_props = (
    ('ShaderNodeBump', 'BUMP', 'Bump'),
    ('ShaderNodeDisplacement', 'DISPLACEMENT', 'Displacement'),
    ('ShaderNodeMapping', 'MAPPING', 'Mapping'),
    ('ShaderNodeNormal', 'NORMAL', 'Normal'),
    ('ShaderNodeNormalMap', 'NORMAL_MAP', 'Normal Map'),
    ('ShaderNodeVectorCurve', 'CURVE_VEC', 'Vector Curves'),
    ('ShaderNodeVectorDisplacement', 'VECTOR_DISPLACEMENT', 'Vector Displacement'),
    ('ShaderNodeVectorTransform', 'VECT_TRANSFORM', 'Vector Transform'),
)
# (rna_type.identifier, type, rna_type.name)
# Keeping mixed case to avoid having to translate entries when adding new nodes in operators.
# Keeping things in alphabetical orde so we don't need to sort later.
shaders_converter_nodes_props = (
    ('ShaderNodeBlackbody', 'BLACKBODY', 'Blackbody'),
    ('ShaderNodeClamp', 'CLAMP', 'Clamp'),
    ('ShaderNodeValToRGB', 'VALTORGB', 'ColorRamp'),
    ('ShaderNodeCombineHSV', 'COMBHSV', 'Combine HSV'),
    ('ShaderNodeCombineRGB', 'COMBRGB', 'Combine RGB'),
    ('ShaderNodeCombineXYZ', 'COMBXYZ', 'Combine XYZ'),
    ('ShaderNodeMapRange', 'MAP_RANGE', 'Map Range'),
    ('ShaderNodeMath', 'MATH', 'Math'),
    ('ShaderNodeRGBToBW', 'RGBTOBW', 'RGB to BW'),
    ('ShaderNodeSeparateRGB', 'SEPRGB', 'Separate RGB'),
    ('ShaderNodeSeparateXYZ', 'SEPXYZ', 'Separate XYZ'),
    ('ShaderNodeSeparateHSV', 'SEPHSV', 'Separate HSV'),
    ('ShaderNodeVectorMath', 'VECT_MATH', 'Vector Math'),
    ('ShaderNodeWavelength', 'WAVELENGTH', 'Wavelength'),
)
# (rna_type.identifier, type, rna_type.name)
# Keeping mixed case to avoid having to translate entries when adding new nodes in operators.
# Keeping things in alphabetical orde so we don't need to sort later.
shaders_layout_nodes_props = (
    ('NodeFrame', 'FRAME', 'Frame'),
    ('NodeReroute', 'REROUTE', 'Reroute'),
)

# compositing nodes
# (rna_type.identifier, type, rna_type.name)
# Keeping mixed case to avoid having to translate entries when adding new nodes in operators.
# Keeping things in alphabetical orde so we don't need to sort later.
compo_input_nodes_props = (
    ('CompositorNodeBokehImage', 'BOKEHIMAGE', 'Bokeh Image'),
    ('CompositorNodeImage', 'IMAGE', 'Image'),
    ('CompositorNodeMask', 'MASK', 'Mask'),
    ('CompositorNodeMovieClip', 'MOVIECLIP', 'Movie Clip'),
    ('CompositorNodeRLayers', 'R_LAYERS', 'Render Layers'),
    ('CompositorNodeRGB', 'RGB', 'RGB'),
    ('CompositorNodeTexture', 'TEXTURE', 'Texture'),
    ('CompositorNodeTime', 'TIME', 'Time'),
    ('CompositorNodeTrackPos', 'TRACKPOS', 'Track Position'),
    ('CompositorNodeValue', 'VALUE', 'Value'),
)
# (rna_type.identifier, type, rna_type.name)
# Keeping mixed case to avoid having to translate entries when adding new nodes in operators.
# Keeping things in alphabetical orde so we don't need to sort later.
compo_output_nodes_props = (
    ('CompositorNodeComposite', 'COMPOSITE', 'Composite'),
    ('CompositorNodeOutputFile', 'OUTPUT_FILE', 'File Output'),
    ('CompositorNodeLevels', 'LEVELS', 'Levels'),
    ('CompositorNodeSplitViewer', 'SPLITVIEWER', 'Split Viewer'),
    ('CompositorNodeViewer', 'VIEWER', 'Viewer'),
)
# (rna_type.identifier, type, rna_type.name)
# Keeping mixed case to avoid having to translate entries when adding new nodes in operators.
# Keeping things in alphabetical orde so we don't need to sort later.
compo_color_nodes_props = (
    ('CompositorNodeAlphaOver', 'ALPHAOVER', 'Alpha Over'),
    ('CompositorNodeBrightContrast', 'BRIGHTCONTRAST', 'Bright/Contrast'),
    ('CompositorNodeColorBalance', 'COLORBALANCE', 'Color Balance'),
    ('CompositorNodeColorCorrection', 'COLORCORRECTION', 'Color Correction'),
    ('CompositorNodeGamma', 'GAMMA', 'Gamma'),
    ('CompositorNodeHueCorrect', 'HUECORRECT', 'Hue Correct'),
    ('CompositorNodeHueSat', 'HUE_SAT', 'Hue Saturation Value'),
    ('CompositorNodeInvert', 'INVERT', 'Invert'),
    ('CompositorNodeMixRGB', 'MIX_RGB', 'Mix'),
    ('CompositorNodeCurveRGB', 'CURVE_RGB', 'RGB Curves'),
    ('CompositorNodeTonemap', 'TONEMAP', 'Tonemap'),
    ('CompositorNodeZcombine', 'ZCOMBINE', 'Z Combine'),
)
# (rna_type.identifier, type, rna_type.name)
# Keeping mixed case to avoid having to translate entries when adding new nodes in operators.
# Keeping things in alphabetical orde so we don't need to sort later.
compo_converter_nodes_props = (
    ('CompositorNodePremulKey', 'PREMULKEY', 'Alpha Convert'),
    ('CompositorNodeValToRGB', 'VALTORGB', 'ColorRamp'),
    ('CompositorNodeCombHSVA', 'COMBHSVA', 'Combine HSVA'),
    ('CompositorNodeCombRGBA', 'COMBRGBA', 'Combine RGBA'),
    ('CompositorNodeCombYCCA', 'COMBYCCA', 'Combine YCbCrA'),
    ('CompositorNodeCombYUVA', 'COMBYUVA', 'Combine YUVA'),
    ('CompositorNodeIDMask', 'ID_MASK', 'ID Mask'),
    ('CompositorNodeMath', 'MATH', 'Math'),
    ('CompositorNodeRGBToBW', 'RGBTOBW', 'RGB to BW'),
    ('CompositorNodeSepRGBA', 'SEPRGBA', 'Separate RGBA'),
    ('CompositorNodeSepHSVA', 'SEPHSVA', 'Separate HSVA'),
    ('CompositorNodeSepYUVA', 'SEPYUVA', 'Separate YUVA'),
    ('CompositorNodeSepYCCA', 'SEPYCCA', 'Separate YCbCrA'),
    ('CompositorNodeSetAlpha', 'SETALPHA', 'Set Alpha'),
    ('CompositorNodeSwitchView', 'VIEWSWITCH', 'View Switch'),
)
# (rna_type.identifier, type, rna_type.name)
# Keeping mixed case to avoid having to translate entries when adding new nodes in operators.
# Keeping things in alphabetical orde so we don't need to sort later.
compo_filter_nodes_props = (
    ('CompositorNodeBilateralblur', 'BILATERALBLUR', 'Bilateral Blur'),
    ('CompositorNodeBlur', 'BLUR', 'Blur'),
    ('CompositorNodeBokehBlur', 'BOKEHBLUR', 'Bokeh Blur'),
    ('CompositorNodeDefocus', 'DEFOCUS', 'Defocus'),
    ('CompositorNodeDenoise', 'DENOISE', 'Denoise'),
    ('CompositorNodeDespeckle', 'DESPECKLE', 'Despeckle'),
    ('CompositorNodeDilateErode', 'DILATEERODE', 'Dilate/Erode'),
    ('CompositorNodeDBlur', 'DBLUR', 'Directional Blur'),
    ('CompositorNodeFilter', 'FILTER', 'Filter'),
    ('CompositorNodeGlare', 'GLARE', 'Glare'),
    ('CompositorNodeInpaint', 'INPAINT', 'Inpaint'),
    ('CompositorNodePixelate', 'PIXELATE', 'Pixelate'),
    ('CompositorNodeSunBeams', 'SUNBEAMS', 'Sun Beams'),
    ('CompositorNodeVecBlur', 'VECBLUR', 'Vector Blur'),
)
# (rna_type.identifier, type, rna_type.name)
# Keeping mixed case to avoid having to translate entries when adding new nodes in operators.
# Keeping things in alphabetical orde so we don't need to sort later.
compo_vector_nodes_props = (
    ('CompositorNodeMapRange', 'MAP_RANGE', 'Map Range'),
    ('CompositorNodeMapValue', 'MAP_VALUE', 'Map Value'),
    ('CompositorNodeNormal', 'NORMAL', 'Normal'),
    ('CompositorNodeNormalize', 'NORMALIZE', 'Normalize'),
    ('CompositorNodeCurveVec', 'CURVE_VEC', 'Vector Curves'),
)
# (rna_type.identifier, type, rna_type.name)
# Keeping mixed case to avoid having to translate entries when adding new nodes in operators.
# Keeping things in alphabetical orde so we don't need to sort later.
compo_matte_nodes_props = (
    ('CompositorNodeBoxMask', 'BOXMASK', 'Box Mask'),
    ('CompositorNodeChannelMatte', 'CHANNEL_MATTE', 'Channel Key'),
    ('CompositorNodeChromaMatte', 'CHROMA_MATTE', 'Chroma Key'),
    ('CompositorNodeColorMatte', 'COLOR_MATTE', 'Color Key'),
    ('CompositorNodeColorSpill', 'COLOR_SPILL', 'Color Spill'),
    ('CompositorNodeCryptomatte', 'CRYPTOMATTE', 'Cryptomatte'),
    ('CompositorNodeDiffMatte', 'DIFF_MATTE', 'Difference Key'),
    ('CompositorNodeDistanceMatte', 'DISTANCE_MATTE', 'Distance Key'),
    ('CompositorNodeDoubleEdgeMask', 'DOUBLEEDGEMASK', 'Double Edge Mask'),
    ('CompositorNodeEllipseMask', 'ELLIPSEMASK', 'Ellipse Mask'),
    ('CompositorNodeKeying', 'KEYING', 'Keying'),
    ('CompositorNodeKeyingScreen', 'KEYINGSCREEN', 'Keying Screen'),
    ('CompositorNodeLumaMatte', 'LUMA_MATTE', 'Luminance Key'),
)
# (rna_type.identifier, type, rna_type.name)
# Keeping mixed case to avoid having to translate entries when adding new nodes in operators.
# Keeping things in alphabetical orde so we don't need to sort later.
compo_distort_nodes_props = (
    ('CompositorNodeCornerPin', 'CORNERPIN', 'Corner Pin'),
    ('CompositorNodeCrop', 'CROP', 'Crop'),
    ('CompositorNodeDisplace', 'DISPLACE', 'Displace'),
    ('CompositorNodeFlip', 'FLIP', 'Flip'),
    ('CompositorNodeLensdist', 'LENSDIST', 'Lens Distortion'),
    ('CompositorNodeMapUV', 'MAP_UV', 'Map UV'),
    ('CompositorNodeMovieDistortion', 'MOVIEDISTORTION', 'Movie Distortion'),
    ('CompositorNodePlaneTrackDeform', 'PLANETRACKDEFORM', 'Plane Track Deform'),
    ('CompositorNodeRotate', 'ROTATE', 'Rotate'),
    ('CompositorNodeScale', 'SCALE', 'Scale'),
    ('CompositorNodeStabilize', 'STABILIZE2D', 'Stabilize 2D'),
    ('CompositorNodeTransform', 'TRANSFORM', 'Transform'),
    ('CompositorNodeTranslate', 'TRANSLATE', 'Translate'),
)
# (rna_type.identifier, type, rna_type.name)
# Keeping mixed case to avoid having to translate entries when adding new nodes in operators.
# Keeping things in alphabetical orde so we don't need to sort later.
compo_layout_nodes_props = (
    ('NodeFrame', 'FRAME', 'Frame'),
    ('NodeReroute', 'REROUTE', 'Reroute'),
    ('CompositorNodeSwitch', 'SWITCH', 'Switch'),
)
# Blender Render material nodes
# (rna_type.identifier, type, rna_type.name)
# Keeping mixed case to avoid having to translate entries when adding new nodes in operators.
blender_mat_input_nodes_props = (
    ('ShaderNodeMaterial', 'MATERIAL', 'Material'),
    ('ShaderNodeCameraData', 'CAMERA', 'Camera Data'),
    ('ShaderNodeLightData', 'LIGHT', 'Light Data'),
    ('ShaderNodeValue', 'VALUE', 'Value'),
    ('ShaderNodeRGB', 'RGB', 'RGB'),
    ('ShaderNodeTexture', 'TEXTURE', 'Texture'),
    ('ShaderNodeGeometry', 'GEOMETRY', 'Geometry'),
    ('ShaderNodeExtendedMaterial', 'MATERIAL_EXT', 'Extended Material'),
)

# (rna_type.identifier, type, rna_type.name)
# Keeping mixed case to avoid having to translate entries when adding new nodes in operators.
blender_mat_output_nodes_props = (
    ('ShaderNodeOutput', 'OUTPUT', 'Output'),
)

# (rna_type.identifier, type, rna_type.name)
# Keeping mixed case to avoid having to translate entries when adding new nodes in operators.
blender_mat_color_nodes_props = (
    ('ShaderNodeMixRGB', 'MIX_RGB', 'MixRGB'),
    ('ShaderNodeRGBCurve', 'CURVE_RGB', 'RGB Curves'),
    ('ShaderNodeInvert', 'INVERT', 'Invert'),
    ('ShaderNodeHueSaturation', 'HUE_SAT', 'Hue/Saturation'),
)

# (rna_type.identifier, type, rna_type.name)
# Keeping mixed case to avoid having to translate entries when adding new nodes in operators.
blender_mat_vector_nodes_props = (
    ('ShaderNodeNormal', 'NORMAL', 'Normal'),
    ('ShaderNodeMapping', 'MAPPING', 'Mapping'),
    ('ShaderNodeVectorCurve', 'CURVE_VEC', 'Vector Curves'),
)

# (rna_type.identifier, type, rna_type.name)
# Keeping mixed case to avoid having to translate entries when adding new nodes in operators.
blender_mat_converter_nodes_props = (
    ('ShaderNodeValToRGB', 'VALTORGB', 'ColorRamp'),
    ('ShaderNodeRGBToBW', 'RGBTOBW', 'RGB to BW'),
    ('ShaderNodeMath', 'MATH', 'Math'),
    ('ShaderNodeVectorMath', 'VECT_MATH', 'Vector Math'),
    ('ShaderNodeSqueeze', 'SQUEEZE', 'Squeeze Value'),
    ('ShaderNodeSeparateRGB', 'SEPRGB', 'Separate RGB'),
    ('ShaderNodeCombineRGB', 'COMBRGB', 'Combine RGB'),
    ('ShaderNodeSeparateHSV', 'SEPHSV', 'Separate HSV'),
    ('ShaderNodeCombineHSV', 'COMBHSV', 'Combine HSV'),
)

# (rna_type.identifier, type, rna_type.name)
# Keeping mixed case to avoid having to translate entries when adding new nodes in operators.
blender_mat_layout_nodes_props = (
    ('NodeReroute', 'REROUTE', 'Reroute'),
)

# Texture Nodes
# (rna_type.identifier, type, rna_type.name)
# Keeping mixed case to avoid having to translate entries when adding new nodes in operators.
texture_input_nodes_props = (
    ('TextureNodeCurveTime', 'CURVE_TIME', 'Curve Time'),
    ('TextureNodeCoordinates', 'COORD', 'Coordinates'),
    ('TextureNodeTexture', 'TEXTURE', 'Texture'),
    ('TextureNodeImage', 'IMAGE', 'Image'),
)

# (rna_type.identifier, type, rna_type.name)
# Keeping mixed case to avoid having to translate entries when adding new nodes in operators.
texture_output_nodes_props = (
    ('TextureNodeOutput', 'OUTPUT', 'Output'),
    ('TextureNodeViewer', 'VIEWER', 'Viewer'),
)

# (rna_type.identifier, type, rna_type.name)
# Keeping mixed case to avoid having to translate entries when adding new nodes in operators.
texture_color_nodes_props = (
    ('TextureNodeMixRGB', 'MIX_RGB', 'Mix RGB'),
    ('TextureNodeCurveRGB', 'CURVE_RGB', 'RGB Curves'),
    ('TextureNodeInvert', 'INVERT', 'Invert'),
    ('TextureNodeHueSaturation', 'HUE_SAT', 'Hue/Saturation'),
    ('TextureNodeCompose', 'COMPOSE', 'Combine RGBA'),
    ('TextureNodeDecompose', 'DECOMPOSE', 'Separate RGBA'),
)

# (rna_type.identifier, type, rna_type.name)
# Keeping mixed case to avoid having to translate entries when adding new nodes in operators.
texture_pattern_nodes_props = (
    ('TextureNodeChecker', 'CHECKER', 'Checker'),
    ('TextureNodeBricks', 'BRICKS', 'Bricks'),
)

# (rna_type.identifier, type, rna_type.name)
# Keeping mixed case to avoid having to translate entries when adding new nodes in operators.
texture_textures_nodes_props = (
    ('TextureNodeTexNoise', 'TEX_NOISE', 'Noise'),
    ('TextureNodeTexDistNoise', 'TEX_DISTNOISE', 'Distorted Noise'),
    ('TextureNodeTexClouds', 'TEX_CLOUDS', 'Clouds'),
    ('TextureNodeTexBlend', 'TEX_BLEND', 'Blend'),
    ('TextureNodeTexVoronoi', 'TEX_VORONOI', 'Voronoi'),
    ('TextureNodeTexMagic', 'TEX_MAGIC', 'Magic'),
    ('TextureNodeTexMarble', 'TEX_MARBLE', 'Marble'),
    ('TextureNodeTexWood', 'TEX_WOOD', 'Wood'),
    ('TextureNodeTexMusgrave', 'TEX_MUSGRAVE', 'Musgrave'),
    ('TextureNodeTexStucci', 'TEX_STUCCI', 'Stucci'),
)

# (rna_type.identifier, type, rna_type.name)
# Keeping mixed case to avoid having to translate entries when adding new nodes in operators.
texture_converter_nodes_props = (
    ('TextureNodeMath', 'MATH', 'Math'),
    ('TextureNodeValToRGB', 'VALTORGB', 'ColorRamp'),
    ('TextureNodeRGBToBW', 'RGBTOBW', 'RGB to BW'),
    ('TextureNodeValToNor', 'VALTONOR', 'Value to Normal'),
    ('TextureNodeDistance', 'DISTANCE', 'Distance'),
)

# (rna_type.identifier, type, rna_type.name)
# Keeping mixed case to avoid having to translate entries when adding new nodes in operators.
texture_distort_nodes_props = (
    ('TextureNodeScale', 'SCALE', 'Scale'),
    ('TextureNodeTranslate', 'TRANSLATE', 'Translate'),
    ('TextureNodeRotate', 'ROTATE', 'Rotate'),
    ('TextureNodeAt', 'AT', 'At'),
)

# (rna_type.identifier, type, rna_type.name)
# Keeping mixed case to avoid having to translate entries when adding new nodes in operators.
texture_layout_nodes_props = (
    ('NodeReroute', 'REROUTE', 'Reroute'),
)
# list of blend types of "Mix" nodes in a form that can be used as 'items' for EnumProperty.
# used list, not tuple for easy merging with other lists.
blend_types = [
    ('MIX', 'Mix', 'Mix Mode'),
    ('ADD', 'Add', 'Add Mode'),
    ('MULTIPLY', 'Multiply', 'Multiply Mode'),
    ('SUBTRACT', 'Subtract', 'Subtract Mode'),
    ('SCREEN', 'Screen', 'Screen Mode'),
    ('DIVIDE', 'Divide', 'Divide Mode'),
    ('DIFFERENCE', 'Difference', 'Difference Mode'),
    ('DARKEN', 'Darken', 'Darken Mode'),
    ('LIGHTEN', 'Lighten', 'Lighten Mode'),
    ('OVERLAY', 'Overlay', 'Overlay Mode'),
    ('DODGE', 'Dodge', 'Dodge Mode'),
    ('BURN', 'Burn', 'Burn Mode'),
    ('HUE', 'Hue', 'Hue Mode'),
    ('SATURATION', 'Saturation', 'Saturation Mode'),
    ('VALUE', 'Value', 'Value Mode'),
    ('COLOR', 'Color', 'Color Mode'),
    ('SOFT_LIGHT', 'Soft Light', 'Soft Light Mode'),
    ('LINEAR_LIGHT', 'Linear Light', 'Linear Light Mode'),
# list of operations of "Math" nodes in a form that can be used as 'items' for EnumProperty.
# used list, not tuple for easy merging with other lists.
operations = [
    ('ADD', 'Add', 'Add Mode'),
    ('SUBTRACT', 'Subtract', 'Subtract Mode'),
    ('MULTIPLY', 'Multiply', 'Multiply Mode'),
    ('DIVIDE', 'Divide', 'Divide Mode'),
    ('MULTIPLY_ADD', 'Multiply Add', 'Multiply Add Mode'),
    ('SINE', 'Sine', 'Sine Mode'),
    ('COSINE', 'Cosine', 'Cosine Mode'),
    ('TANGENT', 'Tangent', 'Tangent Mode'),
    ('ARCSINE', 'Arcsine', 'Arcsine Mode'),
    ('ARCCOSINE', 'Arccosine', 'Arccosine Mode'),
    ('ARCTANGENT', 'Arctangent', 'Arctangent Mode'),
    ('ARCTAN2', 'Arctan2', 'Arctan2 Mode'),
    ('SINH', 'Hyperbolic Sine', 'Hyperbolic Sine Mode'),
    ('COSH', 'Hyperbolic Cosine', 'Hyperbolic Cosine Mode'),
    ('TANH', 'Hyperbolic Tangent', 'Hyperbolic Tangent Mode'),
    ('POWER', 'Power', 'Power Mode'),
    ('LOGARITHM', 'Logatithm', 'Logarithm Mode'),
    ('SQRT', 'Square Root', 'Square Root Mode'),
    ('INVERSE_SQRT', 'Inverse Square Root', 'Inverse Square Root Mode'),
    ('EXPONENT', 'Exponent', 'Exponent Mode'),
    ('MINIMUM', 'Minimum', 'Minimum Mode'),
    ('MAXIMUM', 'Maximum', 'Maximum Mode'),
    ('LESS_THAN', 'Less Than', 'Less Than Mode'),
    ('GREATER_THAN', 'Greater Than', 'Greater Than Mode'),
    ('SIGN', 'Sign', 'Sign Mode'),
    ('COMPARE', 'Compare', 'Compare Mode'),
    ('SMOOTH_MIN', 'Smooth Minimum', 'Smooth Minimum Mode'),
    ('SMOOTH_MAX', 'Smooth Maximum', 'Smooth Maximum Mode'),
    ('FRACT', 'Fraction', 'Fraction Mode'),
    ('MODULO', 'Modulo', 'Modulo Mode'),
    ('SNAP', 'Snap', 'Snap Mode'),
    ('WRAP', 'Wrap', 'Wrap Mode'),
    ('PINGPONG', 'Pingpong', 'Pingpong Mode'),
    ('ABSOLUTE', 'Absolute', 'Absolute Mode'),
    ('ROUND', 'Round', 'Round Mode'),
    ('FLOOR', 'Floor', 'Floor Mode'),
    ('CEIL', 'Ceil', 'Ceil Mode'),
    ('TRUNCATE', 'Truncate', 'Truncate Mode'),
    ('RADIANS', 'To Radians', 'To Radians Mode'),
    ('DEGREES', 'To Degrees', 'To Degrees Mode'),
]

# in NWBatchChangeNodes additional types/operations. Can be used as 'items' for EnumProperty.
# used list, not tuple for easy merging with other lists.
navs = [
    ('CURRENT', 'Current', 'Leave at current state'),
    ('NEXT', 'Next', 'Next blend type/operation'),
    ('PREV', 'Prev', 'Previous blend type/operation'),
]

draw_color_sets = {
    "red_white": (
        (1.0, 1.0, 1.0, 0.7),
        (1.0, 0.0, 0.0, 0.7),
        (0.8, 0.2, 0.2, 1.0)
    ),
    "green": (
        (0.0, 0.0, 0.0, 1.0),
        (0.38, 0.77, 0.38, 1.0),
        (0.38, 0.77, 0.38, 1.0)
    ),
    "yellow": (
        (0.0, 0.0, 0.0, 1.0),
        (0.77, 0.77, 0.16, 1.0),
        (0.77, 0.77, 0.16, 1.0)
    ),
    "purple": (
        (0.0, 0.0, 0.0, 1.0),
        (0.38, 0.38, 0.77, 1.0),
        (0.38, 0.38, 0.77, 1.0)
    ),
    "grey": (
        (0.0, 0.0, 0.0, 1.0),
        (0.63, 0.63, 0.63, 1.0),
        (0.63, 0.63, 0.63, 1.0)
    ),
    "black": (
        (1.0, 1.0, 1.0, 0.7),
        (0.0, 0.0, 0.0, 0.7),
        (0.2, 0.2, 0.2, 1.0)
viewer_socket_name = "tmp_viewer"
def is_visible_socket(socket):
    return not socket.hide and socket.enabled and socket.type != 'CUSTOM'
def nice_hotkey_name(punc):
    # convert the ugly string name into the actual character
    pairs = (
        ('LEFTMOUSE', "LMB"),
        ('MIDDLEMOUSE', "MMB"),
        ('RIGHTMOUSE', "RMB"),
        ('WHEELUPMOUSE', "Wheel Up"),
        ('WHEELDOWNMOUSE', "Wheel Down"),
        ('WHEELINMOUSE', "Wheel In"),
        ('WHEELOUTMOUSE', "Wheel Out"),
        ('ZERO', "0"),
        ('ONE', "1"),
        ('TWO', "2"),
        ('THREE', "3"),
        ('FOUR', "4"),
        ('FIVE', "5"),
        ('SIX', "6"),
        ('SEVEN', "7"),
        ('EIGHT', "8"),
        ('NINE', "9"),
        ('OSKEY', "Super"),
        ('RET', "Enter"),
        ('LINE_FEED', "Enter"),
        ('SEMI_COLON', ";"),
        ('PERIOD', "."),
        ('COMMA', ","),
        ('QUOTE', '"'),
        ('MINUS', "-"),
        ('SLASH', "/"),
        ('BACK_SLASH', "\\"),
        ('EQUAL', "="),
        ('NUMPAD_1', "Numpad 1"),
        ('NUMPAD_2', "Numpad 2"),
        ('NUMPAD_3', "Numpad 3"),
        ('NUMPAD_4', "Numpad 4"),
        ('NUMPAD_5', "Numpad 5"),
        ('NUMPAD_6', "Numpad 6"),
        ('NUMPAD_7', "Numpad 7"),
        ('NUMPAD_8', "Numpad 8"),
        ('NUMPAD_9', "Numpad 9"),
        ('NUMPAD_0', "Numpad 0"),
        ('NUMPAD_PERIOD', "Numpad ."),
        ('NUMPAD_SLASH', "Numpad /"),
        ('NUMPAD_ASTERIX', "Numpad *"),
        ('NUMPAD_MINUS', "Numpad -"),
        ('NUMPAD_ENTER', "Numpad Enter"),
        ('NUMPAD_PLUS', "Numpad +"),
    nice_punc = False
    for (ugly, nice) in pairs:
        if punc == ugly:
            nice_punc = nice
            break
    if not nice_punc:
        nice_punc = punc.replace("_", " ").title()
    return nice_punc


def force_update(context):
    context.space_data.node_tree.update_tag()
    prefs = bpy.context.preferences.system
    return prefs.dpi * prefs.pixel_size / 72
def node_mid_pt(node, axis):
    if axis == 'x':
        d = node.location.x + (node.dimensions.x / 2)
    elif axis == 'y':
        d = node.location.y - (node.dimensions.y / 2)
    else:
        d = 0
    return d


def autolink(node1, node2, links):
    link_made = False

    for outp in node1.outputs:
        for inp in node2.inputs:
            if not inp.is_linked and inp.name == outp.name:
                link_made = True
                links.new(outp, inp)
                return True

    for outp in node1.outputs:
        for inp in node2.inputs:
            if not inp.is_linked and inp.type == outp.type:
                link_made = True
                links.new(outp, inp)
                return True

    # force some connection even if the type doesn't match
    for outp in node1.outputs:
        for inp in node2.inputs:
            if not inp.is_linked:
                link_made = True
                links.new(outp, inp)
                return True

    # even if no sockets are open, force one of matching type
    for outp in node1.outputs:
        for inp in node2.inputs:
            if inp.type == outp.type:
                link_made = True
                links.new(outp, inp)
                return True

    # do something!
    for outp in node1.outputs:
        for inp in node2.inputs:
            link_made = True
            links.new(outp, inp)
            return True

    print("Could not make a link from " + node1.name + " to " + node2.name)
    return link_made


def node_at_pos(nodes, context, event):
    nodes_near_mouse = []
    nodes_under_mouse = []
    target_node = None

    store_mouse_cursor(context, event)
    x, y = context.space_data.cursor_location
    x = x
    y = y

    # Make a list of each corner (and middle of border) for each node.
    # Will be sorted to find nearest point and thus nearest node
    node_points_with_dist = []
    for node in nodes:
        skipnode = False
        if node.type != 'FRAME':  # no point trying to link to a frame node
            locx = node.location.x
            locy = node.location.y
            dimx = node.dimensions.x/dpifac()
            dimy = node.dimensions.y/dpifac()
            if node.parent:
                locx += node.parent.location.x
                locy += node.parent.location.y
                if node.parent.parent:
                    locx += node.parent.parent.location.x
                    locy += node.parent.parent.location.y
                    if node.parent.parent.parent:
                        locx += node.parent.parent.parent.location.x
                        locy += node.parent.parent.parent.location.y
                        if node.parent.parent.parent.parent:
                            # Support three levels or parenting
                            # There's got to be a better way to do this...
                            skipnode = True
            if not skipnode:
                node_points_with_dist.append([node, hypot(x - locx, y - locy)])  # Top Left
                node_points_with_dist.append([node, hypot(x - (locx + dimx), y - locy)])  # Top Right
                node_points_with_dist.append([node, hypot(x - locx, y - (locy - dimy))])  # Bottom Left
                node_points_with_dist.append([node, hypot(x - (locx + dimx), y - (locy - dimy))])  # Bottom Right

                node_points_with_dist.append([node, hypot(x - (locx + (dimx / 2)), y - locy)])  # Mid Top
                node_points_with_dist.append([node, hypot(x - (locx + (dimx / 2)), y - (locy - dimy))])  # Mid Bottom
                node_points_with_dist.append([node, hypot(x - locx, y - (locy - (dimy / 2)))])  # Mid Left
                node_points_with_dist.append([node, hypot(x - (locx + dimx), y - (locy - (dimy / 2)))])  # Mid Right
    nearest_node = sorted(node_points_with_dist, key=lambda k: k[1])[0][0]
        if node.type != 'FRAME' and skipnode == False:
            locx = node.location.x
            locy = node.location.y
            dimx = node.dimensions.x/dpifac()
            dimy = node.dimensions.y/dpifac()
            if node.parent:
                locx += node.parent.location.x
                locy += node.parent.location.y
            if (locx <= x <= locx + dimx) and \
               (locy - dimy <= y <= locy):
                nodes_under_mouse.append(node)

    if len(nodes_under_mouse) == 1:
        if nodes_under_mouse[0] != nearest_node:
            target_node = nodes_under_mouse[0]  # use the node under the mouse if there is one and only one
        else:
            target_node = nearest_node  # else use the nearest node
        target_node = nearest_node
    return target_node


def store_mouse_cursor(context, event):
    space = context.space_data
    v2d = context.region.view2d
    tree = space.edit_tree

    # convert mouse position to the View2D for later node placement
    if context.region.type == 'WINDOW':
        space.cursor_location_from_region(event.mouse_region_x, event.mouse_region_y)
    else:
        space.cursor_location = tree.view_center

def draw_line(x1, y1, x2, y2, size, colour=(1.0, 1.0, 1.0, 0.7)):
    shader = gpu.shader.from_builtin('2D_SMOOTH_COLOR')
    vertices = ((x1, y1), (x2, y2))
    vertex_colors = ((colour[0]+(1.0-colour[0])/4,
                      colour[1]+(1.0-colour[1])/4,
                      colour[2]+(1.0-colour[2])/4,
                      colour[3]+(1.0-colour[3])/4),
                      colour)
    batch = batch_for_shader(shader, 'LINE_STRIP', {"pos": vertices, "color": vertex_colors})
    bgl.glLineWidth(size * dpifac())

    shader.bind()
    batch.draw(shader)


def draw_circle_2d_filled(shader, mx, my, radius, colour=(1.0, 1.0, 1.0, 0.7)):
    radius = radius * dpifac()
    sides = 12
    vertices = [(radius * cos(i * 2 * pi / sides) + mx,
                 radius * sin(i * 2 * pi / sides) + my)
                 for i in range(sides + 1)]
    batch = batch_for_shader(shader, 'TRI_FAN', {"pos": vertices})
    shader.bind()
    shader.uniform_float("color", colour)
    batch.draw(shader)
def draw_rounded_node_border(shader, node, radius=8, colour=(1.0, 1.0, 1.0, 0.7)):
    area_width = bpy.context.area.width - (16*dpifac()) - 1
    bottom_bar = (16*dpifac()) + 1
    radius = radius*dpifac()

    nlocx = (node.location.x+1)*dpifac()
    nlocy = (node.location.y+1)*dpifac()
    ndimx = node.dimensions.x
    ndimy = node.dimensions.y
Greg Zaal's avatar
Greg Zaal committed
    # This is a stupid way to do this... TODO use while loop
    if node.parent:
        nlocx += node.parent.location.x
        nlocy += node.parent.location.y
        if node.parent.parent:
            nlocx += node.parent.parent.location.x
            nlocy += node.parent.parent.location.y
            if node.parent.parent.parent:
                nlocx += node.parent.parent.parent.location.x
                nlocy += node.parent.parent.parent.location.y

Greg Zaal's avatar
Greg Zaal committed
    if node.hide:
        nlocx += -1
        nlocy += 5
    if node.type == 'REROUTE':
        #nlocx += 1
        nlocy -= 1
        ndimx = 0
        ndimy = 0
        radius += 6

    # Top left corner
    mx, my = bpy.context.region.view2d.view_to_region(nlocx, nlocy, clip=False)
    vertices = [(mx,my)]
    for i in range(sides+1):
        if (4<=i<=8):
            if my > bottom_bar and mx < area_width:
                cosine = radius * cos(i * 2 * pi / sides) + mx
                sine = radius * sin(i * 2 * pi / sides) + my
                vertices.append((cosine,sine))
    batch = batch_for_shader(shader, 'TRI_FAN', {"pos": vertices})
    shader.bind()
    shader.uniform_float("color", colour)
    batch.draw(shader)
Greg Zaal's avatar
Greg Zaal committed
    # Top right corner
    mx, my = bpy.context.region.view2d.view_to_region(nlocx + ndimx, nlocy, clip=False)
    vertices = [(mx,my)]
    for i in range(sides+1):
        if (0<=i<=4):
            if my > bottom_bar and mx < area_width:
                cosine = radius * cos(i * 2 * pi / sides) + mx
                sine = radius * sin(i * 2 * pi / sides) + my
                vertices.append((cosine,sine))
    batch = batch_for_shader(shader, 'TRI_FAN', {"pos": vertices})
    shader.bind()
    shader.uniform_float("color", colour)
    batch.draw(shader)
Greg Zaal's avatar
Greg Zaal committed

    # Bottom left corner
    mx, my = bpy.context.region.view2d.view_to_region(nlocx, nlocy - ndimy, clip=False)
    vertices = [(mx,my)]
    for i in range(sides+1):
        if (8<=i<=12):
            if my > bottom_bar and mx < area_width:
                cosine = radius * cos(i * 2 * pi / sides) + mx
                sine = radius * sin(i * 2 * pi / sides) + my
                vertices.append((cosine,sine))
    batch = batch_for_shader(shader, 'TRI_FAN', {"pos": vertices})
    shader.bind()
    shader.uniform_float("color", colour)
    batch.draw(shader)
Greg Zaal's avatar
Greg Zaal committed
    # Bottom right corner
    mx, my = bpy.context.region.view2d.view_to_region(nlocx + ndimx, nlocy - ndimy, clip=False)
    vertices = [(mx,my)]
    for i in range(sides+1):
        if (12<=i<=16):
            if my > bottom_bar and mx < area_width:
                cosine = radius * cos(i * 2 * pi / sides) + mx
                sine = radius * sin(i * 2 * pi / sides) + my
                vertices.append((cosine,sine))
    batch = batch_for_shader(shader, 'TRI_FAN', {"pos": vertices})
    shader.bind()
    shader.uniform_float("color", colour)
    batch.draw(shader)
    # prepare drawing all edges in one batch
    vertices = []
    indices = []
    id_last = 0
Greg Zaal's avatar
Greg Zaal committed
    # Left edge
    m1x, m1y = bpy.context.region.view2d.view_to_region(nlocx, nlocy, clip=False)
    m2x, m2y = bpy.context.region.view2d.view_to_region(nlocx, nlocy - ndimy, clip=False)
    if m1x < area_width and m2x < area_width:
        vertices.extend([(m2x-radius,m2y), (m2x,m2y),
                         (m1x,m1y), (m1x-radius,m1y)])
        indices.extend([(id_last, id_last+1, id_last+3),
                        (id_last+3, id_last+1, id_last+2)])
        id_last += 4
Greg Zaal's avatar
Greg Zaal committed
    # Top edge
    m1x, m1y = bpy.context.region.view2d.view_to_region(nlocx, nlocy, clip=False)
    m2x, m2y = bpy.context.region.view2d.view_to_region(nlocx + ndimx, nlocy, clip=False)
    m1x = min(m1x, area_width)
    m2x = min(m2x, area_width)
    if m1y > bottom_bar and m2y > bottom_bar:
        vertices.extend([(m1x,m1y), (m2x,m1y),
                         (m2x,m1y+radius), (m1x,m1y+radius)])
        indices.extend([(id_last, id_last+1, id_last+3),
                        (id_last+3, id_last+1, id_last+2)])
        id_last += 4
Greg Zaal's avatar
Greg Zaal committed
    # Right edge
    m1x, m1y = bpy.context.region.view2d.view_to_region(nlocx + ndimx, nlocy, clip=False)
    m2x, m2y = bpy.context.region.view2d.view_to_region(nlocx + ndimx, nlocy - ndimy, clip=False)
    m1y = max(m1y, bottom_bar)
    m2y = max(m2y, bottom_bar)
    if m1x < area_width and m2x < area_width:
        vertices.extend([(m1x,m2y), (m1x+radius,m2y),
                         (m1x+radius,m1y), (m1x,m1y)])
        indices.extend([(id_last, id_last+1, id_last+3),
                        (id_last+3, id_last+1, id_last+2)])
        id_last += 4
Greg Zaal's avatar
Greg Zaal committed
    # Bottom edge
    m1x, m1y = bpy.context.region.view2d.view_to_region(nlocx, nlocy-ndimy, clip=False)
    m2x, m2y = bpy.context.region.view2d.view_to_region(nlocx + ndimx, nlocy-ndimy, clip=False)
    m1x = min(m1x, area_width)
    m2x = min(m2x, area_width)
    if m1y > bottom_bar and m2y > bottom_bar:
        vertices.extend([(m1x,m2y), (m2x,m2y),
                         (m2x,m1y-radius), (m1x,m1y-radius)])
        indices.extend([(id_last, id_last+1, id_last+3),
                        (id_last+3, id_last+1, id_last+2)])

    # now draw all edges in one batch
    if len(vertices) != 0:
        batch = batch_for_shader(shader, 'TRIS', {"pos": vertices}, indices=indices)
        shader.bind()
        shader.uniform_float("color", colour)
        batch.draw(shader)
def draw_callback_nodeoutline(self, context, mode):

        bgl.glLineWidth(1)
        bgl.glEnable(bgl.GL_BLEND)
        bgl.glEnable(bgl.GL_LINE_SMOOTH)
        bgl.glHint(bgl.GL_LINE_SMOOTH_HINT, bgl.GL_NICEST)

        nodes, links = get_nodes_links(context)

        shader = gpu.shader.from_builtin('2D_UNIFORM_COLOR')
        if mode == "LINK":
            col_outer = (1.0, 0.2, 0.2, 0.4)
            col_inner = (0.0, 0.0, 0.0, 0.5)
            col_circle_inner = (0.3, 0.05, 0.05, 1.0)
        elif mode == "LINKMENU":
            col_outer = (0.4, 0.6, 1.0, 0.4)
            col_inner = (0.0, 0.0, 0.0, 0.5)
            col_circle_inner = (0.08, 0.15, .3, 1.0)
        elif mode == "MIX":
            col_outer = (0.2, 1.0, 0.2, 0.4)
            col_inner = (0.0, 0.0, 0.0, 0.5)
            col_circle_inner = (0.05, 0.3, 0.05, 1.0)

        m1x = self.mouse_path[0][0]
        m1y = self.mouse_path[0][1]
        m2x = self.mouse_path[-1][0]
        m2y = self.mouse_path[-1][1]

        n1 = nodes[context.scene.NWLazySource]
        n2 = nodes[context.scene.NWLazyTarget]
        if n1 == n2:
            col_outer = (0.4, 0.4, 0.4, 0.4)
            col_inner = (0.0, 0.0, 0.0, 0.5)
            col_circle_inner = (0.2, 0.2, 0.2, 1.0)
        draw_rounded_node_border(shader, n1, radius=6, colour=col_outer)  # outline
        draw_rounded_node_border(shader, n1, radius=5, colour=col_inner)  # inner
        draw_rounded_node_border(shader, n2, radius=6, colour=col_outer)  # outline
        draw_rounded_node_border(shader, n2, radius=5, colour=col_inner)  # inner
        draw_line(m1x, m1y, m2x, m2y, 5, col_outer)  # line outline
        draw_line(m1x, m1y, m2x, m2y, 2, col_inner)  # line inner

        # circle outline
        draw_circle_2d_filled(shader, m1x, m1y, 7, col_outer)
        draw_circle_2d_filled(shader, m2x, m2y, 7, col_outer)
        draw_circle_2d_filled(shader, m1x, m1y, 5, col_circle_inner)
        draw_circle_2d_filled(shader, m2x, m2y, 5, col_circle_inner)
        bgl.glDisable(bgl.GL_LINE_SMOOTH)
def get_active_tree(context):
    tree = context.space_data.node_tree
    # Get nodes from currently edited tree.
    # If user is editing a group, space_data.node_tree is still the base level (outside group).
    # context.active_node is in the group though, so if space_data.node_tree.nodes.active is not
    # the same as context.active_node, the user is in a group.
    # Check recursively until we find the real active node_tree:
    if tree.nodes.active:
        while tree.nodes.active != context.active_node:
            tree = tree.nodes.active.node_tree
            path.append(tree)
    return tree, path
def get_nodes_links(context):
    tree, path = get_active_tree(context)
    return tree.nodes, tree.links
def is_viewer_socket(socket):
    # checks if a internal socket is a valid viewer socket
    return socket.name == viewer_socket_name and socket.NWViewerSocket

def get_internal_socket(socket):
    #get the internal socket from a socket inside or outside the group
    node = socket.node
    if node.type == 'GROUP_OUTPUT':
        source_iterator = node.inputs
        iterator = node.id_data.outputs
    elif node.type == 'GROUP_INPUT':
        source_iterator = node.outputs
        iterator = node.id_data.inputs
    elif hasattr(node, "node_tree"):
        if socket.is_output:
            source_iterator = node.outputs
            iterator = node.node_tree.outputs
        else:
            source_iterator = node.inputs
            iterator = node.node_tree.inputs
    else:
        return None

    for i, s in enumerate(source_iterator):
        if s == socket:
            break
    return iterator[i]

def is_viewer_link(link, output_node):
    if "Emission Viewer" in link.to_node.name or link.to_node == output_node and link.to_socket == output_node.inputs[0]:
        return True
    if link.to_node.type == 'GROUP_OUTPUT':
        socket = get_internal_socket(link.to_socket)
        if is_viewer_socket(socket):
            return True
    return False

def get_group_output_node(tree):
    for node in tree.nodes:
        if node.type == 'GROUP_OUTPUT' and node.is_active_output == True:
            return node

def get_output_location(tree):
    # get right-most location
    sorted_by_xloc = (sorted(tree.nodes, key=lambda x: x.location.x))
    max_xloc_node = sorted_by_xloc[-1]
    if max_xloc_node.name == 'Emission Viewer':
        max_xloc_node = sorted_by_xloc[-2]

    # get average y location
    sum_yloc = 0
    for node in tree.nodes:
        sum_yloc += node.location.y

    loc_x = max_xloc_node.location.x + max_xloc_node.dimensions.x + 80
    loc_y = sum_yloc / len(tree.nodes)
    return loc_x, loc_y

# Principled prefs
class NWPrincipledPreferences(bpy.types.PropertyGroup):
Campbell Barton's avatar
Campbell Barton committed
    base_color: StringProperty(
        name='Base Color',
        default='diffuse diff albedo base col color',
        description='Naming Components for Base Color maps')
Campbell Barton's avatar
Campbell Barton committed
    sss_color: StringProperty(
        name='Subsurface Color',
        default='sss subsurface',
        description='Naming Components for Subsurface Color maps')
Campbell Barton's avatar
Campbell Barton committed
    metallic: StringProperty(
        name='Metallic',
        default='metallic metalness metal mtl',
        description='Naming Components for metallness maps')
Campbell Barton's avatar
Campbell Barton committed
    specular: StringProperty(
        name='Specular',
        default='specularity specular spec spc',
        description='Naming Components for Specular maps')
Campbell Barton's avatar
Campbell Barton committed
    normal: StringProperty(
        name='Normal',
        default='normal nor nrm nrml norm',
        description='Naming Components for Normal maps')
Campbell Barton's avatar
Campbell Barton committed
    bump: StringProperty(
        name='Bump',
        default='bump bmp',
        description='Naming Components for bump maps')
Campbell Barton's avatar
Campbell Barton committed
    rough: StringProperty(
        name='Roughness',
        default='roughness rough rgh',
        description='Naming Components for roughness maps')
Campbell Barton's avatar
Campbell Barton committed
    gloss: StringProperty(
        name='Gloss',
        default='gloss glossy glossiness',
        description='Naming Components for glossy maps')
Campbell Barton's avatar
Campbell Barton committed
    displacement: StringProperty(
        name='Displacement',
        default='displacement displace disp dsp height heightmap',
        description='Naming Components for displacement maps')
# Addon prefs
class NWNodeWrangler(bpy.types.AddonPreferences):
    bl_idname = __name__

Campbell Barton's avatar
Campbell Barton committed
    merge_hide: EnumProperty(
        name="Hide Mix nodes",
        items=(
            ("ALWAYS", "Always", "Always collapse the new merge nodes"),
            ("NON_SHADER", "Non-Shader", "Collapse in all cases except for shaders"),
            ("NEVER", "Never", "Never collapse the new merge nodes")
        ),
        default='NON_SHADER',
        description="When merging nodes with the Ctrl+Numpad0 hotkey (and similar) specify whether to collapse them or show the full node with options expanded")
Campbell Barton's avatar
Campbell Barton committed
    merge_position: EnumProperty(
        name="Mix Node Position",
        items=(
            ("CENTER", "Center", "Place the Mix node between the two nodes"),
            ("BOTTOM", "Bottom", "Place the Mix node at the same height as the lowest node")
        ),
        default='CENTER',
        description="When merging nodes with the Ctrl+Numpad0 hotkey (and similar) specify the position of the new nodes")
Campbell Barton's avatar
Campbell Barton committed
    show_hotkey_list: BoolProperty(
        name="Show Hotkey List",
        default=False,
        description="Expand this box into a list of all the hotkeys for functions in this addon"
    )
Campbell Barton's avatar
Campbell Barton committed
    hotkey_list_filter: StringProperty(
        name="        Filter by Name",
        default="",
        description="Show only hotkeys that have this text in their name"
    )
Campbell Barton's avatar
Campbell Barton committed
    show_principled_lists: BoolProperty(
        name="Show Principled naming tags",
        default=False,
        description="Expand this box into a list of all naming tags for principled texture setup"
    )
Campbell Barton's avatar
Campbell Barton committed
    principled_tags: bpy.props.PointerProperty(type=NWPrincipledPreferences)

    def draw(self, context):
        layout = self.layout
        col = layout.column()
        col.prop(self, "merge_position")
        col.prop(self, "merge_hide")

        col = box.column(align=True)
        col.prop(self, "show_principled_lists", text='Edit tags for auto texture detection in Principled BSDF setup', toggle=True)
        if self.show_principled_lists:
            tags = self.principled_tags

            col.prop(tags, "base_color")
            col.prop(tags, "sss_color")
            col.prop(tags, "metallic")
            col.prop(tags, "specular")
            col.prop(tags, "rough")
            col.prop(tags, "gloss")
            col.prop(tags, "normal")
            col.prop(tags, "bump")
            col.prop(tags, "displacement")
        box = layout.box()
        col = box.column(align=True)
        hotkey_button_name = "Show Hotkey List"
        if self.show_hotkey_list:
            hotkey_button_name = "Hide Hotkey List"
        col.prop(self, "show_hotkey_list", text=hotkey_button_name, toggle=True)
        if self.show_hotkey_list:
            col.prop(self, "hotkey_list_filter", icon="VIEWZOOM")
            col.separator()
            for hotkey in kmi_defs:
Bartek Skorupa's avatar
Bartek Skorupa committed
                if hotkey[7]:
                    hotkey_name = hotkey[7]

                    if self.hotkey_list_filter.lower() in hotkey_name.lower():
                        row = col.row(align=True)
                        row.label(text=hotkey_name)
                        keystr = nice_hotkey_name(hotkey[1])
                        if hotkey[4]:
Bartek Skorupa's avatar
Bartek Skorupa committed
                            keystr = "Shift " + keystr
                        if hotkey[5]:
Bartek Skorupa's avatar
Bartek Skorupa committed
                        if hotkey[3]:
def nw_check(context):
    space = context.space_data
    valid_trees = ["ShaderNodeTree", "CompositorNodeTree", "TextureNodeTree"]
    if space.type == 'NODE_EDITOR' and space.node_tree is not None and space.tree_type in valid_trees:
Campbell Barton's avatar
Campbell Barton committed

    return valid
    @classmethod
    def poll(cls, context):
        return nw_check(context)
# OPERATORS
class NWLazyMix(Operator, NWBase):
    """Add a Mix RGB/Shader node by interactively drawing lines between nodes"""
    bl_idname = "node.nw_lazy_mix"
    bl_label = "Mix Nodes"
    bl_options = {'REGISTER', 'UNDO'}

    def modal(self, context, event):
        context.area.tag_redraw()
        nodes, links = get_nodes_links(context)
        cont = True

        start_pos = [event.mouse_region_x, event.mouse_region_y]

        node1 = None
        if not context.scene.NWBusyDrawing:
            node1 = node_at_pos(nodes, context, event)
            if node1:
                context.scene.NWBusyDrawing = node1.name
        else:
            if context.scene.NWBusyDrawing != 'STOP':
                node1 = nodes[context.scene.NWBusyDrawing]

        context.scene.NWLazySource = node1.name
        context.scene.NWLazyTarget = node_at_pos(nodes, context, event).name

        if event.type == 'MOUSEMOVE':
            self.mouse_path.append((event.mouse_region_x, event.mouse_region_y))

        elif event.type == 'RIGHTMOUSE' and event.value == 'RELEASE':
            end_pos = [event.mouse_region_x, event.mouse_region_y]
            bpy.types.SpaceNodeEditor.draw_handler_remove(self._handle, 'WINDOW')

            node2 = None
            node2 = node_at_pos(nodes, context, event)
            if node2:
                context.scene.NWBusyDrawing = node2.name

            if node1 == node2:
                cont = False

            if cont:
                if node1 and node2:
                    for node in nodes:
                        node.select = False
                    node1.select = True
                    node2.select = True

                    bpy.ops.node.nw_merge_nodes(mode="MIX", merge_type="AUTO")

            context.scene.NWBusyDrawing = ""
            return {'FINISHED'}

        elif event.type == 'ESC':
            print('cancelled')
            bpy.types.SpaceNodeEditor.draw_handler_remove(self._handle, 'WINDOW')
            return {'CANCELLED'}

        return {'RUNNING_MODAL'}

    def invoke(self, context, event):
        if context.area.type == 'NODE_EDITOR':
            # the arguments we pass the the callback
            args = (self, context, 'MIX')
            # Add the region OpenGL drawing callback
            # draw in view space with 'POST_VIEW' and 'PRE_VIEW'
            self._handle = bpy.types.SpaceNodeEditor.draw_handler_add(draw_callback_nodeoutline, args, 'WINDOW', 'POST_PIXEL')

            self.mouse_path = []

            context.window_manager.modal_handler_add(self)
            return {'RUNNING_MODAL'}
        else:
            self.report({'WARNING'}, "View3D not found, cannot run operator")
            return {'CANCELLED'}


class NWLazyConnect(Operator, NWBase):
    """Connect two nodes without clicking a specific socket (automatically determined"""
    bl_idname = "node.nw_lazy_connect"
    bl_label = "Lazy Connect"
    bl_options = {'REGISTER', 'UNDO'}
Campbell Barton's avatar
Campbell Barton committed
    with_menu: BoolProperty()

    def modal(self, context, event):
        context.area.tag_redraw()
        nodes, links = get_nodes_links(context)
        cont = True

        start_pos = [event.mouse_region_x, event.mouse_region_y]

        node1 = None
        if not context.scene.NWBusyDrawing:
            node1 = node_at_pos(nodes, context, event)
            if node1:
                context.scene.NWBusyDrawing = node1.name
        else:
            if context.scene.NWBusyDrawing != 'STOP':
                node1 = nodes[context.scene.NWBusyDrawing]

        context.scene.NWLazySource = node1.name
        context.scene.NWLazyTarget = node_at_pos(nodes, context, event).name

        if event.type == 'MOUSEMOVE':
            self.mouse_path.append((event.mouse_region_x, event.mouse_region_y))

        elif event.type == 'RIGHTMOUSE' and event.value == 'RELEASE':
            end_pos = [event.mouse_region_x, event.mouse_region_y]
            bpy.types.SpaceNodeEditor.draw_handler_remove(self._handle, 'WINDOW')

            node2 = None
            node2 = node_at_pos(nodes, context, event)
            if node2:
                context.scene.NWBusyDrawing = node2.name

            if node1 == node2:
                cont = False

            link_success = False
            if cont:
                if node1 and node2:
                    original_sel = []
                    original_unsel = []
                    for node in nodes:
                        if node.select == True:
                            node.select = False
                            original_sel.append(node)
                        else:
                            original_unsel.append(node)
                    node1.select = True
                    node2.select = True

                    #link_success = autolink(node1, node2, links)
                    if self.with_menu:
                        if len(node1.outputs) > 1 and node2.inputs:
                            bpy.ops.wm.call_menu("INVOKE_DEFAULT", name=NWConnectionListOutputs.bl_idname)
                        elif len(node1.outputs) == 1:
                            bpy.ops.node.nw_call_inputs_menu(from_socket=0)
                    else:
                        link_success = autolink(node1, node2, links)

                    for node in original_sel:
                        node.select = True
                    for node in original_unsel:
                        node.select = False

            if link_success:
                force_update(context)
            context.scene.NWBusyDrawing = ""
            return {'FINISHED'}

        elif event.type == 'ESC':
            bpy.types.SpaceNodeEditor.draw_handler_remove(self._handle, 'WINDOW')
            return {'CANCELLED'}

        return {'RUNNING_MODAL'}

    def invoke(self, context, event):
        if context.area.type == 'NODE_EDITOR':
            nodes, links = get_nodes_links(context)
            node = node_at_pos(nodes, context, event)
            if node:
                context.scene.NWBusyDrawing = node.name

            # the arguments we pass the the callback
            mode = "LINK"
            if self.with_menu:
                mode = "LINKMENU"
            args = (self, context, mode)
            # Add the region OpenGL drawing callback
            # draw in view space with 'POST_VIEW' and 'PRE_VIEW'
            self._handle = bpy.types.SpaceNodeEditor.draw_handler_add(draw_callback_nodeoutline, args, 'WINDOW', 'POST_PIXEL')

            self.mouse_path = []

            context.window_manager.modal_handler_add(self)
            return {'RUNNING_MODAL'}
        else:
            self.report({'WARNING'}, "View3D not found, cannot run operator")
            return {'CANCELLED'}


class NWDeleteUnused(Operator, NWBase):
    """Delete all nodes whose output is not used"""
    bl_idname = 'node.nw_del_unused'
    bl_label = 'Delete Unused Nodes'
    bl_options = {'REGISTER', 'UNDO'}

Campbell Barton's avatar
Campbell Barton committed
    delete_muted: BoolProperty(name="Delete Muted", description="Delete (but reconnect, like Ctrl-X) all muted nodes", default=True)
    delete_frames: BoolProperty(name="Delete Empty Frames", description="Delete all frames that have no nodes inside them", default=True)
    def is_unused_node(self, node):
        end_types = ['OUTPUT_MATERIAL', 'OUTPUT', 'VIEWER', 'COMPOSITE', \
                'SPLITVIEWER', 'OUTPUT_FILE', 'LEVELS', 'OUTPUT_LIGHT', \
                'OUTPUT_WORLD', 'GROUP_INPUT', 'GROUP_OUTPUT', 'FRAME']
        if node.type in end_types:
            return False

        for output in node.outputs:
            if output.links:
                return False
        return True

    @classmethod
    def poll(cls, context):
        valid = False
        if nw_check(context):
            if context.space_data.node_tree.nodes:
                valid = True
        return valid

    def execute(self, context):
        nodes, links = get_nodes_links(context)

        # Store selection
        selection = []
        for node in nodes:
            if node.select == True:
                selection.append(node.name)

        for node in nodes:
            node.select = False

        deleted_nodes = []
        temp_deleted_nodes = []
        del_unused_iterations = len(nodes)
        for it in range(0, del_unused_iterations):
            temp_deleted_nodes = list(deleted_nodes)  # keep record of last iteration
            for node in nodes:
                if self.is_unused_node(node):
                    node.select = True
                    deleted_nodes.append(node.name)
                    bpy.ops.node.delete()

            if temp_deleted_nodes == deleted_nodes:  # stop iterations when there are no more nodes to be deleted
                break

        if self.delete_frames:
            repeat = True
            while repeat:
                frames_in_use = []
                frames = []
                repeat = False
                for node in nodes:
                    if node.parent:
                        frames_in_use.append(node.parent)
                for node in nodes:
                    if node.type == 'FRAME' and node not in frames_in_use:
                        frames.append(node)
                        if node.parent:
                            repeat = True  # repeat for nested frames
                for node in frames:
                    if node not in frames_in_use:
                        node.select = True
                        deleted_nodes.append(node.name)
                bpy.ops.node.delete()

        if self.delete_muted:
            for node in nodes:
                if node.mute:
                    node.select = True
                    deleted_nodes.append(node.name)
            bpy.ops.node.delete_reconnect()

        # get unique list of deleted nodes (iterations would count the same node more than once)
        deleted_nodes = list(set(deleted_nodes))
        for n in deleted_nodes:
            self.report({'INFO'}, "Node " + n + " deleted")
        num_deleted = len(deleted_nodes)
        n = ' node'
        if num_deleted > 1:
            n += 's'
        if num_deleted:
            self.report({'INFO'}, "Deleted " + str(num_deleted) + n)
        else:
            self.report({'INFO'}, "Nothing deleted")

        # Restore selection
        nodes, links = get_nodes_links(context)
        for node in nodes:
            if node.name in selection:
                node.select = True
        return {'FINISHED'}

    def invoke(self, context, event):
        return context.window_manager.invoke_confirm(self, event)


class NWSwapLinks(Operator, NWBase):
    """Swap the output connections of the two selected nodes, or two similar inputs of a single node"""
    bl_idname = 'node.nw_swap_links'
    bl_label = 'Swap Links'
    bl_options = {'REGISTER', 'UNDO'}

    @classmethod
    def poll(cls, context):
        valid = False
        if nw_check(context):
            if context.selected_nodes:
                valid = len(context.selected_nodes) <= 2
        return valid

    def execute(self, context):
        nodes, links = get_nodes_links(context)
        selected_nodes = context.selected_nodes
        n1 = selected_nodes[0]

        # Swap outputs
        if len(selected_nodes) == 2:
            n2 = selected_nodes[1]
            if n1.outputs and n2.outputs:
                n1_outputs = []
                n2_outputs = []

                out_index = 0
                for output in n1.outputs:
                    if output.links:
                        for link in output.links:
                            n1_outputs.append([out_index, link.to_socket])
                            links.remove(link)
                    out_index += 1

                out_index = 0
                for output in n2.outputs:
                    if output.links:
                        for link in output.links:
                            n2_outputs.append([out_index, link.to_socket])
                            links.remove(link)
                    out_index += 1

                for connection in n1_outputs:
                    try:
                        links.new(n2.outputs[connection[0]], connection[1])
                    except:
                        self.report({'WARNING'}, "Some connections have been lost due to differing numbers of output sockets")
                for connection in n2_outputs:
                    try:
                        links.new(n1.outputs[connection[0]], connection[1])
                    except:
                        self.report({'WARNING'}, "Some connections have been lost due to differing numbers of output sockets")
            else:
                if n1.outputs or n2.outputs:
                    self.report({'WARNING'}, "One of the nodes has no outputs!")
                else:
                    self.report({'WARNING'}, "Neither of the nodes have outputs!")

        # Swap Inputs
        elif len(selected_nodes) == 1:
            if n1.inputs:
                types = []
                i=0
                for i1 in n1.inputs:
                    if i1.is_linked:
                        similar_types = 0
                        for i2 in n1.inputs:
                            if i1.type == i2.type and i2.is_linked:
                                similar_types += 1
                        types.append ([i1, similar_types, i])
                    i += 1
                types.sort(key=lambda k: k[1], reverse=True)

                if types:
                    t = types[0]
                    if t[1] == 2:
                        for i2 in n1.inputs:
                            if t[0].type == i2.type == t[0].type and t[0] != i2 and i2.is_linked:
                                pair = [t[0], i2]
                        i1f = pair[0].links[0].from_socket
                        i1t = pair[0].links[0].to_socket
                        i2f = pair[1].links[0].from_socket
                        i2t = pair[1].links[0].to_socket
                        links.new(i1f, i2t)
                        links.new(i2f, i1t)
                    if t[1] == 1:
                        if len(types) == 1:
                            fs = t[0].links[0].from_socket
                            i = t[2]
                            links.remove(t[0].links[0])
                            if i+1 == len(n1.inputs):
                                i = -1
                            i += 1
                            while n1.inputs[i].is_linked:
                                i += 1
                            links.new(fs, n1.inputs[i])
                        elif len(types) == 2:
                            i1f = types[0][0].links[0].from_socket
                            i1t = types[0][0].links[0].to_socket
                            i2f = types[1][0].links[0].from_socket
                            i2t = types[1][0].links[0].to_socket
                            links.new(i1f, i2t)
                            links.new(i2f, i1t)

                else:
                    self.report({'WARNING'}, "This node has no input connections to swap!")
            else:
                self.report({'WARNING'}, "This node has no inputs to swap!")
        force_update(context)
        return {'FINISHED'}


class NWResetBG(Operator, NWBase):
    """Reset the zoom and position of the background image"""
    bl_idname = 'node.nw_bg_reset'
    bl_label = 'Reset Backdrop'
    bl_options = {'REGISTER', 'UNDO'}

    @classmethod
    def poll(cls, context):
        valid = False
        if nw_check(context):
            snode = context.space_data
            valid = snode.tree_type == 'CompositorNodeTree'
        return valid

    def execute(self, context):
        context.space_data.backdrop_zoom = 1
        context.space_data.backdrop_offset[0] = 0
        context.space_data.backdrop_offset[1] = 0
        return {'FINISHED'}


class NWAddAttrNode(Operator, NWBase):
    """Add an Attribute node with this name"""
    bl_idname = 'node.nw_add_attr_node'
    bl_label = 'Add UV map'
    bl_options = {'REGISTER', 'UNDO'}

    def execute(self, context):
        bpy.ops.node.add_node('INVOKE_DEFAULT', use_transform=True, type="ShaderNodeAttribute")
        nodes, links = get_nodes_links(context)
        nodes.active.attribute_name = self.attr_name
        return {'FINISHED'}

class NWEmissionViewer(Operator, NWBase):
    bl_idname = "node.nw_emission_viewer"
    bl_label = "Emission Viewer"
    bl_description = "Connect active node to Emission Shader for shadeless previews"
    bl_options = {'REGISTER', 'UNDO'}

    def __init__(self):
        self.shader_output_type = ""
        self.shader_output_ident = ""
        self.shader_viewer_ident = ""

    @classmethod
    def poll(cls, context):
        if nw_check(context):
            space = context.space_data
                if context.active_node:
                    if context.active_node.type != "OUTPUT_MATERIAL" or context.active_node.type != "OUTPUT_WORLD":
                        return True
                else:
                    return True
        return False
    def ensure_viewer_socket(self, node, socket_type, connect_socket=None):
        #check if a viewer output already exists in a node group otherwise create
        if hasattr(node, "node_tree"):
            index = None
            if len(node.node_tree.outputs):
                free_socket = None
                for i, socket in enumerate(node.node_tree.outputs):
                    if is_viewer_socket(socket) and is_visible_socket(node.outputs[i]) and socket.type == socket_type:
                        #if viewer output is already used but leads to the same socket we can still use it
                        is_used = self.is_socket_used_other_mats(socket)
                        if is_used:
                            if connect_socket == None:
                                continue
                            groupout = get_group_output_node(node.node_tree)
                            groupout_input = groupout.inputs[i]
                            links = groupout_input.links
                            if connect_socket not in [link.from_socket for link in links]:
                                continue
                            index=i
                            break
                        if not free_socket:
                            free_socket = i
                if not index and free_socket:
                    index = free_socket

            if not index:
                #create viewer socket
                node.node_tree.outputs.new(socket_type, viewer_socket_name)
                index = len(node.node_tree.outputs) - 1
                node.node_tree.outputs[index].NWViewerSocket = True
            return index

    def init_shader_variables(self, space, shader_type):
        if shader_type == 'OBJECT':
            if space.id not in [light for light in bpy.data.lights]:  # cannot use bpy.data.lights directly as iterable
                self.shader_output_type = "OUTPUT_MATERIAL"
                self.shader_output_ident = "ShaderNodeOutputMaterial"
                self.shader_viewer_ident = "ShaderNodeEmission"
                self.shader_output_type = "OUTPUT_LIGHT"
                self.shader_output_ident = "ShaderNodeOutputLight"
                self.shader_viewer_ident = "ShaderNodeEmission"
        elif shader_type == 'WORLD':
            self.shader_output_type = "OUTPUT_WORLD"
            self.shader_output_ident = "ShaderNodeOutputWorld"
            self.shader_viewer_ident = "ShaderNodeBackground"

    def get_shader_output_node(self, tree):
        for node in tree.nodes:
            if node.type == self.shader_output_type and node.is_active_output == True:
                return node

    @classmethod
    def ensure_group_output(cls, tree):
        #check if a group output node exists otherwise create
        groupout = get_group_output_node(tree)
        if not groupout:
            groupout = tree.nodes.new('NodeGroupOutput')
            loc_x, loc_y = get_output_location(tree)
            groupout.location.x = loc_x
            groupout.location.y = loc_y
            groupout.select = False
        return groupout

    @classmethod
    def search_sockets(cls, node, sockets, index=None):
        #recursevley scan nodes for viewer sockets and store in list
        for i, input_socket in enumerate(node.inputs):
            if index and i != index:
                continue
            if len(input_socket.links):
                link = input_socket.links[0]
                next_node = link.from_node
                external_socket = link.from_socket
                if hasattr(next_node, "node_tree"):
                    for socket_index, s in enumerate(next_node.outputs):
                        if s == external_socket:
                            break
                    socket = next_node.node_tree.outputs[socket_index]
                    if is_viewer_socket(socket) and socket not in sockets:
                        sockets.append(socket)
                        #continue search inside of node group but restrict socket to where we came from
                        groupout = get_group_output_node(next_node.node_tree)
                        cls.search_sockets(groupout, sockets, index=socket_index)

    @classmethod
    def scan_nodes(cls, tree, sockets):
        # get all viewer sockets in a material tree
        for node in tree.nodes:
            if hasattr(node, "node_tree"):
                for socket in node.node_tree.outputs:
                    if is_viewer_socket(socket) and (socket not in sockets):
                        sockets.append(socket)
                cls.scan_nodes(node.node_tree, sockets)

    def link_leads_to_used_socket(self, link):
        #return True if link leads to a socket that is already used in this material
        socket = get_internal_socket(link.to_socket)
        return (socket and self.is_socket_used_active_mat(socket))

    def is_socket_used_active_mat(self, socket):
        #ensure used sockets in active material is calculated and check given socket
        if not hasattr(self, "used_viewer_sockets_active_mat"):
            self.used_viewer_sockets_active_mat = []
            materialout = self.get_shader_output_node(bpy.context.space_data.node_tree)
            if materialout:
                emission = self.get_viewer_node(materialout)
                self.search_sockets((emission if emission else materialout), self.used_viewer_sockets_active_mat)
        return socket in self.used_viewer_sockets_active_mat

    def is_socket_used_other_mats(self, socket):
        #ensure used sockets in other materials are calculated and check given socket
        if not hasattr(self, "used_viewer_sockets_other_mats"):
            self.used_viewer_sockets_other_mats = []
            for mat in bpy.data.materials:
                if mat.node_tree == bpy.context.space_data.node_tree or not hasattr(mat.node_tree, "nodes"):
                    continue
                # get viewer node
                materialout = self.get_shader_output_node(mat.node_tree)
                if materialout:
                    emission = self.get_viewer_node(materialout)
                    self.search_sockets((emission if emission else materialout), self.used_viewer_sockets_other_mats)
        return socket in self.used_viewer_sockets_other_mats

    @staticmethod
    def get_viewer_node(materialout):
        input_socket = materialout.inputs[0]
        if len(input_socket.links) > 0:
            node = input_socket.links[0].from_node
            if node.type == 'EMISSION' and node.name == "Emission Viewer":
                return node

    def invoke(self, context, event):
        space = context.space_data
        shader_type = space.shader_type
        self.init_shader_variables(space, shader_type)
        shader_types = [x[1] for x in shaders_shader_nodes_props]
        mlocx = event.mouse_region_x
        mlocy = event.mouse_region_y
        select_node = bpy.ops.node.select(mouse_x=mlocx, mouse_y=mlocy, extend=False)
        if 'FINISHED' in select_node:  # only run if mouse click is on a node
            active_tree, path_to_tree = get_active_tree(context)
            nodes, links = active_tree.nodes, active_tree.links
            base_node_tree = space.node_tree
            active = nodes.active
            output_types = [x[1] for x in shaders_output_nodes_props]
            valid = False
                if (active.name != "Emission Viewer") and (active.type not in output_types):
                    for out in active.outputs:
                        if is_visible_socket(out):
                materialout = None  # placeholder node
                #scan through all nodes in tree including nodes inside of groups to find viewer sockets
                self.scan_nodes(base_node_tree, delete_sockets)
                materialout = self.get_shader_output_node(base_node_tree)
                if not materialout:
                    materialout = base_node_tree.nodes.new(self.shader_output_ident)
                    materialout.location = get_output_location(base_node_tree)
                # Analyze outputs, add "Emission Viewer" if needed, make links
                out_i = None
                valid_outputs = []
                for i, out in enumerate(active.outputs):
                    if is_visible_socket(out):
                        valid_outputs.append(i)
                if valid_outputs:
                    out_i = valid_outputs[0]  # Start index of node's outputs
                for i, valid_i in enumerate(valid_outputs):
                    for out_link in active.outputs[valid_i].links:
                        if is_viewer_link(out_link, materialout):
                            if nodes == base_node_tree.nodes or self.link_leads_to_used_socket(out_link):
                                if i < len(valid_outputs) - 1:
                                    out_i = valid_outputs[i + 1]
                                else:
                                    out_i = valid_outputs[0]

                make_links = []  # store sockets for new links
                delete_nodes = [] # store unused nodes to delete in the end
                if active.outputs:
                    # If output type not 'SHADER' - "Emission Viewer" needed
                    if active.outputs[out_i].type != 'SHADER':
                        socket_type = 'NodeSocketColor'
                        # get Emission Viewer node
                        emission_exists = False
                        emission_placeholder = base_node_tree.nodes[0]
                        for node in base_node_tree.nodes:
                            if "Emission Viewer" in node.name:
                                emission_exists = True
                                emission_placeholder = node
                        if not emission_exists:
                            emission = base_node_tree.nodes.new(self.shader_viewer_ident)
                            emission.hide = True
                            emission.location = [materialout.location.x, (materialout.location.y + 40)]
                            emission.label = "Viewer"
                            emission.name = "Emission Viewer"
                            emission.use_custom_color = True
                            emission.color = (0.6, 0.5, 0.4)
                            emission.select = False
                        else:
                            emission = emission_placeholder
                        output_socket = emission.inputs[0]

                        # If Viewer is connected to output by user, don't change those connections (patch by gandalf3)
                        if emission.outputs[0].links.__len__() > 0:
                            if not emission.outputs[0].links[0].to_node == materialout:
                                make_links.append((emission.outputs[0], materialout.inputs[0]))
                        else:
                            make_links.append((emission.outputs[0], materialout.inputs[0]))

                        # Set brightness of viewer to compensate for Film and CM exposure
                        if context.scene.render.engine == 'CYCLES' and hasattr(context.scene, 'cycles'):
                            intensity = 1/context.scene.cycles.film_exposure  # Film exposure is a multiplier
                        intensity /= pow(2, (context.scene.view_settings.exposure))  # CM exposure is measured in stops/EVs (2^x)
                        emission.inputs[1].default_value = intensity

                        # Output type is 'SHADER', no Viewer needed. Delete Viewer if exists.
                        socket_type = 'NodeSocketShader'
                        materialout_index = 1 if active.outputs[out_i].name == "Volume" else 0
                        make_links.append((active.outputs[out_i], materialout.inputs[materialout_index]))
                        output_socket = materialout.inputs[materialout_index]
                        for node in base_node_tree.nodes:
                            if node.name == 'Emission Viewer':
                                delete_nodes.append((base_node_tree, node))
                    for li_from, li_to in make_links:
                        base_node_tree.links.new(li_from, li_to)

                    # Crate links through node groups until we reach the active node
                    tree = base_node_tree
                    link_end = output_socket
                    while tree.nodes.active != active:
                        node = tree.nodes.active
                        index = self.ensure_viewer_socket(node, socket_type, connect_socket=active.outputs[out_i] if node.node_tree.nodes.active == active else None)
                        link_start = node.outputs[index]
                        node_socket = node.node_tree.outputs[index]
                        if node_socket in delete_sockets:
                            delete_sockets.remove(node_socket)
                        tree.links.new(link_start, link_end)
                        # Iterate
                        link_end = self.ensure_group_output(node.node_tree).inputs[index]
                        tree = tree.nodes.active.node_tree
                    tree.links.new(active.outputs[out_i], link_end)

                # Delete sockets
                for socket in delete_sockets:
                    if not self.is_socket_used_other_mats(socket):
                        tree = socket.id_data
                        tree.outputs.remove(socket)

                # Delete nodes
                for tree, node in delete_nodes:
                    tree.nodes.remove(node)
                force_update(context)
            return {'FINISHED'}
        else:
            return {'CANCELLED'}


class NWFrameSelected(Operator, NWBase):
    bl_idname = "node.nw_frame_selected"
    bl_label = "Frame Selected"
    bl_description = "Add a frame node and parent the selected nodes to it"
    bl_options = {'REGISTER', 'UNDO'}

    label_prop: StringProperty(
        name='Label',
        description='The visual name of the frame node',
        default=' '
    )
    color_prop: FloatVectorProperty(
        name="Color",
        description="The color of the frame node",
        default=(0.6, 0.6, 0.6),
        min=0, max=1, step=1, precision=3,
        subtype='COLOR_GAMMA', size=3
    )

    def execute(self, context):
        nodes, links = get_nodes_links(context)
        selected = []
        for node in nodes:
            if node.select == True:
                selected.append(node)

        bpy.ops.node.add_node(type='NodeFrame')
        frm = nodes.active
        frm.label = self.label_prop
        frm.use_custom_color = True
        frm.color = self.color_prop

        for node in selected:
            node.parent = frm

        return {'FINISHED'}


class NWReloadImages(Operator, NWBase):
    bl_idname = "node.nw_reload_images"
    bl_label = "Reload Images"
    bl_description = "Update all the image nodes to match their files on disk"

    def execute(self, context):
        nodes, links = get_nodes_links(context)
        image_types = ["IMAGE", "TEX_IMAGE", "TEX_ENVIRONMENT", "TEXTURE"]
        num_reloaded = 0
        for node in nodes:
            if node.type in image_types:
                if node.type == "TEXTURE":
                    if node.texture:  # node has texture assigned
                        if node.texture.type in ['IMAGE', 'ENVIRONMENT_MAP']:
                            if node.texture.image:  # texture has image assigned
                                node.texture.image.reload()
                                num_reloaded += 1
                else:
                    if node.image:
                        node.image.reload()
                        num_reloaded += 1

        if num_reloaded:
            self.report({'INFO'}, "Reloaded images")
            print("Reloaded " + str(num_reloaded) + " images")
            force_update(context)
            return {'FINISHED'}
        else:
            self.report({'WARNING'}, "No images found to reload in this node tree")
            return {'CANCELLED'}


class NWSwitchNodeType(Operator, NWBase):
    """Switch type of selected nodes """
    bl_idname = "node.nw_swtch_node_type"
    bl_label = "Switch Node Type"
    bl_options = {'REGISTER', 'UNDO'}

Campbell Barton's avatar
Campbell Barton committed
    to_type: EnumProperty(
        name="Switch to type",
        items=list(shaders_input_nodes_props) +
        list(shaders_output_nodes_props) +
        list(shaders_shader_nodes_props) +
        list(shaders_texture_nodes_props) +
        list(shaders_color_nodes_props) +
        list(shaders_vector_nodes_props) +
        list(shaders_converter_nodes_props) +
        list(shaders_layout_nodes_props) +
        list(compo_input_nodes_props) +
        list(compo_output_nodes_props) +
        list(compo_color_nodes_props) +
        list(compo_converter_nodes_props) +
        list(compo_filter_nodes_props) +
        list(compo_vector_nodes_props) +
        list(compo_matte_nodes_props) +
        list(compo_distort_nodes_props) +
        list(compo_layout_nodes_props) +
        list(blender_mat_input_nodes_props) +
        list(blender_mat_output_nodes_props) +
        list(blender_mat_color_nodes_props) +
        list(blender_mat_vector_nodes_props) +
        list(blender_mat_converter_nodes_props) +
        list(blender_mat_layout_nodes_props) +
        list(texture_input_nodes_props) +
        list(texture_output_nodes_props) +
        list(texture_color_nodes_props) +
        list(texture_pattern_nodes_props) +
        list(texture_textures_nodes_props) +
        list(texture_converter_nodes_props) +
        list(texture_distort_nodes_props) +
        list(texture_layout_nodes_props)
    )

    def execute(self, context):
        nodes, links = get_nodes_links(context)
        to_type = self.to_type
        # Those types of nodes will not swap.
        src_excludes = ('NodeFrame')
        # Those attributes of nodes will be copied if possible
        attrs_to_pass = ('color', 'hide', 'label', 'mute', 'parent',
                         'show_options', 'show_preview', 'show_texture',
                         'use_alpha', 'use_clamp', 'use_custom_color', 'location'
                         )
        selected = [n for n in nodes if n.select]
        reselect = []
        for node in [n for n in selected if
                     n.rna_type.identifier not in src_excludes and
                     n.rna_type.identifier != to_type]:
            new_node = nodes.new(to_type)
            for attr in attrs_to_pass:
                if hasattr(node, attr) and hasattr(new_node, attr):
                    setattr(new_node, attr, getattr(node, attr))
            # set image datablock of dst to image of src
            if hasattr(node, 'image') and hasattr(new_node, 'image'):
                if node.image:
                    new_node.image = node.image
            # Special cases
            if new_node.type == 'SWITCH':
                new_node.hide = True
            # Dictionaries: src_sockets and dst_sockets:
            # 'INPUTS': input sockets ordered by type (entry 'MAIN' main type of inputs).
            # 'OUTPUTS': output sockets ordered by type (entry 'MAIN' main type of outputs).
            # in 'INPUTS' and 'OUTPUTS':
            # 'SHADER', 'RGBA', 'VECTOR', 'VALUE' - sockets of those types.
            # socket entry:
            # (index_in_type, socket_index, socket_name, socket_default_value, socket_links)
            src_sockets = {
                'INPUTS': {'SHADER': [], 'RGBA': [], 'VECTOR': [], 'VALUE': [], 'MAIN': None},
                'OUTPUTS': {'SHADER': [], 'RGBA': [], 'VECTOR': [], 'VALUE': [], 'MAIN': None},
            }
            dst_sockets = {
                'INPUTS': {'SHADER': [], 'RGBA': [], 'VECTOR': [], 'VALUE': [], 'MAIN': None},
                'OUTPUTS': {'SHADER': [], 'RGBA': [], 'VECTOR': [], 'VALUE': [], 'MAIN': None},
            }
            types_order_one = 'SHADER', 'RGBA', 'VECTOR', 'VALUE'
            types_order_two = 'SHADER', 'VECTOR', 'RGBA', 'VALUE'
            # check src node to set src_sockets values and dst node to set dst_sockets dict values
            for sockets, nd in ((src_sockets, node), (dst_sockets, new_node)):
                # Check node's inputs and outputs and fill proper entries in "sockets" dict
                for in_out, in_out_name in ((nd.inputs, 'INPUTS'), (nd.outputs, 'OUTPUTS')):
                    # enumerate in inputs, then in outputs
                    # find name, default value and links of socket
                    for i, socket in enumerate(in_out):
                        the_name = socket.name
                        dval = None
                        # Not every socket, especially in outputs has "default_value"
                        if hasattr(socket, 'default_value'):
                            dval = socket.default_value
                        socket_links = []
                        for lnk in socket.links:
                            socket_links.append(lnk)
                        # check type of socket to fill proper keys.
                        for the_type in types_order_one:
                            if socket.type == the_type:
                                # create values for sockets['INPUTS'][the_type] and sockets['OUTPUTS'][the_type]
                                # entry structure: (index_in_type, socket_index, socket_name, socket_default_value, socket_links)
                                sockets[in_out_name][the_type].append((len(sockets[in_out_name][the_type]), i, the_name, dval, socket_links))
                    # Check which of the types in inputs/outputs is considered to be "main".
                    # Set values of sockets['INPUTS']['MAIN'] and sockets['OUTPUTS']['MAIN']
                    for type_check in types_order_one:
                        if sockets[in_out_name][type_check]:
                            sockets[in_out_name]['MAIN'] = type_check
                            break

            matches = {
                'INPUTS': {'SHADER': [], 'RGBA': [], 'VECTOR': [], 'VALUE_NAME': [], 'VALUE': [], 'MAIN': []},
                'OUTPUTS': {'SHADER': [], 'RGBA': [], 'VECTOR': [], 'VALUE_NAME': [], 'VALUE': [], 'MAIN': []},
            }

            for inout, soctype in (
                    ('INPUTS', 'MAIN',),
                    ('INPUTS', 'SHADER',),
                    ('INPUTS', 'RGBA',),
                    ('INPUTS', 'VECTOR',),
                    ('INPUTS', 'VALUE',),
                    ('OUTPUTS', 'MAIN',),
                    ('OUTPUTS', 'SHADER',),
                    ('OUTPUTS', 'RGBA',),
                    ('OUTPUTS', 'VECTOR',),
                    ('OUTPUTS', 'VALUE',),
            ):
                if src_sockets[inout][soctype] and dst_sockets[inout][soctype]:
                    if soctype == 'MAIN':
                        sc = src_sockets[inout][src_sockets[inout]['MAIN']]
                        dt = dst_sockets[inout][dst_sockets[inout]['MAIN']]
                    else:
                        sc = src_sockets[inout][soctype]
                        dt = dst_sockets[inout][soctype]
                    # start with 'dt' to determine number of possibilities.
                    for i, soc in enumerate(dt):
                        # if src main has enough entries - match them with dst main sockets by indexes.
                        if len(sc) > i:
                            matches[inout][soctype].append(((sc[i][1], sc[i][3]), (soc[1], soc[3])))
                        # add 'VALUE_NAME' criterion to inputs.
                        if inout == 'INPUTS' and soctype == 'VALUE':
                            for s in sc:
                                if s[2] == soc[2]:  # if names match
Loading
Loading full blame…