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

1import datetime 

2import enum 

3import io 

4import os 

5import struct 

6import zipfile 

7from xml.etree import ElementTree 

8 

9import numpy 

10 

11from . import __about__ as metadata, base 

12from .utils import b 

13 

14try: 

15 from . import _speedups 

16except ImportError: # pragma: no cover 

17 _speedups = None 

18 

19 

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 

28 

29 

30# For backwards compatibility, leave the original references 

31AUTOMATIC = Mode.AUTOMATIC 

32ASCII = Mode.ASCII 

33BINARY = Mode.BINARY 

34 

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}' 

45 

46 

47class BaseStl(base.BaseMesh): 

48 

49 @classmethod 

50 def load(cls, fh, mode=AUTOMATIC, speedups=True): 

51 '''Load Mesh from STL file 

52 

53 Automatically detects binary versus ascii STL files. 

54 

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 

61 

62 if isinstance(header, str): # pragma: no branch 

63 header = b(header) 

64 

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) 

87 

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) 

101 

102 return name, data 

103 

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) 

116 

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 

129 

130 name = header.strip() 

131 

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() 

139 

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)) 

147 

148 lines = b(header).split(b'\n') 

149 

150 def get(prefix=''): 

151 prefix = b(prefix).lower() 

152 

153 if lines: 

154 raw_line = lines.pop(0) 

155 else: 

156 raise RuntimeError(recoverable[0], 'Unable to find more lines') 

157 

158 if not lines: 

159 recoverable[0] = False 

160 

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) 

164 

165 raw_line = raw_line.strip() 

166 line = raw_line.lower() 

167 if line == b(''): 

168 return get(prefix) 

169 

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 

179 

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 ) 

189 

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) 

199 

200 line = get() 

201 if not lines: 

202 raise RuntimeError( 

203 recoverable[0], 

204 'No lines found, impossible to read' 

205 ) 

206 

207 # Yield the name 

208 yield line[5:].strip() 

209 

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 

231 

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) 

247 

248 def save(self, filename, fh=None, mode=AUTOMATIC, update_normals=True): 

249 '''Save the STL to a (binary) file 

250 

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. 

253 

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() 

262 

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) 

283 

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 ) 

291 

292 name = self.name 

293 if not name: 

294 name = os.path.split(filename)[-1] 

295 

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 

304 

305 def _write_ascii(self, fh, name): 

306 try: 

307 fh.fileno() 

308 speedups = self.speedups 

309 except io.UnsupportedOperation: 

310 speedups = False 

311 

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)) 

317 

318 p('solid %s' % name, file=fh) 

319 

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) 

329 

330 p('endsolid %s' % name, file=fh) 

331 

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 ) 

340 

341 # Make it exactly 80 characters 

342 return header[:80].ljust(80, ' ') 

343 

344 def _write_binary(self, fh, name): 

345 header = self.get_header(name) 

346 packed = struct.pack('<i', self.data.size) 

347 

348 if isinstance(fh, io.TextIOWrapper): # pragma: no cover 

349 packed = str(packed) 

350 else: 

351 header = b(header) 

352 packed = b(packed) 

353 

354 fh.write(header) 

355 fh.write(packed) 

356 

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) 

363 

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`') 

370 

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 

377 

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` 

382 

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 ) 

393 

394 return cls( 

395 data, calculate_normals, name=name, 

396 speedups=speedups, **kwargs 

397 ) 

398 

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 

405 

406 Note: mode is hardcoded to ascii since binary stl files do not support 

407 the multi format 

408 

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 

419 

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 ) 

432 

433 finally: 

434 if close: 

435 fh.close() 

436 

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 

443 

444 Note: mode is hardcoded to ascii since binary stl files do not support 

445 the multi format 

446 

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 ) 

463 

464 data = numpy.concatenate([mesh.data for mesh in meshes]) 

465 return cls(data, calculate_normals=calculate_normals, **kwargs) 

466 

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 

478 

479 assert model, 'No 3D model found in %s' % filename 

480 with zip.open(model.lstrip('/')) as fh: 

481 root = ElementTree.parse(fh).getroot() 

482 

483 elements = root.findall('./{*}resources/{*}object/{*}mesh') 

484 for mesh_element in elements: # pragma: no branch 

485 triangles = [] 

486 vertices = [] 

487 

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']]) 

496 

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 ) 

509 

510 mesh = cls(numpy.zeros(len(triangles), dtype=cls.dtype)) 

511 mesh.vectors[:] = numpy.array(triangles) 

512 yield mesh 

513 

514 

515StlMesh = BaseStl.from_file