implementation of adding spirals (add_curve_spirals) in edit mode

2 files changed, 181 insertions, 37 deletions

diff --git a/add_curve_extra_objects/__init__.py b/add_curve_extra_objects/__init__.py
index d7fbe2a4..1653ad86 100644
--- a/add_curve_extra_objects/__init__.py
+++ b/add_curve_extra_objects/__init__.py
@@ -247,10 +247,10 @@ def menu_func(self, context):
     layout = self.layout
 
     layout.operator_menu_enum("curve.curveaceous_galore", "ProfileType", icon='CURVE_DATA')
+    layout.operator_menu_enum("curve.spirals", "spiral_type", icon='CURVE_DATA')
     if context.mode != 'OBJECT':
         # fix in D2142 will allow to work in EDIT_CURVE
         return None
-    layout.operator_menu_enum("curve.spirals", "spiral_type", icon='CURVE_DATA')
     layout.separator()
 
     layout.menu(INFO_MT_curve_knots_add.bl_idname, text="Knots", icon='CURVE_DATA')
diff --git a/add_curve_extra_objects/add_curve_spirals.py b/add_curve_extra_objects/add_curve_spirals.py
index 7bb225d0..96eb8aba 100644
--- a/add_curve_extra_objects/add_curve_spirals.py
+++ b/add_curve_extra_objects/add_curve_spirals.py
@@ -24,6 +24,11 @@ from bpy.props import (
         BoolProperty,
         FloatProperty,
         IntProperty,
+        FloatVectorProperty
+        )
+from mathutils import (
+        Vector,
+        Matrix,
         )
 from math import (
         sin, cos, pi
@@ -58,7 +63,7 @@ def make_spiral(props, context):
     step_z = props.z_scale / (steps - 1)  # z increase in one step
 
     verts = []
-    verts.extend([props.radius, 0, 0, 1])
+    verts.append([props.radius, 0, 0])
 
     cur_phi = 0
     cur_z = 0
@@ -82,7 +87,7 @@ def make_spiral(props, context):
 
         px = cur_rad * cos(cur_phi)
         py = cur_rad * sin(cur_phi)
-        verts.extend([px, py, cur_z, 1])
+        verts.append([px, py, cur_z])
 
     return verts
 
@@ -104,7 +109,7 @@ def make_spiral_spheric(props, context):
     step_theta = pi / (steps - 1)  # theta increase in one step (pi == 180 deg)
 
     verts = []
-    verts.extend([0, 0, -props.radius, 1])  # First vertex at south pole
+    verts.append([0, 0, -props.radius])  # First vertex at south pole
 
     cur_phi = 0
     cur_theta = -pi / 2  # Beginning at south pole
@@ -115,7 +120,7 @@ def make_spiral_spheric(props, context):
         py = props.radius * cos(cur_theta) * sin(cur_phi)
         pz = props.radius * sin(cur_theta)
 
-        verts.extend([px, py, pz, 1])
+        verts.append([px, py, pz])
         cur_theta += step_theta
         cur_phi += step_phi
 
@@ -157,7 +162,7 @@ def make_spiral_torus(props, context):
               sin(props.curves_number * cur_theta)
         pz = cur_inner_rad * sin(cur_phi) + cur_z
 
-        verts.extend([px, py, pz, 1])
+        verts.append([px, py, pz])
 
         if props.touch and cur_phi >= n_cycle * 2 * pi:
             step_z = ((n_cycle + 1) * props.dif_inner_radius +
@@ -172,8 +177,50 @@ def make_spiral_torus(props, context):
 
     return verts
 
-
-def draw_curve(props, context):
+# ------------------------------------------------------------
+# calculates the matrix for the new object
+# depending on user pref
+
+def align_matrix(context, location):
+    loc = Matrix.Translation(location)
+    obj_align = context.preferences.edit.object_align
+    if (context.space_data.type == 'VIEW_3D' and
+            obj_align == 'VIEW'):
+        rot = context.space_data.region_3d.view_matrix.to_3x3().inverted().to_4x4()
+    else:
+        rot = Matrix()
+    align_matrix = loc @ rot
+
+    return align_matrix
+
+# ------------------------------------------------------------
+# get array of vertcoordinates according to splinetype
+def vertsToPoints(Verts, splineType):
+
+    # main vars
+    vertArray = []
+
+    # array for BEZIER spline output (V3)
+    if splineType == 'BEZIER':
+        for v in Verts:
+            vertArray += v
+
+    # array for nonBEZIER output (V4)
+    else:
+        for v in Verts:
+            vertArray += v
+            if splineType == 'NURBS':
+                # for nurbs w=1
+                vertArray.append(1)
+            else:
+                # for poly w=0
+                vertArray.append(0)
+    return vertArray
+
+def draw_curve(props, context, align_matrix):
+    # output splineType 'POLY' 'NURBS' 'BEZIER'
+    splineType = props.curve_type
+    
     if props.spiral_type == 'ARCH':
         verts = make_spiral(props, context)
     if props.spiral_type == 'LOG':
@@ -183,22 +230,44 @@ def draw_curve(props, context):
     if props.spiral_type == 'TORUS':
         verts = make_spiral_torus(props, context)
 
-    curve_data = bpy.data.curves.new(name='Spiral', type='CURVE')
-    curve_data.dimensions = '3D'
-
-    spline = curve_data.splines.new(type=props.curve_type)
-    """
-    if props.curve_type == 0:
-        spline = curve_data.splines.new(type='POLY')
-    elif props.curve_type == 1:
-        spline = curve_data.splines.new(type='NURBS')
-    """
-    spline.points.add(len(verts) * 0.25 - 1)
-    # Add only one quarter of points as elements in verts,
-    # because verts looks like: "x,y,z,?,x,y,z,?,x,..."
-    spline.points.foreach_set('co', verts)
-    new_obj = object_data_add(context, curve_data)
-
+    # create object
+    if bpy.context.mode == 'EDIT_CURVE':
+        Curve = context.active_object
+        spline = Curve.data.splines.new(type=splineType)          # spline
+    else:
+        # create curve
+        newCurve = bpy.data.curves.new(name='Spiral', type='CURVE')  # curvedatablock
+        spline = newCurve.splines.new(type=splineType)          # spline
+
+        # set curveOptions
+        newCurve.dimensions = props.shape
+        
+        # create object with newCurve
+        Curve = object_data_add(context, newCurve)  # place in active scene
+        Curve.select_set(True)
+        
+    Curve.matrix_world = align_matrix  # apply matrix
+    Curve.rotation_euler = props.rotation_euler
+        
+    # set curveOptions
+    spline.use_cyclic_u = props.use_cyclic_u
+    spline.use_endpoint_u = props.endp_u
+    spline.order_u = props.order_u
+    
+    # turn verts into array
+    vertArray = vertsToPoints(verts, splineType)
+        
+    # create spline from vertarray
+    if splineType == 'BEZIER':
+        spline.bezier_points.add(int(len(vertArray) * 0.33))
+        spline.bezier_points.foreach_set('co', vertArray)
+        for point in spline.bezier_points:
+            point.handle_right_type = props.handleType
+            point.handle_left_type = props.handleType
+    else:
+        spline.points.add(int(len(vertArray) * 0.25 - 1))
+        spline.points.foreach_set('co', vertArray)
+        spline.use_endpoint_u = False
 
 class CURVE_OT_spirals(Operator):
     bl_idname = "curve.spirals"
@@ -206,6 +275,9 @@ class CURVE_OT_spirals(Operator):
     bl_description = "Create different types of spirals"
     bl_options = {'REGISTER', 'UNDO', 'PRESET'}
 
+    # align_matrix for the invoke
+    align_matrix : Matrix()
+    
     spiral_type : EnumProperty(
             items=[('ARCH', "Archemedian", "Archemedian"),
                    ("LOG", "Logarithmic", "Logarithmic"),
@@ -215,13 +287,6 @@ class CURVE_OT_spirals(Operator):
             name="Spiral Type",
             description="Type of spiral to add"
             )
-    curve_type : EnumProperty(
-            items=[('POLY', "Poly", "PolyLine"),
-                   ("NURBS", "NURBS", "NURBS")],
-            default='POLY',
-            name="Curve Type",
-            description="Type of spline to use"
-            )
     spiral_direction : EnumProperty(
             items=[('COUNTER_CLOCKWISE', "Counter Clockwise",
                     "Wind in a counter clockwise direction"),
@@ -295,6 +360,62 @@ class CURVE_OT_spirals(Operator):
             default=False,
             description="No empty spaces between cycles"
             )
+    # Curve Options
+    shapeItems = [
+        ('2D', "2D", "2D shape Curve"),
+        ('3D', "3D", "3D shape Curve")]
+    shape : EnumProperty(
+            name="2D / 3D",
+            items=shapeItems,
+            description="2D or 3D Curve",
+            default='3D'
+            )
+    curve_type : EnumProperty(
+            name="Output splines",
+            description="Type of splines to output",
+            items=[
+            ('POLY', "Poly", "Poly Spline type"),
+            ('NURBS', "Nurbs", "Nurbs Spline type"),
+            ('BEZIER', "Bezier", "Bezier Spline type")],
+            default='POLY'
+            )
+    use_cyclic_u : BoolProperty(
+            name="Cyclic",
+            default=False,
+            description="make curve closed"
+            )
+    endp_u : BoolProperty(
+            name="Use endpoint u",
+            default=True,
+            description="stretch to endpoints"
+            )
+    order_u : IntProperty(
+            name="Order u",
+            default=4,
+            min=2, soft_min=2,
+            max=6, soft_max=6,
+            description="Order of nurbs spline"
+            )
+    handleType : EnumProperty(
+            name="Handle type",
+            default='VECTOR',
+            description="Bezier handles type",
+            items=[
+            ('VECTOR', "Vector", "Vector type Bezier handles"),
+            ('AUTO', "Auto", "Automatic type Bezier handles")]
+            )
+    startlocation : FloatVectorProperty(
+            name="",
+            description="Start location",
+            default=(0.0, 0.0, 0.0),
+            subtype='TRANSLATION'
+            )
+    rotation_euler : FloatVectorProperty(
+            name="",
+            description="Rotation",
+            default=(0.0, 0.0, 0.0),
+            subtype='EULER'
+            )
 
     def draw(self, context):
         layout = self.layout
@@ -305,12 +426,11 @@ class CURVE_OT_spirals(Operator):
         row = col.row(align=True)
         row.menu("OBJECT_MT_spiral_curve_presets",
                  text=bpy.types.OBJECT_MT_spiral_curve_presets.bl_label)
-        row.operator("curve_extras.spiral_presets", text="")
-        op = row.operator("curve_extras.spiral_presets", text="")
+        row.operator("curve_extras.spiral_presets", text=" + ")
+        op = row.operator("curve_extras.spiral_presets", text=" - ")
         op.remove_active = True
 
         layout.prop(self, "spiral_type")
-        layout.prop(self, "curve_type")
         layout.prop(self, "spiral_direction")
 
         col = layout.column(align=True)
@@ -360,16 +480,40 @@ class CURVE_OT_spirals(Operator):
             col.prop(self, "dif_radius", text="Increase of Torus Radius")
             col.prop(self, "dif_inner_radius", text="Increase of Inner Radius")
 
+        row = layout.row()
+        row.prop(self, "shape", expand=True)
+        
+        # output options
+        col = layout.column()
+        col.label(text="Output Curve Type:")
+        col.row().prop(self, "curve_type", expand=True)
+        
+        if self.curve_type == 'NURBS':
+            col.prop(self, "order_u")
+        elif self.curve_type == 'BEZIER':
+            col.row().prop(self, 'handleType', expand=True)
+
+        col = layout.column()
+        col.row().prop(self, "use_cyclic_u", expand=True)
+        
+        box = layout.box()
+        box.label(text="Location:")
+        box.prop(self, "startlocation")
+        box = layout.box()
+        box.label(text="Rotation:")
+        box.prop(self, "rotation_euler")
+
     @classmethod
     def poll(cls, context):
         return context.scene is not None
 
     def execute(self, context):
         time_start = time.time()
-        draw_curve(self, context)
+        self.align_matrix = align_matrix(context, self.startlocation)
+        draw_curve(self, context, self.align_matrix)
 
-        self.report({'INFO'},
-                    "Drawing Spiral Finished: %.4f sec" % (time.time() - time_start))
+        #self.report({'INFO'},
+                    #"Drawing Spiral Finished: %.4f sec" % (time.time() - time_start))
 
         return {'FINISHED'}