Coverage for /Volumes/workspace/numpy-stl/stl/stl.py: 100%
245 statements
« prev ^ index » next coverage.py v6.5.0, created at 2022-12-14 23:18 +0100
« prev ^ index » next coverage.py v6.5.0, created at 2022-12-14 23:18 +0100
1import datetime
2import enum
3import io
4import os
5import struct
6import zipfile
7from xml.etree import ElementTree
9import numpy
11from . import __about__ as metadata, base
12from .utils import b
14try:
15 from . import _speedups
16except ImportError: # pragma: no cover
17 _speedups = None
20class Mode(enum.IntEnum):
21 #: Automatically detect whether the output is a TTY, if so, write ASCII
22 #: otherwise write BINARY
23 AUTOMATIC = 0
24 #: Force writing ASCII
25 ASCII = 1
26 #: Force writing BINARY
27 BINARY = 2
30# For backwards compatibility, leave the original references
31AUTOMATIC = Mode.AUTOMATIC
32ASCII = Mode.ASCII
33BINARY = Mode.BINARY
35#: Amount of bytes to read while using buffered reading
36BUFFER_SIZE = 4096
37#: The amount of bytes in the header field
38HEADER_SIZE = 80
39#: The amount of bytes in the count field
40COUNT_SIZE = 4
41#: The maximum amount of triangles we can read from binary files
42MAX_COUNT = 1e8
43#: The header format, can be safely monkeypatched. Limited to 80 characters
44HEADER_FORMAT = '{package_name} ({version}) {now} {name}'
47class BaseStl(base.BaseMesh):
49 @classmethod
50 def load(cls, fh, mode=AUTOMATIC, speedups=True):
51 '''Load Mesh from STL file
53 Automatically detects binary versus ascii STL files.
55 :param file fh: The file handle to open
56 :param int mode: Automatically detect the filetype or force binary
57 '''
58 header = fh.read(HEADER_SIZE)
59 if not header:
60 return
62 if isinstance(header, str): # pragma: no branch
63 header = b(header)
65 if mode is AUTOMATIC:
66 if header.lstrip().lower().startswith(b'solid'):
67 try:
68 name, data = cls._load_ascii(
69 fh, header, speedups=speedups
70 )
71 except RuntimeError as exception:
72 print('exception', exception)
73 (recoverable, e) = exception.args
74 # If we didn't read beyond the header the stream is still
75 # readable through the binary reader
76 if recoverable:
77 name, data = cls._load_binary(
78 fh, header,
79 check_size=False
80 )
81 else:
82 # Apparently we've read beyond the header. Let's try
83 # seeking :)
84 # Note that this fails when reading from stdin, we
85 # can't recover from that.
86 fh.seek(HEADER_SIZE)
88 # Since we know this is a seekable file now and we're
89 # not 100% certain it's binary, check the size while
90 # reading
91 name, data = cls._load_binary(
92 fh, header,
93 check_size=True
94 )
95 else:
96 name, data = cls._load_binary(fh, header)
97 elif mode is ASCII:
98 name, data = cls._load_ascii(fh, header, speedups=speedups)
99 else:
100 name, data = cls._load_binary(fh, header)
102 return name, data
104 @classmethod
105 def _load_binary(cls, fh, header, check_size=False):
106 # Read the triangle count
107 count_data = fh.read(COUNT_SIZE)
108 if len(count_data) != COUNT_SIZE:
109 count = 0
110 else:
111 count, = struct.unpack('<i', b(count_data))
112 # raise RuntimeError()
113 assert count < MAX_COUNT, ('File too large, got %d triangles which '
114 'exceeds the maximum of %d') % (
115 count, MAX_COUNT)
117 if check_size:
118 try:
119 # Check the size of the file
120 fh.seek(0, os.SEEK_END)
121 raw_size = fh.tell() - HEADER_SIZE - COUNT_SIZE
122 expected_count = int(raw_size / cls.dtype.itemsize)
123 assert expected_count == count, ('Expected %d vectors but '
124 'header indicates %d') % (
125 expected_count, count)
126 fh.seek(HEADER_SIZE + COUNT_SIZE)
127 except IOError: # pragma: no cover
128 pass
130 name = header.strip()
132 # Read the rest of the binary data
133 try:
134 return name, numpy.fromfile(fh, dtype=cls.dtype, count=count)
135 except io.UnsupportedOperation:
136 data = numpy.frombuffer(fh.read(), dtype=cls.dtype, count=count)
137 # Copy to make the buffer writable
138 return name, data.copy()
140 @classmethod
141 def _ascii_reader(cls, fh, header):
142 if b'\n' in header:
143 recoverable = [True]
144 else:
145 recoverable = [False]
146 header += b(fh.read(BUFFER_SIZE))
148 lines = b(header).split(b'\n')
150 def get(prefix=''):
151 prefix = b(prefix).lower()
153 if lines:
154 raw_line = lines.pop(0)
155 else:
156 raise RuntimeError(recoverable[0], 'Unable to find more lines')
158 if not lines:
159 recoverable[0] = False
161 # Read more lines and make sure we prepend any old data
162 lines[:] = b(fh.read(BUFFER_SIZE)).split(b'\n')
163 raw_line += lines.pop(0)
165 raw_line = raw_line.strip()
166 line = raw_line.lower()
167 if line == b(''):
168 return get(prefix)
170 if prefix:
171 if line.startswith(prefix):
172 values = line.replace(prefix, b(''), 1).strip().split()
173 elif line.startswith(b('endsolid')) \
174 or line.startswith(b('end solid')):
175 # go back to the beginning of new solid part
176 size_unprocessedlines = sum(
177 len(line) + 1 for line in lines
178 ) - 1
180 if size_unprocessedlines > 0:
181 position = fh.tell()
182 fh.seek(position - size_unprocessedlines)
183 raise StopIteration()
184 else:
185 raise RuntimeError(
186 recoverable[0],
187 '%r should start with %r' % (line, prefix)
188 )
190 if len(values) == 3:
191 return [float(v) for v in values]
192 else: # pragma: no cover
193 raise RuntimeError(
194 recoverable[0],
195 'Incorrect value %r' % line
196 )
197 else:
198 return b(raw_line)
200 line = get()
201 if not lines:
202 raise RuntimeError(
203 recoverable[0],
204 'No lines found, impossible to read'
205 )
207 # Yield the name
208 yield line[5:].strip()
210 while True:
211 # Read from the header lines first, until that point we can recover
212 # and go to the binary option. After that we cannot due to
213 # unseekable files such as sys.stdin
214 #
215 # Numpy doesn't support any non-file types so wrapping with a
216 # buffer and/or StringIO does not work.
217 try:
218 normals = get('facet normal')
219 assert get().lower() == b('outer loop')
220 v0 = get('vertex')
221 v1 = get('vertex')
222 v2 = get('vertex')
223 assert get().lower() == b('endloop')
224 assert get().lower() == b('endfacet')
225 attrs = 0
226 yield (normals, (v0, v1, v2), attrs)
227 except AssertionError as e: # pragma: no cover
228 raise RuntimeError(recoverable[0], e)
229 except StopIteration:
230 return
232 @classmethod
233 def _load_ascii(cls, fh, header, speedups=True):
234 # Speedups does not support non file-based streams
235 try:
236 fh.fileno()
237 except io.UnsupportedOperation:
238 speedups = False
239 # The speedups module is covered by travis but it can't be tested in
240 # all environments, this makes coverage checks easier
241 if _speedups and speedups: # pragma: no cover
242 return _speedups.ascii_read(fh, header)
243 else:
244 iterator = cls._ascii_reader(fh, header)
245 name = next(iterator)
246 return name, numpy.fromiter(iterator, dtype=cls.dtype)
248 def save(self, filename, fh=None, mode=AUTOMATIC, update_normals=True):
249 '''Save the STL to a (binary) file
251 If mode is :py:data:`AUTOMATIC` an :py:data:`ASCII` file will be
252 written if the output is a TTY and a :py:data:`BINARY` file otherwise.
254 :param str filename: The file to load
255 :param file fh: The file handle to open
256 :param int mode: The mode to write, default is :py:data:`AUTOMATIC`.
257 :param bool update_normals: Whether to update the normals
258 '''
259 assert filename, 'Filename is required for the STL headers'
260 if update_normals:
261 self.update_normals()
263 if mode is AUTOMATIC:
264 # Try to determine if the file is a TTY.
265 if fh:
266 try:
267 if os.isatty(fh.fileno()): # pragma: no cover
268 write = self._write_ascii
269 else:
270 write = self._write_binary
271 except IOError:
272 # If TTY checking fails then it's an io.BytesIO() (or one
273 # of its siblings from io). Assume binary.
274 write = self._write_binary
275 else:
276 write = self._write_binary
277 elif mode is BINARY:
278 write = self._write_binary
279 elif mode is ASCII:
280 write = self._write_ascii
281 else:
282 raise ValueError('Mode %r is invalid' % mode)
284 if isinstance(fh, io.TextIOBase):
285 # Provide a more helpful error if the user mistakenly
286 # assumes ASCII files should be text files.
287 raise TypeError(
288 "File handles should be in binary mode - even when"
289 " writing an ASCII STL."
290 )
292 name = self.name
293 if not name:
294 name = os.path.split(filename)[-1]
296 try:
297 if fh:
298 write(fh, name)
299 else:
300 with open(filename, 'wb') as fh:
301 write(fh, name)
302 except IOError: # pragma: no cover
303 pass
305 def _write_ascii(self, fh, name):
306 try:
307 fh.fileno()
308 speedups = self.speedups
309 except io.UnsupportedOperation:
310 speedups = False
312 if _speedups and speedups: # pragma: no cover
313 _speedups.ascii_write(fh, b(name), self.data)
314 else:
315 def p(s, file):
316 file.write(b('%s\n' % s))
318 p('solid %s' % name, file=fh)
320 for row in self.data:
321 vectors = row['vectors']
322 p('facet normal %r %r %r' % tuple(row['normals']), file=fh)
323 p(' outer loop', file=fh)
324 p(' vertex %r %r %r' % tuple(vectors[0]), file=fh)
325 p(' vertex %r %r %r' % tuple(vectors[1]), file=fh)
326 p(' vertex %r %r %r' % tuple(vectors[2]), file=fh)
327 p(' endloop', file=fh)
328 p('endfacet', file=fh)
330 p('endsolid %s' % name, file=fh)
332 def get_header(self, name):
333 # Format the header
334 header = HEADER_FORMAT.format(
335 package_name=metadata.__package_name__,
336 version=metadata.__version__,
337 now=datetime.datetime.now(),
338 name=name,
339 )
341 # Make it exactly 80 characters
342 return header[:80].ljust(80, ' ')
344 def _write_binary(self, fh, name):
345 header = self.get_header(name)
346 packed = struct.pack('<i', self.data.size)
348 if isinstance(fh, io.TextIOWrapper): # pragma: no cover
349 packed = str(packed)
350 else:
351 header = b(header)
352 packed = b(packed)
354 fh.write(header)
355 fh.write(packed)
357 if isinstance(fh, io.BufferedWriter):
358 # Write to a true file.
359 self.data.tofile(fh)
360 else:
361 # Write to a pseudo buffer.
362 fh.write(self.data.data)
364 # In theory this should no longer be possible but I'll leave it here
365 # anyway...
366 if self.data.size: # pragma: no cover
367 assert fh.tell() > 84, (
368 'numpy silently refused to write our file. Note that writing '
369 'to `StringIO` objects is not supported by `numpy`')
371 @classmethod
372 def from_file(
373 cls, filename, calculate_normals=True, fh=None,
374 mode=Mode.AUTOMATIC, speedups=True, **kwargs
375 ):
376 '''Load a mesh from a STL file
378 :param str filename: The file to load
379 :param bool calculate_normals: Whether to update the normals
380 :param file fh: The file handle to open
381 :param dict kwargs: The same as for :py:class:`stl.mesh.Mesh`
383 '''
384 if fh:
385 name, data = cls.load(
386 fh, mode=mode, speedups=speedups
387 )
388 else:
389 with open(filename, 'rb') as fh:
390 name, data = cls.load(
391 fh, mode=mode, speedups=speedups
392 )
394 return cls(
395 data, calculate_normals, name=name,
396 speedups=speedups, **kwargs
397 )
399 @classmethod
400 def from_multi_file(
401 cls, filename, calculate_normals=True, fh=None,
402 mode=Mode.AUTOMATIC, speedups=True, **kwargs
403 ):
404 '''Load multiple meshes from a STL file
406 Note: mode is hardcoded to ascii since binary stl files do not support
407 the multi format
409 :param str filename: The file to load
410 :param bool calculate_normals: Whether to update the normals
411 :param file fh: The file handle to open
412 :param dict kwargs: The same as for :py:class:`stl.mesh.Mesh`
413 '''
414 if fh:
415 close = False
416 else:
417 fh = open(filename, 'rb')
418 close = True
420 try:
421 raw_data = cls.load(fh, mode=mode, speedups=speedups)
422 while raw_data:
423 name, data = raw_data
424 yield cls(
425 data, calculate_normals, name=name,
426 speedups=speedups, **kwargs
427 )
428 raw_data = cls.load(
429 fh, mode=ASCII,
430 speedups=speedups
431 )
433 finally:
434 if close:
435 fh.close()
437 @classmethod
438 def from_files(
439 cls, filenames, calculate_normals=True, mode=Mode.AUTOMATIC,
440 speedups=True, **kwargs
441 ):
442 '''Load multiple meshes from STL files into a single mesh
444 Note: mode is hardcoded to ascii since binary stl files do not support
445 the multi format
447 :param list(str) filenames: The files to load
448 :param bool calculate_normals: Whether to update the normals
449 :param file fh: The file handle to open
450 :param dict kwargs: The same as for :py:class:`stl.mesh.Mesh`
451 '''
452 meshes = []
453 for filename in filenames:
454 meshes.append(
455 cls.from_file(
456 filename,
457 calculate_normals=calculate_normals,
458 mode=mode,
459 speedups=speedups,
460 **kwargs
461 )
462 )
464 data = numpy.concatenate([mesh.data for mesh in meshes])
465 return cls(data, calculate_normals=calculate_normals, **kwargs)
467 @classmethod
468 def from_3mf_file(cls, filename, calculate_normals=True, **kwargs):
469 with zipfile.ZipFile(filename) as zip:
470 with zip.open('_rels/.rels') as rels_fh:
471 model = None
472 root = ElementTree.parse(rels_fh).getroot()
473 for child in root: # pragma: no branch
474 type_ = child.attrib.get('Type', '')
475 if type_.endswith('3dmodel'): # pragma: no branch
476 model = child.attrib.get('Target', '')
477 break
479 assert model, 'No 3D model found in %s' % filename
480 with zip.open(model.lstrip('/')) as fh:
481 root = ElementTree.parse(fh).getroot()
483 elements = root.findall('./{*}resources/{*}object/{*}mesh')
484 for mesh_element in elements: # pragma: no branch
485 triangles = []
486 vertices = []
488 for element in mesh_element:
489 tag = element.tag
490 if tag.endswith('vertices'):
491 # Collect all the vertices
492 for vertice in element:
493 a = {k: float(v) for k, v in
494 vertice.attrib.items()}
495 vertices.append([a['x'], a['y'], a['z']])
497 elif tag.endswith('triangles'): # pragma: no branch
498 # Map the triangles to the vertices and collect
499 for triangle in element:
500 a = {k: int(v) for k, v in
501 triangle.attrib.items()}
502 triangles.append(
503 [
504 vertices[a['v1']],
505 vertices[a['v2']],
506 vertices[a['v3']],
507 ]
508 )
510 mesh = cls(numpy.zeros(len(triangles), dtype=cls.dtype))
511 mesh.vectors[:] = numpy.array(triangles)
512 yield mesh
515StlMesh = BaseStl.from_file