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
21"""
22Import and export STL files
23
24Used as a blender script, it load all the stl files in the scene:
25
26blender --python stl_utils.py -- file1.stl file2.stl file3.stl ...
27"""
28
29# TODO: endien
30
31
32class ListDict(dict):
33    """
34    Set struct with order.
35
36    You can:
37       - insert data into without doubles
38       - get the list of data in insertion order with self.list
39
40    Like collections.OrderedDict, but quicker, can be replaced if
41    ODict is optimised.
42    """
43
44    def __init__(self):
45        dict.__init__(self)
46        self.list = []
47        self._len = 0
48
49    def add(self, item):
50        """
51        Add a value to the Set, return its position in it.
52        """
53        value = self.setdefault(item, self._len)
54        if value == self._len:
55            self.list.append(item)
56            self._len += 1
57
58        return value
59
60
61# an stl binary file is
62# - 80 bytes of description
63# - 4 bytes of size (unsigned int)
64# - size triangles :
65#
66#   - 12 bytes of normal
67#   - 9 * 4 bytes of coordinate (3*3 floats)
68#   - 2 bytes of garbage (usually 0)
69BINARY_HEADER = 80
70BINARY_STRIDE = 12 * 4 + 2
71
72
73def _header_version():
74    import bpy
75    return "Exported from Blender-" + bpy.app.version_string
76
77
78def _is_ascii_file(data):
79    """
80    This function returns True if the data represents an ASCII file.
81
82    Please note that a False value does not necessary means that the data
83    represents a binary file. It can be a (very *RARE* in real life, but
84    can easily be forged) ascii file.
85    """
86
87    import os
88    import struct
89
90    # Skip header...
91    data.seek(BINARY_HEADER)
92    size = struct.unpack('<I', data.read(4))[0]
93    # Use seek() method to get size of the file.
94    data.seek(0, os.SEEK_END)
95    file_size = data.tell()
96    # Reset to the start of the file.
97    data.seek(0)
98
99    if size == 0:  # Odds to get that result from an ASCII file are null...
100        print("WARNING! Reported size (facet number) is 0, assuming invalid binary STL file.")
101        return False  # Assume binary in this case.
102
103    return (file_size != BINARY_HEADER + 4 + BINARY_STRIDE * size)
104
105
106def _binary_read(data):
107    # Skip header...
108
109    import os
110    import struct
111
112    data.seek(BINARY_HEADER)
113    size = struct.unpack('<I', data.read(4))[0]
114
115    if size == 0:
116        # Workaround invalid crap.
117        data.seek(0, os.SEEK_END)
118        file_size = data.tell()
119        # Reset to after-the-size in the file.
120        data.seek(BINARY_HEADER + 4)
121
122        file_size -= BINARY_HEADER + 4
123        size = file_size // BINARY_STRIDE
124        print("WARNING! Reported size (facet number) is 0, inferring %d facets from file size." % size)
125
126    # We read 4096 elements at once, avoids too much calls to read()!
127    CHUNK_LEN = 4096
128    chunks = [CHUNK_LEN] * (size // CHUNK_LEN)
129    chunks.append(size % CHUNK_LEN)
130
131    unpack = struct.Struct('<12f').unpack_from
132    for chunk_len in chunks:
133        if chunk_len == 0:
134            continue
135        buf = data.read(BINARY_STRIDE * chunk_len)
136        for i in range(chunk_len):
137            # read the normal and points coordinates of each triangle
138            pt = unpack(buf, BINARY_STRIDE * i)
139            yield pt[:3], (pt[3:6], pt[6:9], pt[9:])
140
141
142def _ascii_read(data):
143    # an stl ascii file is like
144    # HEADER: solid some name
145    # for each face:
146    #
147    #     facet normal x y z
148    #     outerloop
149    #          vertex x y z
150    #          vertex x y z
151    #          vertex x y z
152    #     endloop
153    #     endfacet
154
155    # strip header
156    data.readline()
157
158    curr_nor = None
159
160    for l in data:
161        l = l.lstrip()
162        if l.startswith(b'facet'):
163            curr_nor = tuple(map(float, l.split()[2:]))
164        # if we encounter a vertex, read next 2
165        if l.startswith(b'vertex'):
166            yield curr_nor, [tuple(map(float, l_item.split()[1:])) for l_item in (l, data.readline(), data.readline())]
167
168
169def _binary_write(filepath, faces):
170    import struct
171    import itertools
172    from mathutils.geometry import normal
173
174    with open(filepath, 'wb') as data:
175        fw = data.write
176        # header
177        # we write padding at header beginning to avoid to
178        # call len(list(faces)) which may be expensive
179        fw(struct.calcsize('<80sI') * b'\0')
180
181        # 3 vertex == 9f
182        pack = struct.Struct('<9f').pack
183
184        # number of vertices written
185        nb = 0
186
187        for face in faces:
188            # calculate face normal
189            # write normal + vertexes + pad as attributes
190            fw(struct.pack('<3f', *normal(*face)) + pack(*itertools.chain.from_iterable(face)))
191            # attribute byte count (unused)
192            fw(b'\0\0')
193            nb += 1
194
195        # header, with correct value now
196        data.seek(0)
197        fw(struct.pack('<80sI', _header_version().encode('ascii'), nb))
198
199
200def _ascii_write(filepath, faces):
201    from mathutils.geometry import normal
202
203    with open(filepath, 'w') as data:
204        fw = data.write
205        header = _header_version()
206        fw('solid %s\n' % header)
207
208        for face in faces:
209            # calculate face normal
210            fw('facet normal %f %f %f\nouter loop\n' % normal(*face)[:])
211            for vert in face:
212                fw('vertex %f %f %f\n' % vert[:])
213            fw('endloop\nendfacet\n')
214
215        fw('endsolid %s\n' % header)
216
217
218def write_stl(filepath="", faces=(), ascii=False):
219    """
220    Write a stl file from faces,
221
222    filepath
223       output filepath
224
225    faces
226       iterable of tuple of 3 vertex, vertex is tuple of 3 coordinates as float
227
228    ascii
229       save the file in ascii format (very huge)
230    """
231    (_ascii_write if ascii else _binary_write)(filepath, faces)
232
233
234def read_stl(filepath):
235    """
236    Return the triangles and points of an stl binary file.
237
238    Please note that this process can take lot of time if the file is
239    huge (~1m30 for a 1 Go stl file on an quad core i7).
240
241    - returns a tuple(triangles, triangles' normals, points).
242
243      triangles
244          A list of triangles, each triangle as a tuple of 3 index of
245          point in *points*.
246
247      triangles' normals
248          A list of vectors3 (tuples, xyz).
249
250      points
251          An indexed list of points, each point is a tuple of 3 float
252          (xyz).
253
254    Example of use:
255
256       >>> tris, tri_nors, pts = read_stl(filepath)
257       >>> pts = list(pts)
258       >>>
259       >>> # print the coordinate of the triangle n
260       >>> print(pts[i] for i in tris[n])
261    """
262    import time
263    start_time = time.process_time()
264
265    tris, tri_nors, pts = [], [], ListDict()
266
267    with open(filepath, 'rb') as data:
268        # check for ascii or binary
269        gen = _ascii_read if _is_ascii_file(data) else _binary_read
270
271        for nor, pt in gen(data):
272            # Add the triangle and the point.
273            # If the point is already in the list of points, the
274            # index returned by pts.add() will be the one from the
275            # first equal point inserted.
276            tris.append([pts.add(p) for p in pt])
277            tri_nors.append(nor)
278
279    print('Import finished in %.4f sec.' % (time.process_time() - start_time))
280
281    return tris, tri_nors, pts.list
282
283
284if __name__ == '__main__':
285    import sys
286    import bpy
287    from io_mesh_stl import blender_utils
288
289    filepaths = sys.argv[sys.argv.index('--') + 1:]
290
291    for filepath in filepaths:
292        objName = bpy.path.display_name(filepath)
293        tris, pts = read_stl(filepath)
294
295        blender_utils.create_and_link_mesh(objName, tris, pts)
296