ImageFile.py 22 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609610611612613614615616617618619620621622623624625626627628629630631632633634635636637638639640641642643644645646647648649650651652653654655656657658659660661662663664665666667668669670671672673674675676677678679680681682683684685686687688689690691692693694695696697698699700701702703704705706707708709710711712713714715716717718719720721722723724725726727728729730731732733734735736737738739740741742743744745746747748
  1. #
  2. # The Python Imaging Library.
  3. # $Id$
  4. #
  5. # base class for image file handlers
  6. #
  7. # history:
  8. # 1995-09-09 fl Created
  9. # 1996-03-11 fl Fixed load mechanism.
  10. # 1996-04-15 fl Added pcx/xbm decoders.
  11. # 1996-04-30 fl Added encoders.
  12. # 1996-12-14 fl Added load helpers
  13. # 1997-01-11 fl Use encode_to_file where possible
  14. # 1997-08-27 fl Flush output in _save
  15. # 1998-03-05 fl Use memory mapping for some modes
  16. # 1999-02-04 fl Use memory mapping also for "I;16" and "I;16B"
  17. # 1999-05-31 fl Added image parser
  18. # 2000-10-12 fl Set readonly flag on memory-mapped images
  19. # 2002-03-20 fl Use better messages for common decoder errors
  20. # 2003-04-21 fl Fall back on mmap/map_buffer if map is not available
  21. # 2003-10-30 fl Added StubImageFile class
  22. # 2004-02-25 fl Made incremental parser more robust
  23. #
  24. # Copyright (c) 1997-2004 by Secret Labs AB
  25. # Copyright (c) 1995-2004 by Fredrik Lundh
  26. #
  27. # See the README file for information on usage and redistribution.
  28. #
  29. import io
  30. import itertools
  31. import struct
  32. import sys
  33. from . import Image
  34. from ._util import is_path
  35. MAXBLOCK = 65536
  36. SAFEBLOCK = 1024 * 1024
  37. LOAD_TRUNCATED_IMAGES = False
  38. """Whether or not to load truncated image files. User code may change this."""
  39. ERRORS = {
  40. -1: "image buffer overrun error",
  41. -2: "decoding error",
  42. -3: "unknown error",
  43. -8: "bad configuration",
  44. -9: "out of memory error",
  45. }
  46. """
  47. Dict of known error codes returned from :meth:`.PyDecoder.decode`,
  48. :meth:`.PyEncoder.encode` :meth:`.PyEncoder.encode_to_pyfd` and
  49. :meth:`.PyEncoder.encode_to_file`.
  50. """
  51. #
  52. # --------------------------------------------------------------------
  53. # Helpers
  54. def raise_oserror(error):
  55. try:
  56. message = Image.core.getcodecstatus(error)
  57. except AttributeError:
  58. message = ERRORS.get(error)
  59. if not message:
  60. message = f"decoder error {error}"
  61. raise OSError(message + " when reading image file")
  62. def _tilesort(t):
  63. # sort on offset
  64. return t[2]
  65. #
  66. # --------------------------------------------------------------------
  67. # ImageFile base class
  68. class ImageFile(Image.Image):
  69. """Base class for image file format handlers."""
  70. def __init__(self, fp=None, filename=None):
  71. super().__init__()
  72. self._min_frame = 0
  73. self.custom_mimetype = None
  74. self.tile = None
  75. """ A list of tile descriptors, or ``None`` """
  76. self.readonly = 1 # until we know better
  77. self.decoderconfig = ()
  78. self.decodermaxblock = MAXBLOCK
  79. if is_path(fp):
  80. # filename
  81. self.fp = open(fp, "rb")
  82. self.filename = fp
  83. self._exclusive_fp = True
  84. else:
  85. # stream
  86. self.fp = fp
  87. self.filename = filename
  88. # can be overridden
  89. self._exclusive_fp = None
  90. try:
  91. try:
  92. self._open()
  93. except (
  94. IndexError, # end of data
  95. TypeError, # end of data (ord)
  96. KeyError, # unsupported mode
  97. EOFError, # got header but not the first frame
  98. struct.error,
  99. ) as v:
  100. raise SyntaxError(v) from v
  101. if not self.mode or self.size[0] <= 0 or self.size[1] <= 0:
  102. raise SyntaxError("not identified by this driver")
  103. except BaseException:
  104. # close the file only if we have opened it this constructor
  105. if self._exclusive_fp:
  106. self.fp.close()
  107. raise
  108. def get_format_mimetype(self):
  109. if self.custom_mimetype:
  110. return self.custom_mimetype
  111. if self.format is not None:
  112. return Image.MIME.get(self.format.upper())
  113. def verify(self):
  114. """Check file integrity"""
  115. # raise exception if something's wrong. must be called
  116. # directly after open, and closes file when finished.
  117. if self._exclusive_fp:
  118. self.fp.close()
  119. self.fp = None
  120. def load(self):
  121. """Load image data based on tile list"""
  122. if self.tile is None:
  123. raise OSError("cannot load this image")
  124. pixel = Image.Image.load(self)
  125. if not self.tile:
  126. return pixel
  127. self.map = None
  128. use_mmap = self.filename and len(self.tile) == 1
  129. # As of pypy 2.1.0, memory mapping was failing here.
  130. use_mmap = use_mmap and not hasattr(sys, "pypy_version_info")
  131. readonly = 0
  132. # look for read/seek overrides
  133. try:
  134. read = self.load_read
  135. # don't use mmap if there are custom read/seek functions
  136. use_mmap = False
  137. except AttributeError:
  138. read = self.fp.read
  139. try:
  140. seek = self.load_seek
  141. use_mmap = False
  142. except AttributeError:
  143. seek = self.fp.seek
  144. if use_mmap:
  145. # try memory mapping
  146. decoder_name, extents, offset, args = self.tile[0]
  147. if (
  148. decoder_name == "raw"
  149. and len(args) >= 3
  150. and args[0] == self.mode
  151. and args[0] in Image._MAPMODES
  152. ):
  153. try:
  154. # use mmap, if possible
  155. import mmap
  156. with open(self.filename) as fp:
  157. self.map = mmap.mmap(fp.fileno(), 0, access=mmap.ACCESS_READ)
  158. self.im = Image.core.map_buffer(
  159. self.map, self.size, decoder_name, offset, args
  160. )
  161. readonly = 1
  162. # After trashing self.im,
  163. # we might need to reload the palette data.
  164. if self.palette:
  165. self.palette.dirty = 1
  166. except (AttributeError, OSError, ImportError):
  167. self.map = None
  168. self.load_prepare()
  169. err_code = -3 # initialize to unknown error
  170. if not self.map:
  171. # sort tiles in file order
  172. self.tile.sort(key=_tilesort)
  173. try:
  174. # FIXME: This is a hack to handle TIFF's JpegTables tag.
  175. prefix = self.tile_prefix
  176. except AttributeError:
  177. prefix = b""
  178. # Remove consecutive duplicates that only differ by their offset
  179. self.tile = [
  180. list(tiles)[-1]
  181. for _, tiles in itertools.groupby(
  182. self.tile, lambda tile: (tile[0], tile[1], tile[3])
  183. )
  184. ]
  185. for decoder_name, extents, offset, args in self.tile:
  186. seek(offset)
  187. decoder = Image._getdecoder(
  188. self.mode, decoder_name, args, self.decoderconfig
  189. )
  190. try:
  191. decoder.setimage(self.im, extents)
  192. if decoder.pulls_fd:
  193. decoder.setfd(self.fp)
  194. err_code = decoder.decode(b"")[1]
  195. else:
  196. b = prefix
  197. while True:
  198. try:
  199. s = read(self.decodermaxblock)
  200. except (IndexError, struct.error) as e:
  201. # truncated png/gif
  202. if LOAD_TRUNCATED_IMAGES:
  203. break
  204. else:
  205. raise OSError("image file is truncated") from e
  206. if not s: # truncated jpeg
  207. if LOAD_TRUNCATED_IMAGES:
  208. break
  209. else:
  210. raise OSError(
  211. "image file is truncated "
  212. f"({len(b)} bytes not processed)"
  213. )
  214. b = b + s
  215. n, err_code = decoder.decode(b)
  216. if n < 0:
  217. break
  218. b = b[n:]
  219. finally:
  220. # Need to cleanup here to prevent leaks
  221. decoder.cleanup()
  222. self.tile = []
  223. self.readonly = readonly
  224. self.load_end()
  225. if self._exclusive_fp and self._close_exclusive_fp_after_loading:
  226. self.fp.close()
  227. self.fp = None
  228. if not self.map and not LOAD_TRUNCATED_IMAGES and err_code < 0:
  229. # still raised if decoder fails to return anything
  230. raise_oserror(err_code)
  231. return Image.Image.load(self)
  232. def load_prepare(self):
  233. # create image memory if necessary
  234. if not self.im or self.im.mode != self.mode or self.im.size != self.size:
  235. self.im = Image.core.new(self.mode, self.size)
  236. # create palette (optional)
  237. if self.mode == "P":
  238. Image.Image.load(self)
  239. def load_end(self):
  240. # may be overridden
  241. pass
  242. # may be defined for contained formats
  243. # def load_seek(self, pos):
  244. # pass
  245. # may be defined for blocked formats (e.g. PNG)
  246. # def load_read(self, bytes):
  247. # pass
  248. def _seek_check(self, frame):
  249. if (
  250. frame < self._min_frame
  251. # Only check upper limit on frames if additional seek operations
  252. # are not required to do so
  253. or (
  254. not (hasattr(self, "_n_frames") and self._n_frames is None)
  255. and frame >= self.n_frames + self._min_frame
  256. )
  257. ):
  258. raise EOFError("attempt to seek outside sequence")
  259. return self.tell() != frame
  260. class StubImageFile(ImageFile):
  261. """
  262. Base class for stub image loaders.
  263. A stub loader is an image loader that can identify files of a
  264. certain format, but relies on external code to load the file.
  265. """
  266. def _open(self):
  267. raise NotImplementedError("StubImageFile subclass must implement _open")
  268. def load(self):
  269. loader = self._load()
  270. if loader is None:
  271. raise OSError(f"cannot find loader for this {self.format} file")
  272. image = loader.load(self)
  273. assert image is not None
  274. # become the other object (!)
  275. self.__class__ = image.__class__
  276. self.__dict__ = image.__dict__
  277. return image.load()
  278. def _load(self):
  279. """(Hook) Find actual image loader."""
  280. raise NotImplementedError("StubImageFile subclass must implement _load")
  281. class Parser:
  282. """
  283. Incremental image parser. This class implements the standard
  284. feed/close consumer interface.
  285. """
  286. incremental = None
  287. image = None
  288. data = None
  289. decoder = None
  290. offset = 0
  291. finished = 0
  292. def reset(self):
  293. """
  294. (Consumer) Reset the parser. Note that you can only call this
  295. method immediately after you've created a parser; parser
  296. instances cannot be reused.
  297. """
  298. assert self.data is None, "cannot reuse parsers"
  299. def feed(self, data):
  300. """
  301. (Consumer) Feed data to the parser.
  302. :param data: A string buffer.
  303. :exception OSError: If the parser failed to parse the image file.
  304. """
  305. # collect data
  306. if self.finished:
  307. return
  308. if self.data is None:
  309. self.data = data
  310. else:
  311. self.data = self.data + data
  312. # parse what we have
  313. if self.decoder:
  314. if self.offset > 0:
  315. # skip header
  316. skip = min(len(self.data), self.offset)
  317. self.data = self.data[skip:]
  318. self.offset = self.offset - skip
  319. if self.offset > 0 or not self.data:
  320. return
  321. n, e = self.decoder.decode(self.data)
  322. if n < 0:
  323. # end of stream
  324. self.data = None
  325. self.finished = 1
  326. if e < 0:
  327. # decoding error
  328. self.image = None
  329. raise_oserror(e)
  330. else:
  331. # end of image
  332. return
  333. self.data = self.data[n:]
  334. elif self.image:
  335. # if we end up here with no decoder, this file cannot
  336. # be incrementally parsed. wait until we've gotten all
  337. # available data
  338. pass
  339. else:
  340. # attempt to open this file
  341. try:
  342. with io.BytesIO(self.data) as fp:
  343. im = Image.open(fp)
  344. except OSError:
  345. # traceback.print_exc()
  346. pass # not enough data
  347. else:
  348. flag = hasattr(im, "load_seek") or hasattr(im, "load_read")
  349. if flag or len(im.tile) != 1:
  350. # custom load code, or multiple tiles
  351. self.decode = None
  352. else:
  353. # initialize decoder
  354. im.load_prepare()
  355. d, e, o, a = im.tile[0]
  356. im.tile = []
  357. self.decoder = Image._getdecoder(im.mode, d, a, im.decoderconfig)
  358. self.decoder.setimage(im.im, e)
  359. # calculate decoder offset
  360. self.offset = o
  361. if self.offset <= len(self.data):
  362. self.data = self.data[self.offset :]
  363. self.offset = 0
  364. self.image = im
  365. def __enter__(self):
  366. return self
  367. def __exit__(self, *args):
  368. self.close()
  369. def close(self):
  370. """
  371. (Consumer) Close the stream.
  372. :returns: An image object.
  373. :exception OSError: If the parser failed to parse the image file either
  374. because it cannot be identified or cannot be
  375. decoded.
  376. """
  377. # finish decoding
  378. if self.decoder:
  379. # get rid of what's left in the buffers
  380. self.feed(b"")
  381. self.data = self.decoder = None
  382. if not self.finished:
  383. raise OSError("image was incomplete")
  384. if not self.image:
  385. raise OSError("cannot parse this image")
  386. if self.data:
  387. # incremental parsing not possible; reopen the file
  388. # not that we have all data
  389. with io.BytesIO(self.data) as fp:
  390. try:
  391. self.image = Image.open(fp)
  392. finally:
  393. self.image.load()
  394. return self.image
  395. # --------------------------------------------------------------------
  396. def _save(im, fp, tile, bufsize=0):
  397. """Helper to save image based on tile list
  398. :param im: Image object.
  399. :param fp: File object.
  400. :param tile: Tile list.
  401. :param bufsize: Optional buffer size
  402. """
  403. im.load()
  404. if not hasattr(im, "encoderconfig"):
  405. im.encoderconfig = ()
  406. tile.sort(key=_tilesort)
  407. # FIXME: make MAXBLOCK a configuration parameter
  408. # It would be great if we could have the encoder specify what it needs
  409. # But, it would need at least the image size in most cases. RawEncode is
  410. # a tricky case.
  411. bufsize = max(MAXBLOCK, bufsize, im.size[0] * 4) # see RawEncode.c
  412. try:
  413. fh = fp.fileno()
  414. fp.flush()
  415. exc = None
  416. except (AttributeError, io.UnsupportedOperation) as e:
  417. exc = e
  418. for e, b, o, a in tile:
  419. if o > 0:
  420. fp.seek(o)
  421. encoder = Image._getencoder(im.mode, e, a, im.encoderconfig)
  422. try:
  423. encoder.setimage(im.im, b)
  424. if encoder.pushes_fd:
  425. encoder.setfd(fp)
  426. l, s = encoder.encode_to_pyfd()
  427. else:
  428. if exc:
  429. # compress to Python file-compatible object
  430. while True:
  431. l, s, d = encoder.encode(bufsize)
  432. fp.write(d)
  433. if s:
  434. break
  435. else:
  436. # slight speedup: compress to real file object
  437. s = encoder.encode_to_file(fh, bufsize)
  438. if s < 0:
  439. raise OSError(f"encoder error {s} when writing image file") from exc
  440. finally:
  441. encoder.cleanup()
  442. if hasattr(fp, "flush"):
  443. fp.flush()
  444. def _safe_read(fp, size):
  445. """
  446. Reads large blocks in a safe way. Unlike fp.read(n), this function
  447. doesn't trust the user. If the requested size is larger than
  448. SAFEBLOCK, the file is read block by block.
  449. :param fp: File handle. Must implement a <b>read</b> method.
  450. :param size: Number of bytes to read.
  451. :returns: A string containing <i>size</i> bytes of data.
  452. Raises an OSError if the file is truncated and the read cannot be completed
  453. """
  454. if size <= 0:
  455. return b""
  456. if size <= SAFEBLOCK:
  457. data = fp.read(size)
  458. if len(data) < size:
  459. raise OSError("Truncated File Read")
  460. return data
  461. data = []
  462. remaining_size = size
  463. while remaining_size > 0:
  464. block = fp.read(min(remaining_size, SAFEBLOCK))
  465. if not block:
  466. break
  467. data.append(block)
  468. remaining_size -= len(block)
  469. if sum(len(d) for d in data) < size:
  470. raise OSError("Truncated File Read")
  471. return b"".join(data)
  472. class PyCodecState:
  473. def __init__(self):
  474. self.xsize = 0
  475. self.ysize = 0
  476. self.xoff = 0
  477. self.yoff = 0
  478. def extents(self):
  479. return self.xoff, self.yoff, self.xoff + self.xsize, self.yoff + self.ysize
  480. class PyCodec:
  481. def __init__(self, mode, *args):
  482. self.im = None
  483. self.state = PyCodecState()
  484. self.fd = None
  485. self.mode = mode
  486. self.init(args)
  487. def init(self, args):
  488. """
  489. Override to perform codec specific initialization
  490. :param args: Array of args items from the tile entry
  491. :returns: None
  492. """
  493. self.args = args
  494. def cleanup(self):
  495. """
  496. Override to perform codec specific cleanup
  497. :returns: None
  498. """
  499. pass
  500. def setfd(self, fd):
  501. """
  502. Called from ImageFile to set the Python file-like object
  503. :param fd: A Python file-like object
  504. :returns: None
  505. """
  506. self.fd = fd
  507. def setimage(self, im, extents=None):
  508. """
  509. Called from ImageFile to set the core output image for the codec
  510. :param im: A core image object
  511. :param extents: a 4 tuple of (x0, y0, x1, y1) defining the rectangle
  512. for this tile
  513. :returns: None
  514. """
  515. # following c code
  516. self.im = im
  517. if extents:
  518. (x0, y0, x1, y1) = extents
  519. else:
  520. (x0, y0, x1, y1) = (0, 0, 0, 0)
  521. if x0 == 0 and x1 == 0:
  522. self.state.xsize, self.state.ysize = self.im.size
  523. else:
  524. self.state.xoff = x0
  525. self.state.yoff = y0
  526. self.state.xsize = x1 - x0
  527. self.state.ysize = y1 - y0
  528. if self.state.xsize <= 0 or self.state.ysize <= 0:
  529. raise ValueError("Size cannot be negative")
  530. if (
  531. self.state.xsize + self.state.xoff > self.im.size[0]
  532. or self.state.ysize + self.state.yoff > self.im.size[1]
  533. ):
  534. raise ValueError("Tile cannot extend outside image")
  535. class PyDecoder(PyCodec):
  536. """
  537. Python implementation of a format decoder. Override this class and
  538. add the decoding logic in the :meth:`decode` method.
  539. See :ref:`Writing Your Own File Codec in Python<file-codecs-py>`
  540. """
  541. _pulls_fd = False
  542. @property
  543. def pulls_fd(self):
  544. return self._pulls_fd
  545. def decode(self, buffer):
  546. """
  547. Override to perform the decoding process.
  548. :param buffer: A bytes object with the data to be decoded.
  549. :returns: A tuple of ``(bytes consumed, errcode)``.
  550. If finished with decoding return -1 for the bytes consumed.
  551. Err codes are from :data:`.ImageFile.ERRORS`.
  552. """
  553. raise NotImplementedError()
  554. def set_as_raw(self, data, rawmode=None):
  555. """
  556. Convenience method to set the internal image from a stream of raw data
  557. :param data: Bytes to be set
  558. :param rawmode: The rawmode to be used for the decoder.
  559. If not specified, it will default to the mode of the image
  560. :returns: None
  561. """
  562. if not rawmode:
  563. rawmode = self.mode
  564. d = Image._getdecoder(self.mode, "raw", rawmode)
  565. d.setimage(self.im, self.state.extents())
  566. s = d.decode(data)
  567. if s[0] >= 0:
  568. raise ValueError("not enough image data")
  569. if s[1] != 0:
  570. raise ValueError("cannot decode image data")
  571. class PyEncoder(PyCodec):
  572. """
  573. Python implementation of a format encoder. Override this class and
  574. add the decoding logic in the :meth:`encode` method.
  575. See :ref:`Writing Your Own File Codec in Python<file-codecs-py>`
  576. """
  577. _pushes_fd = False
  578. @property
  579. def pushes_fd(self):
  580. return self._pushes_fd
  581. def encode(self, bufsize):
  582. """
  583. Override to perform the encoding process.
  584. :param bufsize: Buffer size.
  585. :returns: A tuple of ``(bytes encoded, errcode, bytes)``.
  586. If finished with encoding return 1 for the error code.
  587. Err codes are from :data:`.ImageFile.ERRORS`.
  588. """
  589. raise NotImplementedError()
  590. def encode_to_pyfd(self):
  591. """
  592. If ``pushes_fd`` is ``True``, then this method will be used,
  593. and ``encode()`` will only be called once.
  594. :returns: A tuple of ``(bytes consumed, errcode)``.
  595. Err codes are from :data:`.ImageFile.ERRORS`.
  596. """
  597. if not self.pushes_fd:
  598. return 0, -8 # bad configuration
  599. bytes_consumed, errcode, data = self.encode(0)
  600. if data:
  601. self.fd.write(data)
  602. return bytes_consumed, errcode
  603. def encode_to_file(self, fh, bufsize):
  604. """
  605. :param fh: File handle.
  606. :param bufsize: Buffer size.
  607. :returns: If finished successfully, return 0.
  608. Otherwise, return an error code. Err codes are from
  609. :data:`.ImageFile.ERRORS`.
  610. """
  611. errcode = 0
  612. while errcode == 0:
  613. status, errcode, buf = self.encode(bufsize)
  614. if status > 0:
  615. fh.write(buf[status:])
  616. return errcode