Hide keyboard shortcuts

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) 

3 

4import io 

5import os 

6import enum 

7import numpy 

8import struct 

9import datetime 

10 

11from . import base 

12from . import __about__ as metadata 

13from .utils import b 

14from .utils import s 

15 

16try: 

17 from . import _speedups 

18except ImportError: # pragma: no cover 

19 _speedups = None 

20 

21 

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 

30 

31 

32# For backwards compatibility, leave the original references 

33AUTOMATIC = Mode.AUTOMATIC 

34ASCII = Mode.ASCII 

35BINARY = Mode.BINARY 

36 

37 

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

48 

49 

50class BaseStl(base.BaseMesh): 

51 

52 @classmethod 

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

54 '''Load Mesh from STL file 

55 

56 Automatically detects binary versus ascii STL files. 

57 

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 

64 

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

66 header = b(header) 

67 

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) 

86 

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) 

98 

99 return name, data 

100 

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) 

113 

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 

126 

127 name = header.strip() 

128 

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

136 

137 @classmethod 

138 def _ascii_reader(cls, fh, header): 

139 recoverable = [True] 

140 line_separator = b'\n' 

141 

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

151 

152 lines = b(header).split(line_separator) 

153 

154 def get(prefix=''): 

155 prefix = b(prefix).lower() 

156 

157 if lines: 

158 raw_line = lines.pop(0) 

159 else: 

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

161 

162 if not lines: 

163 recoverable[0] = False 

164 

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) 

168 

169 raw_line = raw_line.strip() 

170 line = raw_line.lower() 

171 if line == b(''): 

172 return get(prefix) 

173 

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 

181 

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

190 

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) 

198 

199 line = get() 

200 if not lines: 

201 raise RuntimeError(recoverable[0], 

202 'No lines found, impossible to read') 

203 

204 # Yield the name 

205 yield line[5:].strip() 

206 

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 

228 

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) 

244 

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

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

247 

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. 

250 

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

259 

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) 

280 

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

287 

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 

297 

298 def _write_ascii(self, fh, name): 

299 try: 

300 fh.fileno() 

301 speedups = self.speedups 

302 except io.UnsupportedOperation: 

303 speedups = False 

304 

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

310 

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

312 

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) 

322 

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

324 

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 ) 

333 

334 # Make it exactly 80 characters 

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

336 

337 def _write_binary(self, fh, name): 

338 header = self.get_header(name) 

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

340 

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

342 packed = str(packed) 

343 else: 

344 header = b(header) 

345 packed = b(packed) 

346 

347 fh.write(header) 

348 fh.write(packed) 

349 

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) 

356 

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

363 

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 

368 

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` 

373 

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) 

382 

383 return cls(data, calculate_normals, name=name, 

384 speedups=speedups, **kwargs) 

385 

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 

390 

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

392 the multi format 

393 

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 

404 

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) 

413 

414 finally: 

415 if close: 

416 fh.close() 

417 

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 

422 

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

424 the multi format 

425 

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

439 

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

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

442 

443 

444StlMesh = BaseStl.from_file