filewrapper.py 4.3 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121
  1. # SPDX-FileCopyrightText: 2015 Eric Larson
  2. #
  3. # SPDX-License-Identifier: Apache-2.0
  4. from __future__ import annotations
  5. import mmap
  6. from tempfile import NamedTemporaryFile
  7. from typing import TYPE_CHECKING, Any, Callable
  8. if TYPE_CHECKING:
  9. from collections.abc import Buffer
  10. from http.client import HTTPResponse
  11. class CallbackFileWrapper:
  12. """
  13. Small wrapper around a fp object which will tee everything read into a
  14. buffer, and when that file is closed it will execute a callback with the
  15. contents of that buffer.
  16. All attributes are proxied to the underlying file object.
  17. This class uses members with a double underscore (__) leading prefix so as
  18. not to accidentally shadow an attribute.
  19. The data is stored in a temporary file until it is all available. As long
  20. as the temporary files directory is disk-based (sometimes it's a
  21. memory-backed-``tmpfs`` on Linux), data will be unloaded to disk if memory
  22. pressure is high. For small files the disk usually won't be used at all,
  23. it'll all be in the filesystem memory cache, so there should be no
  24. performance impact.
  25. """
  26. def __init__(
  27. self, fp: HTTPResponse, callback: Callable[[Buffer], None] | None
  28. ) -> None:
  29. self.__buf = NamedTemporaryFile("rb+", delete=True)
  30. self.__fp = fp
  31. self.__callback = callback
  32. def __getattr__(self, name: str) -> Any:
  33. # The vagaries of garbage collection means that self.__fp is
  34. # not always set. By using __getattribute__ and the private
  35. # name[0] allows looking up the attribute value and raising an
  36. # AttributeError when it doesn't exist. This stop things from
  37. # infinitely recursing calls to getattr in the case where
  38. # self.__fp hasn't been set.
  39. #
  40. # [0] https://docs.python.org/2/reference/expressions.html#atom-identifiers
  41. fp = self.__getattribute__("_CallbackFileWrapper__fp")
  42. return getattr(fp, name)
  43. def __is_fp_closed(self) -> bool:
  44. try:
  45. return self.__fp.fp is None
  46. except AttributeError:
  47. pass
  48. try:
  49. closed: bool = self.__fp.closed
  50. return closed
  51. except AttributeError:
  52. pass
  53. # We just don't cache it then.
  54. # TODO: Add some logging here...
  55. return False
  56. def _close(self) -> None:
  57. result: Buffer
  58. if self.__callback:
  59. if self.__buf.tell() == 0:
  60. # Empty file:
  61. result = b""
  62. else:
  63. # Return the data without actually loading it into memory,
  64. # relying on Python's buffer API and mmap(). mmap() just gives
  65. # a view directly into the filesystem's memory cache, so it
  66. # doesn't result in duplicate memory use.
  67. self.__buf.seek(0, 0)
  68. result = memoryview(
  69. mmap.mmap(self.__buf.fileno(), 0, access=mmap.ACCESS_READ)
  70. )
  71. self.__callback(result)
  72. # We assign this to None here, because otherwise we can get into
  73. # really tricky problems where the CPython interpreter dead locks
  74. # because the callback is holding a reference to something which
  75. # has a __del__ method. Setting this to None breaks the cycle
  76. # and allows the garbage collector to do it's thing normally.
  77. self.__callback = None
  78. # Closing the temporary file releases memory and frees disk space.
  79. # Important when caching big files.
  80. self.__buf.close()
  81. def read(self, amt: int | None = None) -> bytes:
  82. data: bytes = self.__fp.read(amt)
  83. if data:
  84. # We may be dealing with b'', a sign that things are over:
  85. # it's passed e.g. after we've already closed self.__buf.
  86. self.__buf.write(data)
  87. if self.__is_fp_closed():
  88. self._close()
  89. return data
  90. def _safe_read(self, amt: int) -> bytes:
  91. data: bytes = self.__fp._safe_read(amt) # type: ignore[attr-defined]
  92. if amt == 2 and data == b"\r\n":
  93. # urllib executes this read to toss the CRLF at the end
  94. # of the chunk.
  95. return data
  96. self.__buf.write(data)
  97. if self.__is_fp_closed():
  98. self._close()
  99. return data