Coverage for /Volumes/workspace/numpy-stl/stl/stl.py : 100%

Hot-keys on this page
r m x p toggle line displays
j k next/prev highlighted chunk
0 (zero) top of page
1 (one) first highlighted chunk
1from __future__ import (absolute_import, division, print_function,
2 unicode_literals)
4import io
5import os
6import enum
7import numpy
8import struct
9import datetime
11from . import base
12from . import __about__ as metadata
13from .utils import b
14from .utils import s
16try:
17 from . import _speedups
18except ImportError: # pragma: no cover
19 _speedups = None
22class Mode(enum.IntEnum):
23 #: Automatically detect whether the output is a TTY, if so, write ASCII
24 #: otherwise write BINARY
25 AUTOMATIC = 0
26 #: Force writing ASCII
27 ASCII = 1
28 #: Force writing BINARY
29 BINARY = 2
32# For backwards compatibility, leave the original references
33AUTOMATIC = Mode.AUTOMATIC
34ASCII = Mode.ASCII
35BINARY = Mode.BINARY
38#: Amount of bytes to read while using buffered reading
39BUFFER_SIZE = 4096
40#: The amount of bytes in the header field
41HEADER_SIZE = 80
42#: The amount of bytes in the count field
43COUNT_SIZE = 4
44#: The maximum amount of triangles we can read from binary files
45MAX_COUNT = 1e8
46#: The header format, can be safely monkeypatched. Limited to 80 characters
47HEADER_FORMAT = '{package_name} ({version}) {now} {name}'
50class BaseStl(base.BaseMesh):
52 @classmethod
53 def load(cls, fh, mode=AUTOMATIC, speedups=True):
54 '''Load Mesh from STL file
56 Automatically detects binary versus ascii STL files.
58 :param file fh: The file handle to open
59 :param int mode: Automatically detect the filetype or force binary
60 '''
61 header = fh.read(HEADER_SIZE)
62 if not header:
63 return
65 if isinstance(header, str): # pragma: no branch
66 header = b(header)
68 if mode is AUTOMATIC:
69 if header.lstrip().lower().startswith(b'solid'):
70 try:
71 name, data = cls._load_ascii(
72 fh, header, speedups=speedups)
73 except RuntimeError as exception:
74 (recoverable, e) = exception.args
75 # If we didn't read beyond the header the stream is still
76 # readable through the binary reader
77 if recoverable:
78 name, data = cls._load_binary(fh, header,
79 check_size=False)
80 else:
81 # Apparently we've read beyond the header. Let's try
82 # seeking :)
83 # Note that this fails when reading from stdin, we
84 # can't recover from that.
85 fh.seek(HEADER_SIZE)
87 # Since we know this is a seekable file now and we're
88 # not 100% certain it's binary, check the size while
89 # reading
90 name, data = cls._load_binary(fh, header,
91 check_size=True)
92 else:
93 name, data = cls._load_binary(fh, header)
94 elif mode is ASCII:
95 name, data = cls._load_ascii(fh, header, speedups=speedups)
96 else:
97 name, data = cls._load_binary(fh, header)
99 return name, data
101 @classmethod
102 def _load_binary(cls, fh, header, check_size=False):
103 # Read the triangle count
104 count_data = fh.read(COUNT_SIZE)
105 if len(count_data) != COUNT_SIZE:
106 count = 0
107 else:
108 count, = struct.unpack(s('<i'), b(count_data))
109 # raise RuntimeError()
110 assert count < MAX_COUNT, ('File too large, got %d triangles which '
111 'exceeds the maximum of %d') % (
112 count, MAX_COUNT)
114 if check_size:
115 try:
116 # Check the size of the file
117 fh.seek(0, os.SEEK_END)
118 raw_size = fh.tell() - HEADER_SIZE - COUNT_SIZE
119 expected_count = int(raw_size / cls.dtype.itemsize)
120 assert expected_count == count, ('Expected %d vectors but '
121 'header indicates %d') % (
122 expected_count, count)
123 fh.seek(HEADER_SIZE + COUNT_SIZE)
124 except IOError: # pragma: no cover
125 pass
127 name = header.strip()
129 # Read the rest of the binary data
130 try:
131 return name, numpy.fromfile(fh, dtype=cls.dtype, count=count)
132 except io.UnsupportedOperation:
133 data = numpy.frombuffer(fh.read(), dtype=cls.dtype, count=count)
134 # Copy to make the buffer writable
135 return name, data.copy()
137 @classmethod
138 def _ascii_reader(cls, fh, header):
139 recoverable = [True]
140 line_separator = b'\n'
142 if b'\r\n' in header:
143 line_separator = b'\r\n'
144 elif b'\n' in header:
145 pass
146 elif b'\r' in header:
147 line_separator = b'\r'
148 else:
149 recoverable = [False]
150 header += b(fh.read(BUFFER_SIZE))
152 lines = b(header).split(line_separator)
154 def get(prefix=''):
155 prefix = b(prefix).lower()
157 if lines:
158 raw_line = lines.pop(0)
159 else:
160 raise RuntimeError(recoverable[0], 'Unable to find more lines')
162 if not lines:
163 recoverable[0] = False
165 # Read more lines and make sure we prepend any old data
166 lines[:] = b(fh.read(BUFFER_SIZE)).split(line_separator)
167 raw_line += lines.pop(0)
169 raw_line = raw_line.strip()
170 line = raw_line.lower()
171 if line == b(''):
172 return get(prefix)
174 if prefix:
175 if line.startswith(prefix):
176 values = line.replace(prefix, b(''), 1).strip().split()
177 elif line.startswith(b('endsolid')):
178 # go back to the beginning of new solid part
179 size_unprocessedlines = sum(
180 len(line) + 1 for line in lines) - 1
182 if size_unprocessedlines > 0:
183 position = fh.tell()
184 fh.seek(position - size_unprocessedlines)
185 raise StopIteration()
186 else:
187 raise RuntimeError(
188 recoverable[0],
189 '%r should start with %r' % (line, prefix))
191 if len(values) == 3:
192 return [float(v) for v in values]
193 else: # pragma: no cover
194 raise RuntimeError(recoverable[0],
195 'Incorrect value %r' % line)
196 else:
197 return b(raw_line)
199 line = get()
200 if not lines:
201 raise RuntimeError(recoverable[0],
202 'No lines found, impossible to read')
204 # Yield the name
205 yield line[5:].strip()
207 while True:
208 # Read from the header lines first, until that point we can recover
209 # and go to the binary option. After that we cannot due to
210 # unseekable files such as sys.stdin
211 #
212 # Numpy doesn't support any non-file types so wrapping with a
213 # buffer and/or StringIO does not work.
214 try:
215 normals = get('facet normal')
216 assert get().lower() == b('outer loop')
217 v0 = get('vertex')
218 v1 = get('vertex')
219 v2 = get('vertex')
220 assert get().lower() == b('endloop')
221 assert get().lower() == b('endfacet')
222 attrs = 0
223 yield (normals, (v0, v1, v2), attrs)
224 except AssertionError as e: # pragma: no cover
225 raise RuntimeError(recoverable[0], e)
226 except StopIteration:
227 return
229 @classmethod
230 def _load_ascii(cls, fh, header, speedups=True):
231 # Speedups does not support non file-based streams
232 try:
233 fh.fileno()
234 except io.UnsupportedOperation:
235 speedups = False
236 # The speedups module is covered by travis but it can't be tested in
237 # all environments, this makes coverage checks easier
238 if _speedups and speedups: # pragma: no cover
239 return _speedups.ascii_read(fh, header)
240 else:
241 iterator = cls._ascii_reader(fh, header)
242 name = next(iterator)
243 return name, numpy.fromiter(iterator, dtype=cls.dtype)
245 def save(self, filename, fh=None, mode=AUTOMATIC, update_normals=True):
246 '''Save the STL to a (binary) file
248 If mode is :py:data:`AUTOMATIC` an :py:data:`ASCII` file will be
249 written if the output is a TTY and a :py:data:`BINARY` file otherwise.
251 :param str filename: The file to load
252 :param file fh: The file handle to open
253 :param int mode: The mode to write, default is :py:data:`AUTOMATIC`.
254 :param bool update_normals: Whether to update the normals
255 '''
256 assert filename, 'Filename is required for the STL headers'
257 if update_normals:
258 self.update_normals()
260 if mode is AUTOMATIC:
261 # Try to determine if the file is a TTY.
262 if fh:
263 try:
264 if os.isatty(fh.fileno()): # pragma: no cover
265 write = self._write_ascii
266 else:
267 write = self._write_binary
268 except IOError:
269 # If TTY checking fails then it's an io.BytesIO() (or one
270 # of its siblings from io). Assume binary.
271 write = self._write_binary
272 else:
273 write = self._write_binary
274 elif mode is BINARY:
275 write = self._write_binary
276 elif mode is ASCII:
277 write = self._write_ascii
278 else:
279 raise ValueError('Mode %r is invalid' % mode)
281 if isinstance(fh, io.TextIOBase):
282 # Provide a more helpful error if the user mistakenly
283 # assumes ASCII files should be text files.
284 raise TypeError(
285 "File handles should be in binary mode - even when"
286 " writing an ASCII STL.")
288 name = os.path.split(filename)[-1]
289 try:
290 if fh:
291 write(fh, name)
292 else:
293 with open(filename, 'wb') as fh:
294 write(fh, filename)
295 except IOError: # pragma: no cover
296 pass
298 def _write_ascii(self, fh, name):
299 try:
300 fh.fileno()
301 speedups = self.speedups
302 except io.UnsupportedOperation:
303 speedups = False
305 if _speedups and speedups: # pragma: no cover
306 _speedups.ascii_write(fh, b(name), self.data)
307 else:
308 def p(s, file):
309 file.write(b('%s\n' % s))
311 p('solid %s' % name, file=fh)
313 for row in self.data:
314 vectors = row['vectors']
315 p('facet normal %f %f %f' % tuple(row['normals']), file=fh)
316 p(' outer loop', file=fh)
317 p(' vertex %f %f %f' % tuple(vectors[0]), file=fh)
318 p(' vertex %f %f %f' % tuple(vectors[1]), file=fh)
319 p(' vertex %f %f %f' % tuple(vectors[2]), file=fh)
320 p(' endloop', file=fh)
321 p('endfacet', file=fh)
323 p('endsolid %s' % name, file=fh)
325 def get_header(self, name):
326 # Format the header
327 header = HEADER_FORMAT.format(
328 package_name=metadata.__package_name__,
329 version=metadata.__version__,
330 now=datetime.datetime.now(),
331 name=name,
332 )
334 # Make it exactly 80 characters
335 return header[:80].ljust(80, ' ')
337 def _write_binary(self, fh, name):
338 header = self.get_header(name)
339 packed = struct.pack(s('<i'), self.data.size)
341 if isinstance(fh, io.TextIOWrapper): # pragma: no cover
342 packed = str(packed)
343 else:
344 header = b(header)
345 packed = b(packed)
347 fh.write(header)
348 fh.write(packed)
350 if isinstance(fh, io.BufferedWriter):
351 # Write to a true file.
352 self.data.tofile(fh)
353 else:
354 # Write to a pseudo buffer.
355 fh.write(self.data.data)
357 # In theory this should no longer be possible but I'll leave it here
358 # anyway...
359 if self.data.size: # pragma: no cover
360 assert fh.tell() > 84, (
361 'numpy silently refused to write our file. Note that writing '
362 'to `StringIO` objects is not supported by `numpy`')
364 @classmethod
365 def from_file(cls, filename, calculate_normals=True, fh=None,
366 mode=Mode.AUTOMATIC, speedups=True, **kwargs):
367 '''Load a mesh from a STL file
369 :param str filename: The file to load
370 :param bool calculate_normals: Whether to update the normals
371 :param file fh: The file handle to open
372 :param dict kwargs: The same as for :py:class:`stl.mesh.Mesh`
374 '''
375 if fh:
376 name, data = cls.load(
377 fh, mode=mode, speedups=speedups)
378 else:
379 with open(filename, 'rb') as fh:
380 name, data = cls.load(
381 fh, mode=mode, speedups=speedups)
383 return cls(data, calculate_normals, name=name,
384 speedups=speedups, **kwargs)
386 @classmethod
387 def from_multi_file(cls, filename, calculate_normals=True, fh=None,
388 mode=Mode.AUTOMATIC, speedups=True, **kwargs):
389 '''Load multiple meshes from a STL file
391 Note: mode is hardcoded to ascii since binary stl files do not support
392 the multi format
394 :param str filename: The file to load
395 :param bool calculate_normals: Whether to update the normals
396 :param file fh: The file handle to open
397 :param dict kwargs: The same as for :py:class:`stl.mesh.Mesh`
398 '''
399 if fh:
400 close = False
401 else:
402 fh = open(filename, 'rb')
403 close = True
405 try:
406 raw_data = cls.load(fh, mode=mode, speedups=speedups)
407 while raw_data:
408 name, data = raw_data
409 yield cls(data, calculate_normals, name=name,
410 speedups=speedups, **kwargs)
411 raw_data = cls.load(fh, mode=ASCII,
412 speedups=speedups)
414 finally:
415 if close:
416 fh.close()
418 @classmethod
419 def from_files(cls, filenames, calculate_normals=True, mode=Mode.AUTOMATIC,
420 speedups=True, **kwargs):
421 '''Load multiple meshes from a STL file
423 Note: mode is hardcoded to ascii since binary stl files do not support
424 the multi format
426 :param list(str) filenames: The files to load
427 :param bool calculate_normals: Whether to update the normals
428 :param file fh: The file handle to open
429 :param dict kwargs: The same as for :py:class:`stl.mesh.Mesh`
430 '''
431 meshes = []
432 for filename in filenames:
433 meshes.append(cls.from_file(
434 filename,
435 calculate_normals=calculate_normals,
436 mode=mode,
437 speedups=speedups,
438 **kwargs))
440 data = numpy.concatenate([mesh.data for mesh in meshes])
441 return cls(data, calculate_normals=calculate_normals, **kwargs)
444StlMesh = BaseStl.from_file