__init__.py 10 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357
  1. import asyncio
  2. from functools import partial, singledispatch
  3. from io import BufferedRandom, BufferedReader, BufferedWriter, FileIO, TextIOBase
  4. from tempfile import NamedTemporaryFile as syncNamedTemporaryFile
  5. from tempfile import SpooledTemporaryFile as syncSpooledTemporaryFile
  6. from tempfile import TemporaryDirectory as syncTemporaryDirectory
  7. from tempfile import TemporaryFile as syncTemporaryFile
  8. from tempfile import _TemporaryFileWrapper as syncTemporaryFileWrapper
  9. from ..base import AiofilesContextManager
  10. from ..threadpool.binary import AsyncBufferedIOBase, AsyncBufferedReader, AsyncFileIO
  11. from ..threadpool.text import AsyncTextIOWrapper
  12. from .temptypes import AsyncSpooledTemporaryFile, AsyncTemporaryDirectory
  13. import sys
  14. __all__ = [
  15. "NamedTemporaryFile",
  16. "TemporaryFile",
  17. "SpooledTemporaryFile",
  18. "TemporaryDirectory",
  19. ]
  20. # ================================================================
  21. # Public methods for async open and return of temp file/directory
  22. # objects with async interface
  23. # ================================================================
  24. if sys.version_info >= (3, 12):
  25. def NamedTemporaryFile(
  26. mode="w+b",
  27. buffering=-1,
  28. encoding=None,
  29. newline=None,
  30. suffix=None,
  31. prefix=None,
  32. dir=None,
  33. delete=True,
  34. delete_on_close=True,
  35. loop=None,
  36. executor=None,
  37. ):
  38. """Async open a named temporary file"""
  39. return AiofilesContextManager(
  40. _temporary_file(
  41. named=True,
  42. mode=mode,
  43. buffering=buffering,
  44. encoding=encoding,
  45. newline=newline,
  46. suffix=suffix,
  47. prefix=prefix,
  48. dir=dir,
  49. delete=delete,
  50. delete_on_close=delete_on_close,
  51. loop=loop,
  52. executor=executor,
  53. )
  54. )
  55. else:
  56. def NamedTemporaryFile(
  57. mode="w+b",
  58. buffering=-1,
  59. encoding=None,
  60. newline=None,
  61. suffix=None,
  62. prefix=None,
  63. dir=None,
  64. delete=True,
  65. loop=None,
  66. executor=None,
  67. ):
  68. """Async open a named temporary file"""
  69. return AiofilesContextManager(
  70. _temporary_file(
  71. named=True,
  72. mode=mode,
  73. buffering=buffering,
  74. encoding=encoding,
  75. newline=newline,
  76. suffix=suffix,
  77. prefix=prefix,
  78. dir=dir,
  79. delete=delete,
  80. loop=loop,
  81. executor=executor,
  82. )
  83. )
  84. def TemporaryFile(
  85. mode="w+b",
  86. buffering=-1,
  87. encoding=None,
  88. newline=None,
  89. suffix=None,
  90. prefix=None,
  91. dir=None,
  92. loop=None,
  93. executor=None,
  94. ):
  95. """Async open an unnamed temporary file"""
  96. return AiofilesContextManager(
  97. _temporary_file(
  98. named=False,
  99. mode=mode,
  100. buffering=buffering,
  101. encoding=encoding,
  102. newline=newline,
  103. suffix=suffix,
  104. prefix=prefix,
  105. dir=dir,
  106. loop=loop,
  107. executor=executor,
  108. )
  109. )
  110. def SpooledTemporaryFile(
  111. max_size=0,
  112. mode="w+b",
  113. buffering=-1,
  114. encoding=None,
  115. newline=None,
  116. suffix=None,
  117. prefix=None,
  118. dir=None,
  119. loop=None,
  120. executor=None,
  121. ):
  122. """Async open a spooled temporary file"""
  123. return AiofilesContextManager(
  124. _spooled_temporary_file(
  125. max_size=max_size,
  126. mode=mode,
  127. buffering=buffering,
  128. encoding=encoding,
  129. newline=newline,
  130. suffix=suffix,
  131. prefix=prefix,
  132. dir=dir,
  133. loop=loop,
  134. executor=executor,
  135. )
  136. )
  137. def TemporaryDirectory(suffix=None, prefix=None, dir=None, loop=None, executor=None):
  138. """Async open a temporary directory"""
  139. return AiofilesContextManagerTempDir(
  140. _temporary_directory(
  141. suffix=suffix, prefix=prefix, dir=dir, loop=loop, executor=executor
  142. )
  143. )
  144. # =========================================================
  145. # Internal coroutines to open new temp files/directories
  146. # =========================================================
  147. if sys.version_info >= (3, 12):
  148. async def _temporary_file(
  149. named=True,
  150. mode="w+b",
  151. buffering=-1,
  152. encoding=None,
  153. newline=None,
  154. suffix=None,
  155. prefix=None,
  156. dir=None,
  157. delete=True,
  158. delete_on_close=True,
  159. loop=None,
  160. executor=None,
  161. max_size=0,
  162. ):
  163. """Async method to open a temporary file with async interface"""
  164. if loop is None:
  165. loop = asyncio.get_running_loop()
  166. if named:
  167. cb = partial(
  168. syncNamedTemporaryFile,
  169. mode=mode,
  170. buffering=buffering,
  171. encoding=encoding,
  172. newline=newline,
  173. suffix=suffix,
  174. prefix=prefix,
  175. dir=dir,
  176. delete=delete,
  177. delete_on_close=delete_on_close,
  178. )
  179. else:
  180. cb = partial(
  181. syncTemporaryFile,
  182. mode=mode,
  183. buffering=buffering,
  184. encoding=encoding,
  185. newline=newline,
  186. suffix=suffix,
  187. prefix=prefix,
  188. dir=dir,
  189. )
  190. f = await loop.run_in_executor(executor, cb)
  191. # Wrap based on type of underlying IO object
  192. if type(f) is syncTemporaryFileWrapper:
  193. # _TemporaryFileWrapper was used (named files)
  194. result = wrap(f.file, f, loop=loop, executor=executor)
  195. result._closer = f._closer
  196. return result
  197. else:
  198. # IO object was returned directly without wrapper
  199. return wrap(f, f, loop=loop, executor=executor)
  200. else:
  201. async def _temporary_file(
  202. named=True,
  203. mode="w+b",
  204. buffering=-1,
  205. encoding=None,
  206. newline=None,
  207. suffix=None,
  208. prefix=None,
  209. dir=None,
  210. delete=True,
  211. loop=None,
  212. executor=None,
  213. max_size=0,
  214. ):
  215. """Async method to open a temporary file with async interface"""
  216. if loop is None:
  217. loop = asyncio.get_running_loop()
  218. if named:
  219. cb = partial(
  220. syncNamedTemporaryFile,
  221. mode=mode,
  222. buffering=buffering,
  223. encoding=encoding,
  224. newline=newline,
  225. suffix=suffix,
  226. prefix=prefix,
  227. dir=dir,
  228. delete=delete,
  229. )
  230. else:
  231. cb = partial(
  232. syncTemporaryFile,
  233. mode=mode,
  234. buffering=buffering,
  235. encoding=encoding,
  236. newline=newline,
  237. suffix=suffix,
  238. prefix=prefix,
  239. dir=dir,
  240. )
  241. f = await loop.run_in_executor(executor, cb)
  242. # Wrap based on type of underlying IO object
  243. if type(f) is syncTemporaryFileWrapper:
  244. # _TemporaryFileWrapper was used (named files)
  245. result = wrap(f.file, f, loop=loop, executor=executor)
  246. # add delete property
  247. result.delete = f.delete
  248. return result
  249. else:
  250. # IO object was returned directly without wrapper
  251. return wrap(f, f, loop=loop, executor=executor)
  252. async def _spooled_temporary_file(
  253. max_size=0,
  254. mode="w+b",
  255. buffering=-1,
  256. encoding=None,
  257. newline=None,
  258. suffix=None,
  259. prefix=None,
  260. dir=None,
  261. loop=None,
  262. executor=None,
  263. ):
  264. """Open a spooled temporary file with async interface"""
  265. if loop is None:
  266. loop = asyncio.get_running_loop()
  267. cb = partial(
  268. syncSpooledTemporaryFile,
  269. max_size=max_size,
  270. mode=mode,
  271. buffering=buffering,
  272. encoding=encoding,
  273. newline=newline,
  274. suffix=suffix,
  275. prefix=prefix,
  276. dir=dir,
  277. )
  278. f = await loop.run_in_executor(executor, cb)
  279. # Single interface provided by SpooledTemporaryFile for all modes
  280. return AsyncSpooledTemporaryFile(f, loop=loop, executor=executor)
  281. async def _temporary_directory(
  282. suffix=None, prefix=None, dir=None, loop=None, executor=None
  283. ):
  284. """Async method to open a temporary directory with async interface"""
  285. if loop is None:
  286. loop = asyncio.get_running_loop()
  287. cb = partial(syncTemporaryDirectory, suffix, prefix, dir)
  288. f = await loop.run_in_executor(executor, cb)
  289. return AsyncTemporaryDirectory(f, loop=loop, executor=executor)
  290. class AiofilesContextManagerTempDir(AiofilesContextManager):
  291. """With returns the directory location, not the object (matching sync lib)"""
  292. async def __aenter__(self):
  293. self._obj = await self._coro
  294. return self._obj.name
  295. @singledispatch
  296. def wrap(base_io_obj, file, *, loop=None, executor=None):
  297. """Wrap the object with interface based on type of underlying IO"""
  298. raise TypeError("Unsupported IO type: {}".format(base_io_obj))
  299. @wrap.register(TextIOBase)
  300. def _(base_io_obj, file, *, loop=None, executor=None):
  301. return AsyncTextIOWrapper(file, loop=loop, executor=executor)
  302. @wrap.register(BufferedWriter)
  303. def _(base_io_obj, file, *, loop=None, executor=None):
  304. return AsyncBufferedIOBase(file, loop=loop, executor=executor)
  305. @wrap.register(BufferedReader)
  306. @wrap.register(BufferedRandom)
  307. def _(base_io_obj, file, *, loop=None, executor=None):
  308. return AsyncBufferedReader(file, loop=loop, executor=executor)
  309. @wrap.register(FileIO)
  310. def _(base_io_obj, file, *, loop=None, executor=None):
  311. return AsyncFileIO(file, loop=loop, executor=executor)