vertex layer comp node.py

bm_vlayer_compnode.py

# ##### 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 #####

import bpy
from bpy.props import IntProperty, EnumProperty, BoolProperty, FloatProperty
import bmesh
import bmesh.ops

from sverchok.node_tree import SverchCustomTreeNode
from sverchok.data_structure import updateNode, match_long_repeat, fullList
from sverchok.utils.sv_bmesh_utils import bmesh_from_pydata, pydata_from_bmesh

def set_data_for_layer(bm, data, layer):
    for i in range(len(bm.verts)):
        bm.verts[i][data] = layer[i] or 0.1


def shrink_geometry(bm, dist, layers):

    # should detect float/int/
    made_layers = []
    for idx, layer in enumerate(layers):
        
        first_element = layer[0] or 0.2
        layer_name = '_layer' + str(idx)
        layers = bm.verts.layers
        current_type = ''

        if isinstance(first_element, float):
            current_type = 'float'
            data_layer = layers.float.new(current_type + layer_name)
        elif isinstance(first_element, int):
            current_type = 'int'
            data_layer = layers.int.new(current_type + layer_name)

        made_layers.append(data_layer)
        bm.verts.ensure_lookup_table()
        set_data_for_layer(bm, data_layer, layer)

    bmesh.ops.remove_doubles(bm, verts=bm.verts[:], dist=dist)
    bm.verts.ensure_lookup_table()
    
    verts, edges, faces = pydata_from_bmesh(bm)
    data_out = [verts, edges, faces]
    for layer_name in made_layers:
        # is there a foreach for this?
        # something like
        #    bm.verts.layers.foreach_get(layer_name, receiver_list)
        data_out.append([bm.verts[i][layer_name] for i in range(len(bm.verts))])

    return data_out


class SvBmeshRemoveDoublesKeepDatalayer(bpy.types.Node, SverchCustomTreeNode):
    ''' Limited Dissolve '''
    bl_idname = 'SvBmeshRemoveDoublesKeepDatalayer'
    bl_label = 'bm rm Doubles KeepDatalayers'
    bl_icon = 'OUTLINER_OB_EMPTY'

    distance = FloatProperty(default=0.001, min=0.0, update=updateNode)
    num_extra_sockets = IntProperty(default=0)

    def sv_init(self, context):
        for idx, sockets in enumerate([self.inputs, self.outputs]):
            sockets.new('VerticesSocket', 'Verts')
            sockets.new('StringsSocket', 'Edges')
            sockets.new('StringsSocket', 'Polys')
            sockets.new('StringsSocket', 'Data A')
            sockets.new('StringsSocket', 'Data B')


    def draw_buttons(self, context, layout):
        layout.prop(self, "distance")

    def process(self):

        in_linked = lambda x: self.inputs[x].is_linked
        out_linked = lambda x: self.outputs[x].is_linked

        if not in_linked('Verts'):
            return
        else:
            if not (in_linked('Polys') or in_linked('Edges')):
                return

        if not (any(out_linked(x) for x in ['Verts', 'Edges', 'Polys'])):
            return

        input_data = [socket.sv_get(default=[[]]) for socket in self.inputs]
        all_data = match_long_repeat(input_data)

        collected = [[], [], [], [], []]
        for verts, edges, faces, d1, d2 in zip(*all_data):
            bm = bmesh_from_pydata(verts, edges, faces)
            for idx, content in enumerate(shrink_geometry(bm, self.distance, [d1, d2])):
                collected[idx].append(content)
            bm.free()

        for idx, socket in enumerate(self.outputs):
            socket.sv_set(collected[idx])


def register():
    bpy.utils.register_class(SvBmeshRemoveDoublesKeepDatalayer)


def unregister():
    bpy.utils.unregister_class(SvBmeshRemoveDoublesKeepDatalayer)