node_Lathe.py

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

import math
from math import pi, radians

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, angle, steps):
    bm = bmesh_from_pydata(verts, edges, [])
    geom = bm.verts[:] + bm.edges[:]

    spin(
        bm, geom=geom, cent=cent, axis=axis,
        dvec=dvec, angle=angle, 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', update=updateNode)
    dist = FloatProperty(name="merge distance", default=0.0001, update=updateNode)
    Degrees = FloatProperty(name="Degrees", default=360.0, update=updateNode)
    Steps = IntProperty(name="Steps", default=20, update=updateNode)
    cent = FloatVectorProperty(name='cent', size=3, update=updateNode)
    dvec = FloatVectorProperty(name='dvec', size=3, update=updateNode)
    axis = FloatVectorProperty(
        name='axis', 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')

        # 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', 'Degrees', 'Degrees').prop_name = 'Degrees'
        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)
        row.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],
            ['Degrees', 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)

        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 get_final_values_from_data(self, data):
        cent, axis, dvec, Degrees, Steps = data

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

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

        defaults = {
            '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(),
            'angle': radians(Degrees),
            'steps': max(Steps, 0)
        }
        return defaults

    def nothing_to_process(self):
        if not ('Poly' in self.outputs):
            return True

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

    def update(self):
        if self.nothing_to_process():
            return

        mverts, *mrest = self.get_sockets_data()

        def yield_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, *data = yield_socket_data(idx)
            final_values = self.get_final_values_from_data(data)
            final_values['verts'] = Verts
            final_values['edges'] = edges

            v, p = get_lathed_geometry(self, **final_values)
            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)