node_Lathe.py

import bpy
import bmesh
import mathutils
from mathutils import Vector, Matrix
from bpy.props import BoolProperty, FloatProperty, FloatVectorProperty

from node_s import *
from util import *

import sv_bmesh_utils
from sv_bmesh_utils import bmesh_from_pydata


def get_lathed_geometry(node, verts, edges, cent, axis, dvec, space, angle, steps):

    # steps = max(steps, 0)
    bm = bmesh_from_pydata(verts, edges, [])
    geom = bm.verts[:] + bm.edges[:]

    bmesh.ops.spin(bm,
                   geom=geom, cent=cent, axis=axis, dvec=dvec, angle=angle,
                   space=space, steps=steps, use_duplicate=False)

    if node.remove_doubles:
        bmesh.ops.remove_doubles(bm, verts=bm.verts[:], dist=node.dist)

    v = [v.co[:] for v in bm.verts]
    p = [[i.index for i in p.verts] for p in bm.faces[:]]
    bm.free()
    return v, p


class SvLatheNode(Node, SverchCustomTreeNode):

    bl_idname = 'SvLatheNode'
    bl_label = 'Sv Lathe Node'
    bl_icon = 'OUTLINER_OB_EMPTY'

    remove_doubles = BoolProperty(
        name='merge', description='Remove doubles, merge',
        default=False,
        update=updateNode)

    dist = FloatProperty(name="merge distance", default=0.0001, update=updateNode)
    angle = FloatProperty(name="angle", default=1.0, update=updateNode)
    steps = IntProperty(name="steps", default=20, update=updateNode)
    cent = FloatVectorProperty(name='cent', subtype='XYZ', size=3, update=updateNode)
    dvec = FloatVectorProperty(name='dvec', subtype='XYZ', size=3, update=updateNode)
    axis = FloatVectorProperty(
        name='axis', subtype='XYZ', size=3, update=updateNode, default=(0, 0, 1))

    def init(self, context):

        # so many inputs!!
        self.inputs.new('VerticesSocket', 'Verts', 'Verts')
        self.inputs.new('StringsSocket', 'Edges', 'Edges')
        self.inputs.new('MatrixSocket', 'Matrix', 'Matrix')

        # these have default props
        self.inputs.new('VerticesSocket', 'cent', 'cent').prop_name = 'cent'
        self.inputs.new('VerticesSocket', 'axis', 'axis').prop_name = 'axis'
        self.inputs.new('VerticesSocket', 'dvec', 'dvec').prop_name = 'dvec'
        self.inputs.new('StringsSocket', 'Angle', 'Angle').prop_name = 'angle'
        self.inputs.new('StringsSocket', 'Steps', 'Steps').prop_name = 'steps'

        # outputs, sparse!
        self.outputs.new('VerticesSocket', 'Verts', 'Verts')
        self.outputs.new('StringsSocket', 'Poly', 'Poly')

    def draw_buttons(self, context, layout):
        row = layout.row(align=True)
        split = row.split()
        col1 = split.column()
        col1.prop(self, "remove_doubles", text="merge")

    def get_corrected_data(self, socket_name, socket_type):
        inputs = self.inputs
        socket = inputs[socket_name].links[0].from_socket
        if isinstance(socket, socket_type):
            socket_in = SvGetSocketAnyType(self, inputs[socket_name])
            return dataCorrect(socket_in)
        else:
            return []

    def get_sockets_data(self):
        inputs = self.inputs
        mverts = self.get_corrected_data('Verts', VerticesSocket)

        socket_inputs = [
            ['Edges', StringsSocket],
            ['cent', VerticesSocket],
            ['axis', VerticesSocket],
            ['dvec', VerticesSocket],
            ['Matrix', MatrixSocket],
            ['Angle', StringsSocket],
            ['Steps', StringsSocket]
        ]

        all_socket_data = [mverts]
        for named_socket, socket_type in socket_inputs:
            socket_data = []
            if named_socket in inputs and inputs[named_socket].links:
                socket_data = self.get_corrected_data(named_socket, socket_type)
            all_socket_data.append(socket_data)

        # woah, suck on this one!
        return all_socket_data

    def get_structure(self, stype, sindex):
        if not stype:
            return []

        try:
            j = stype[sindex]
        except IndexError:
            j = []
        finally:
            return j

    def update(self):
        inputs = self.inputs
        outputs = self.outputs

        if not ('Poly' in outputs):
            return

        if not (inputs['Verts'].links and outputs['Verts'].links):
            return

        mverts, *mrest = self.get_sockets_data()

        def get_non_vert_socket_data(obj_index):
            for geom in mrest:
                yield self.get_structure(geom, obj_index)

        verts_out, faces_out = [], []
        for idx, Verts in enumerate(mverts):
            if not Verts:
                continue

            edges, cent, axis, dvec, space, angle, steps = \
                get_non_vert_socket_data(idx)

            if isinstance(angle, list) and len(angle)>0:
                angle = angle[0]
            else:
                angle = self.angle

            if isinstance(steps, list) and len(steps)>0:
                steps = steps[0]
            else:
                steps = self.steps

            defaults = {
                'verts': Verts,
                'edges': edges,
                'cent': Vector(cent[0]) if cent else Vector(),
                'axis': Vector(axis[0]) if axis else Vector((0, 0, 1)),
                'dvec': Vector(dvec[0]) if dvec else Vector(),
                'space': Matrix(space) if space else Matrix(),
                'angle': angle,
                'steps': steps
            }

            if not defaults['edges']:
                del defaults['edges']

            v, p = get_lathed_geometry(self, **defaults)
            verts_out.append(v)
            faces_out.append(p)

        SvSetSocketAnyType(self, 'Verts', verts_out)
        SvSetSocketAnyType(self, 'Poly', faces_out)

    def update_socket(self, context):
        self.update()


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


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