- dir() now works for collection functions
[blender-staging.git] / release / scripts / op / 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., 59 Temple Place - Suite 330, Boston, MA  02111-1307, USA.
16 #
17 # ##### END GPL LICENSE BLOCK #####
18
19 # <pep8-80 compliant>
20 import bpy
21 import Mathutils
22 from math import cos, sin, pi
23
24
25 def add_torus(major_rad, minor_rad, major_seg, minor_seg):
26     Vector = Mathutils.Vector
27     Quaternion = Mathutils.Quaternion
28
29     PI_2 = pi * 2
30     z_axis = (0, 0, 1)
31
32     verts = []
33     faces = []
34     i1 = 0
35     tot_verts = major_seg * minor_seg
36     for major_index in range(major_seg):
37         quat = Quaternion(z_axis, (major_index / major_seg) * PI_2)
38
39         for minor_index in range(minor_seg):
40             angle = 2 * pi * minor_index / minor_seg
41
42             vec = Vector(major_rad + (cos(angle) * minor_rad), 0.0,
43                         (sin(angle) * minor_rad)) * quat
44
45             verts.extend([vec.x, vec.y, vec.z])
46
47             if minor_index + 1 == minor_seg:
48                 i2 = (major_index) * minor_seg
49                 i3 = i1 + minor_seg
50                 i4 = i2 + minor_seg
51
52             else:
53                 i2 = i1 + 1
54                 i3 = i1 + minor_seg
55                 i4 = i3 + 1
56
57             if i2 >= tot_verts:
58                 i2 = i2 - tot_verts
59             if i3 >= tot_verts:
60                 i3 = i3 - tot_verts
61             if i4 >= tot_verts:
62                 i4 = i4 - tot_verts
63
64             # stupid eekadoodle
65             if i2:
66                 faces.extend([i1, i3, i4, i2])
67             else:
68                 faces.extend([i2, i1, i3, i4])
69
70             i1 += 1
71
72     return verts, faces
73
74 from bpy.props import *
75
76
77 class AddTorus(bpy.types.Operator):
78     '''Add a torus mesh.'''
79     bl_idname = "mesh.primitive_torus_add"
80     bl_label = "Add Torus"
81     bl_register = True
82     bl_undo = True
83
84     major_radius = FloatProperty(name="Major Radius",
85             description="Number of segments for the main ring of the torus",
86             default=1.0, min=0.01, max=100.0)
87     minor_radius = FloatProperty(name="Minor Radius",
88             description="Number of segments for the minor ring of the torus",
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=16, min=3, max=256)
96
97     def execute(self, context):
98
99         verts_loc, faces = add_torus(self.properties.major_radius,
100                                     self.properties.minor_radius,
101                                     self.properties.major_segments,
102                                     self.properties.minor_segments)
103
104         mesh = bpy.data.add_mesh("Torus")
105
106         mesh.add_geometry(int(len(verts_loc) / 3), 0, int(len(faces) / 4))
107         mesh.verts.foreach_set("co", verts_loc)
108         mesh.faces.foreach_set("verts_raw", faces)
109
110         scene = context.scene
111
112         # ugh
113         for ob in scene.objects:
114             ob.selected = False
115
116         mesh.update()
117         ob_new = bpy.data.add_object('MESH', "Torus")
118         ob_new.data = mesh
119         scene.objects.link(ob_new)
120         scene.objects.active = ob_new
121         ob_new.selected = True
122
123         ob_new.location = tuple(context.scene.cursor_location)
124
125         return ('FINISHED',)
126
127 # Register the operator
128 bpy.ops.add(AddTorus)
129
130 # Add to a menu
131 import dynamic_menu
132
133 menu_func = (lambda self, context: self.layout.itemO(AddTorus.bl_idname,
134                                         text="Torus", icon='ICON_MESH_DONUT'))
135
136 menu_item = dynamic_menu.add(bpy.types.INFO_MT_mesh_add, menu_func)
137
138 if __name__ == "__main__":
139     bpy.ops.mesh.primitive_torus_add()