move generic bpy helper modules into bpy_extras.
[blender.git] / release / scripts / startup / bl_operators / add_mesh_torus.py
1 # ##### BEGIN GPL LICENSE BLOCK #####
2 #
3 #  This program is free software; you can redistribute it and/or
4 #  modify it under the terms of the GNU General Public License
5 #  as published by the Free Software Foundation; either version 2
6 #  of the License, or (at your option) any later version.
7 #
8 #  This program is distributed in the hope that it will be useful,
9 #  but WITHOUT ANY WARRANTY; without even the implied warranty of
10 #  MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
11 #  GNU General Public License for more details.
12 #
13 #  You should have received a copy of the GNU General Public License
14 #  along with this program; if not, write to the Free Software Foundation,
15 #  Inc., 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301, USA.
16 #
17 # ##### END GPL LICENSE BLOCK #####
18
19 # <pep8 compliant>
20 import bpy
21 import mathutils
22
23
24 def add_torus(major_rad, minor_rad, major_seg, minor_seg):
25     from math import cos, sin, pi
26
27     Vector = mathutils.Vector
28     Quaternion = mathutils.Quaternion
29
30     PI_2 = pi * 2.0
31     z_axis = 0.0, 0.0, 1.0
32
33     verts = []
34     faces = []
35     i1 = 0
36     tot_verts = major_seg * minor_seg
37     for major_index in range(major_seg):
38         quat = Quaternion(z_axis, (major_index / major_seg) * PI_2)
39
40         for minor_index in range(minor_seg):
41             angle = 2 * pi * minor_index / minor_seg
42
43             vec = Vector((major_rad + (cos(angle) * minor_rad), 0.0,
44                         (sin(angle) * minor_rad))) * quat
45
46             verts.extend(vec[:])
47
48             if minor_index + 1 == minor_seg:
49                 i2 = (major_index) * minor_seg
50                 i3 = i1 + minor_seg
51                 i4 = i2 + minor_seg
52
53             else:
54                 i2 = i1 + 1
55                 i3 = i1 + minor_seg
56                 i4 = i3 + 1
57
58             if i2 >= tot_verts:
59                 i2 = i2 - tot_verts
60             if i3 >= tot_verts:
61                 i3 = i3 - tot_verts
62             if i4 >= tot_verts:
63                 i4 = i4 - tot_verts
64
65             # stupid eekadoodle
66             if i2:
67                 faces.extend([i1, i3, i4, i2])
68             else:
69                 faces.extend([i2, i1, i3, i4])
70
71             i1 += 1
72
73     return verts, faces
74
75 from bpy.props import FloatProperty, IntProperty, BoolProperty, FloatVectorProperty
76
77
78 class AddTorus(bpy.types.Operator):
79     '''Add a torus mesh'''
80     bl_idname = "mesh.primitive_torus_add"
81     bl_label = "Add Torus"
82     bl_options = {'REGISTER', 'UNDO'}
83
84     major_radius = FloatProperty(name="Major Radius",
85             description="Radius from the origin to the center of the cross sections",
86             default=1.0, min=0.01, max=100.0)
87     minor_radius = FloatProperty(name="Minor Radius",
88             description="Radius of the torus' cross section",
89             default=0.25, min=0.01, max=100.0)
90     major_segments = IntProperty(name="Major Segments",
91             description="Number of segments for the main ring of the torus",
92             default=48, min=3, max=256)
93     minor_segments = IntProperty(name="Minor Segments",
94             description="Number of segments for the minor ring of the torus",
95             default=12, min=3, max=256)
96     use_abso = BoolProperty(name="Use Int+Ext Controls",
97             description="Use the Int / Ext controls for torus dimensions",
98             default=False)
99     abso_major_rad = FloatProperty(name="Exterior Radius",
100             description="Total Exterior Radius of the torus",
101             default=1.0, min=0.01, max=100.0)
102     abso_minor_rad = FloatProperty(name="Inside Radius",
103             description="Total Interior Radius of the torus",
104             default=0.5, min=0.01, max=100.0)
105
106     # generic transform props
107     view_align = BoolProperty(name="Align to View",
108             default=False)
109     location = FloatVectorProperty(name="Location",
110             subtype='TRANSLATION')
111     rotation = FloatVectorProperty(name="Rotation",
112             subtype='EULER')
113
114     def execute(self, context):
115
116         if self.use_abso == True:
117             extra_helper = (self.abso_major_rad - self.abso_minor_rad) * 0.5
118             self.major_radius = self.abso_minor_rad + extra_helper
119             self.minor_radius = extra_helper
120
121         verts_loc, faces = add_torus(self.major_radius,
122                                     self.minor_radius,
123                                     self.major_segments,
124                                     self.minor_segments)
125
126         mesh = bpy.data.meshes.new("Torus")
127
128         mesh.vertices.add(len(verts_loc) // 3)
129         mesh.faces.add(len(faces) // 4)
130
131         mesh.vertices.foreach_set("co", verts_loc)
132         mesh.faces.foreach_set("vertices_raw", faces)
133         mesh.update()
134
135         from bpy_extras import object_utils
136         object_utils.object_data_add(context, mesh, operator=self)
137
138         return {'FINISHED'}