1from __future__ import annotations
2
3import os
4import sys
5import tempfile
6from collections.abc import Iterable
7from io import BytesIO, TextIOWrapper
8from types import TracebackType
9from typing import (
10 TYPE_CHECKING,
11 Any,
12 AnyStr,
13 Generic,
14 overload,
15)
16
17from .. import to_thread
18from .._core._fileio import AsyncFile
19from ..lowlevel import checkpoint_if_cancelled
20from ._tasks import CancelScope
21
22if TYPE_CHECKING:
23 from _typeshed import OpenBinaryMode, OpenTextMode, ReadableBuffer, WriteableBuffer
24
25
26class TemporaryFile(Generic[AnyStr]):
27 """
28 An asynchronous temporary file that is automatically created and cleaned up.
29
30 This class provides an asynchronous context manager interface to a temporary file.
31 The file is created using Python's standard `tempfile.TemporaryFile` function in a
32 background thread, and is wrapped as an asynchronous file using `AsyncFile`.
33
34 :param mode: The mode in which the file is opened. Defaults to "w+b".
35 :param buffering: The buffering policy (-1 means the default buffering).
36 :param encoding: The encoding used to decode or encode the file. Only applicable in
37 text mode.
38 :param newline: Controls how universal newlines mode works (only applicable in text
39 mode).
40 :param suffix: The suffix for the temporary file name.
41 :param prefix: The prefix for the temporary file name.
42 :param dir: The directory in which the temporary file is created.
43 :param errors: The error handling scheme used for encoding/decoding errors.
44 """
45
46 _async_file: AsyncFile[AnyStr]
47
48 @overload
49 def __init__(
50 self: TemporaryFile[bytes],
51 mode: OpenBinaryMode = ...,
52 buffering: int = ...,
53 encoding: str | None = ...,
54 newline: str | None = ...,
55 suffix: str | None = ...,
56 prefix: str | None = ...,
57 dir: str | None = ...,
58 *,
59 errors: str | None = ...,
60 ): ...
61 @overload
62 def __init__(
63 self: TemporaryFile[str],
64 mode: OpenTextMode,
65 buffering: int = ...,
66 encoding: str | None = ...,
67 newline: str | None = ...,
68 suffix: str | None = ...,
69 prefix: str | None = ...,
70 dir: str | None = ...,
71 *,
72 errors: str | None = ...,
73 ): ...
74
75 def __init__(
76 self,
77 mode: OpenTextMode | OpenBinaryMode = "w+b",
78 buffering: int = -1,
79 encoding: str | None = None,
80 newline: str | None = None,
81 suffix: str | None = None,
82 prefix: str | None = None,
83 dir: str | None = None,
84 *,
85 errors: str | None = None,
86 ) -> None:
87 self.mode = mode
88 self.buffering = buffering
89 self.encoding = encoding
90 self.newline = newline
91 self.suffix: str | None = suffix
92 self.prefix: str | None = prefix
93 self.dir: str | None = dir
94 self.errors = errors
95
96 async def __aenter__(self) -> AsyncFile[AnyStr]:
97 fp = await to_thread.run_sync(
98 lambda: tempfile.TemporaryFile(
99 self.mode,
100 self.buffering,
101 self.encoding,
102 self.newline,
103 self.suffix,
104 self.prefix,
105 self.dir,
106 errors=self.errors,
107 )
108 )
109 self._async_file = AsyncFile(fp)
110 return self._async_file
111
112 async def __aexit__(
113 self,
114 exc_type: type[BaseException] | None,
115 exc_value: BaseException | None,
116 traceback: TracebackType | None,
117 ) -> None:
118 await self._async_file.aclose()
119
120
121class NamedTemporaryFile(Generic[AnyStr]):
122 """
123 An asynchronous named temporary file that is automatically created and cleaned up.
124
125 This class provides an asynchronous context manager for a temporary file with a
126 visible name in the file system. It uses Python's standard
127 :func:`~tempfile.NamedTemporaryFile` function and wraps the file object with
128 :class:`AsyncFile` for asynchronous operations.
129
130 :param mode: The mode in which the file is opened. Defaults to "w+b".
131 :param buffering: The buffering policy (-1 means the default buffering).
132 :param encoding: The encoding used to decode or encode the file. Only applicable in
133 text mode.
134 :param newline: Controls how universal newlines mode works (only applicable in text
135 mode).
136 :param suffix: The suffix for the temporary file name.
137 :param prefix: The prefix for the temporary file name.
138 :param dir: The directory in which the temporary file is created.
139 :param delete: Whether to delete the file when it is closed.
140 :param errors: The error handling scheme used for encoding/decoding errors.
141 :param delete_on_close: (Python 3.12+) Whether to delete the file on close.
142 """
143
144 _async_file: AsyncFile[AnyStr]
145
146 @overload
147 def __init__(
148 self: NamedTemporaryFile[bytes],
149 mode: OpenBinaryMode = ...,
150 buffering: int = ...,
151 encoding: str | None = ...,
152 newline: str | None = ...,
153 suffix: str | None = ...,
154 prefix: str | None = ...,
155 dir: str | None = ...,
156 delete: bool = ...,
157 *,
158 errors: str | None = ...,
159 delete_on_close: bool = ...,
160 ): ...
161 @overload
162 def __init__(
163 self: NamedTemporaryFile[str],
164 mode: OpenTextMode,
165 buffering: int = ...,
166 encoding: str | None = ...,
167 newline: str | None = ...,
168 suffix: str | None = ...,
169 prefix: str | None = ...,
170 dir: str | None = ...,
171 delete: bool = ...,
172 *,
173 errors: str | None = ...,
174 delete_on_close: bool = ...,
175 ): ...
176
177 def __init__(
178 self,
179 mode: OpenBinaryMode | OpenTextMode = "w+b",
180 buffering: int = -1,
181 encoding: str | None = None,
182 newline: str | None = None,
183 suffix: str | None = None,
184 prefix: str | None = None,
185 dir: str | None = None,
186 delete: bool = True,
187 *,
188 errors: str | None = None,
189 delete_on_close: bool = True,
190 ) -> None:
191 self._params: dict[str, Any] = {
192 "mode": mode,
193 "buffering": buffering,
194 "encoding": encoding,
195 "newline": newline,
196 "suffix": suffix,
197 "prefix": prefix,
198 "dir": dir,
199 "delete": delete,
200 "errors": errors,
201 }
202 if sys.version_info >= (3, 12):
203 self._params["delete_on_close"] = delete_on_close
204
205 async def __aenter__(self) -> AsyncFile[AnyStr]:
206 fp = await to_thread.run_sync(
207 lambda: tempfile.NamedTemporaryFile(**self._params)
208 )
209 self._async_file = AsyncFile(fp)
210 return self._async_file
211
212 async def __aexit__(
213 self,
214 exc_type: type[BaseException] | None,
215 exc_value: BaseException | None,
216 traceback: TracebackType | None,
217 ) -> None:
218 await self._async_file.aclose()
219
220
221class SpooledTemporaryFile(AsyncFile[AnyStr]):
222 """
223 An asynchronous spooled temporary file that starts in memory and is spooled to disk.
224
225 This class provides an asynchronous interface to a spooled temporary file, much like
226 Python's standard :class:`~tempfile.SpooledTemporaryFile`. It supports asynchronous
227 write operations and provides a method to force a rollover to disk.
228
229 :param max_size: Maximum size in bytes before the file is rolled over to disk.
230 :param mode: The mode in which the file is opened. Defaults to "w+b".
231 :param buffering: The buffering policy (-1 means the default buffering).
232 :param encoding: The encoding used to decode or encode the file (text mode only).
233 :param newline: Controls how universal newlines mode works (text mode only).
234 :param suffix: The suffix for the temporary file name.
235 :param prefix: The prefix for the temporary file name.
236 :param dir: The directory in which the temporary file is created.
237 :param errors: The error handling scheme used for encoding/decoding errors.
238 """
239
240 _rolled: bool = False
241
242 @overload
243 def __init__(
244 self: SpooledTemporaryFile[bytes],
245 max_size: int = ...,
246 mode: OpenBinaryMode = ...,
247 buffering: int = ...,
248 encoding: str | None = ...,
249 newline: str | None = ...,
250 suffix: str | None = ...,
251 prefix: str | None = ...,
252 dir: str | None = ...,
253 *,
254 errors: str | None = ...,
255 ): ...
256 @overload
257 def __init__(
258 self: SpooledTemporaryFile[str],
259 max_size: int = ...,
260 mode: OpenTextMode = ...,
261 buffering: int = ...,
262 encoding: str | None = ...,
263 newline: str | None = ...,
264 suffix: str | None = ...,
265 prefix: str | None = ...,
266 dir: str | None = ...,
267 *,
268 errors: str | None = ...,
269 ): ...
270
271 def __init__(
272 self,
273 max_size: int = 0,
274 mode: OpenBinaryMode | OpenTextMode = "w+b",
275 buffering: int = -1,
276 encoding: str | None = None,
277 newline: str | None = None,
278 suffix: str | None = None,
279 prefix: str | None = None,
280 dir: str | None = None,
281 *,
282 errors: str | None = None,
283 ) -> None:
284 self._tempfile_params: dict[str, Any] = {
285 "mode": mode,
286 "buffering": buffering,
287 "encoding": encoding,
288 "newline": newline,
289 "suffix": suffix,
290 "prefix": prefix,
291 "dir": dir,
292 "errors": errors,
293 }
294 self._max_size = max_size
295 if "b" in mode:
296 super().__init__(BytesIO()) # type: ignore[arg-type]
297 else:
298 super().__init__(
299 TextIOWrapper( # type: ignore[arg-type]
300 BytesIO(),
301 encoding=encoding,
302 errors=errors,
303 newline=newline,
304 write_through=True,
305 )
306 )
307
308 async def aclose(self) -> None:
309 if not self._rolled:
310 self._fp.close()
311 return
312
313 await super().aclose()
314
315 async def _check(self) -> None:
316 if self._rolled or self._fp.tell() <= self._max_size:
317 return
318
319 await self.rollover()
320
321 async def rollover(self) -> None:
322 if self._rolled:
323 return
324
325 self._rolled = True
326 buffer = self._fp
327 buffer.seek(0)
328 self._fp = await to_thread.run_sync(
329 lambda: tempfile.TemporaryFile(**self._tempfile_params)
330 )
331 await self.write(buffer.read())
332 buffer.close()
333
334 @property
335 def closed(self) -> bool:
336 return self._fp.closed
337
338 async def read(self, size: int = -1) -> AnyStr:
339 if not self._rolled:
340 await checkpoint_if_cancelled()
341 return self._fp.read(size)
342
343 return await super().read(size) # type: ignore[return-value]
344
345 async def read1(self: SpooledTemporaryFile[bytes], size: int = -1) -> bytes:
346 if not self._rolled:
347 await checkpoint_if_cancelled()
348 return self._fp.read1(size)
349
350 return await super().read1(size)
351
352 async def readline(self) -> AnyStr:
353 if not self._rolled:
354 await checkpoint_if_cancelled()
355 return self._fp.readline()
356
357 return await super().readline() # type: ignore[return-value]
358
359 async def readlines(self) -> list[AnyStr]:
360 if not self._rolled:
361 await checkpoint_if_cancelled()
362 return self._fp.readlines()
363
364 return await super().readlines() # type: ignore[return-value]
365
366 async def readinto(self: SpooledTemporaryFile[bytes], b: WriteableBuffer) -> int:
367 if not self._rolled:
368 await checkpoint_if_cancelled()
369 return self._fp.readinto(b)
370
371 return await super().readinto(b)
372
373 async def readinto1(self: SpooledTemporaryFile[bytes], b: WriteableBuffer) -> int:
374 if not self._rolled:
375 await checkpoint_if_cancelled()
376 return self._fp.readinto1(b)
377
378 return await super().readinto1(b)
379
380 async def seek(self, offset: int, whence: int | None = os.SEEK_SET) -> int:
381 if not self._rolled:
382 await checkpoint_if_cancelled()
383 return self._fp.seek(offset, whence)
384
385 return await super().seek(offset, whence)
386
387 async def tell(self) -> int:
388 if not self._rolled:
389 await checkpoint_if_cancelled()
390 return self._fp.tell()
391
392 return await super().tell()
393
394 async def truncate(self, size: int | None = None) -> int:
395 if not self._rolled:
396 await checkpoint_if_cancelled()
397 return self._fp.truncate(size)
398
399 return await super().truncate(size)
400
401 @overload
402 async def write(self: SpooledTemporaryFile[bytes], b: ReadableBuffer) -> int: ...
403 @overload
404 async def write(self: SpooledTemporaryFile[str], b: str) -> int: ...
405
406 async def write(self, b: ReadableBuffer | str) -> int:
407 """
408 Asynchronously write data to the spooled temporary file.
409
410 If the file has not yet been rolled over, the data is written synchronously,
411 and a rollover is triggered if the size exceeds the maximum size.
412
413 :param s: The data to write.
414 :return: The number of bytes written.
415 :raises RuntimeError: If the underlying file is not initialized.
416
417 """
418 if not self._rolled:
419 await checkpoint_if_cancelled()
420 result = self._fp.write(b)
421 await self._check()
422 return result
423
424 return await super().write(b) # type: ignore[misc]
425
426 @overload
427 async def writelines(
428 self: SpooledTemporaryFile[bytes], lines: Iterable[ReadableBuffer]
429 ) -> None: ...
430 @overload
431 async def writelines(
432 self: SpooledTemporaryFile[str], lines: Iterable[str]
433 ) -> None: ...
434
435 async def writelines(self, lines: Iterable[str] | Iterable[ReadableBuffer]) -> None:
436 """
437 Asynchronously write a list of lines to the spooled temporary file.
438
439 If the file has not yet been rolled over, the lines are written synchronously,
440 and a rollover is triggered if the size exceeds the maximum size.
441
442 :param lines: An iterable of lines to write.
443 :raises RuntimeError: If the underlying file is not initialized.
444
445 """
446 if not self._rolled:
447 await checkpoint_if_cancelled()
448 result = self._fp.writelines(lines)
449 await self._check()
450 return result
451
452 return await super().writelines(lines) # type: ignore[misc]
453
454
455class TemporaryDirectory(Generic[AnyStr]):
456 """
457 An asynchronous temporary directory that is created and cleaned up automatically.
458
459 This class provides an asynchronous context manager for creating a temporary
460 directory. It wraps Python's standard :class:`~tempfile.TemporaryDirectory` to
461 perform directory creation and cleanup operations in a background thread.
462
463 :param suffix: Suffix to be added to the temporary directory name.
464 :param prefix: Prefix to be added to the temporary directory name.
465 :param dir: The parent directory where the temporary directory is created.
466 :param ignore_cleanup_errors: Whether to ignore errors during cleanup
467 :param delete: Whether to delete the directory upon closing (Python 3.12+).
468 """
469
470 def __init__(
471 self,
472 suffix: AnyStr | None = None,
473 prefix: AnyStr | None = None,
474 dir: AnyStr | None = None,
475 *,
476 ignore_cleanup_errors: bool = False,
477 delete: bool = True,
478 ) -> None:
479 self.suffix: AnyStr | None = suffix
480 self.prefix: AnyStr | None = prefix
481 self.dir: AnyStr | None = dir
482 self.ignore_cleanup_errors = ignore_cleanup_errors
483 self.delete = delete
484
485 self._tempdir: tempfile.TemporaryDirectory | None = None
486
487 async def __aenter__(self) -> str:
488 params: dict[str, Any] = {
489 "suffix": self.suffix,
490 "prefix": self.prefix,
491 "dir": self.dir,
492 "ignore_cleanup_errors": self.ignore_cleanup_errors,
493 }
494 if sys.version_info >= (3, 12):
495 params["delete"] = self.delete
496
497 self._tempdir = await to_thread.run_sync(
498 lambda: tempfile.TemporaryDirectory(**params)
499 )
500 return await to_thread.run_sync(self._tempdir.__enter__)
501
502 async def __aexit__(
503 self,
504 exc_type: type[BaseException] | None,
505 exc_value: BaseException | None,
506 traceback: TracebackType | None,
507 ) -> None:
508 if self._tempdir is not None:
509 with CancelScope(shield=True):
510 await to_thread.run_sync(
511 self._tempdir.__exit__, exc_type, exc_value, traceback
512 )
513
514 async def cleanup(self) -> None:
515 if self._tempdir is not None:
516 await to_thread.run_sync(self._tempdir.cleanup)
517
518
519@overload
520async def mkstemp(
521 suffix: str | None = None,
522 prefix: str | None = None,
523 dir: str | None = None,
524 text: bool = False,
525) -> tuple[int, str]: ...
526
527
528@overload
529async def mkstemp(
530 suffix: bytes | None = None,
531 prefix: bytes | None = None,
532 dir: bytes | None = None,
533 text: bool = False,
534) -> tuple[int, bytes]: ...
535
536
537async def mkstemp(
538 suffix: AnyStr | None = None,
539 prefix: AnyStr | None = None,
540 dir: AnyStr | None = None,
541 text: bool = False,
542) -> tuple[int, str | bytes]:
543 """
544 Asynchronously create a temporary file and return an OS-level handle and the file
545 name.
546
547 This function wraps `tempfile.mkstemp` and executes it in a background thread.
548
549 :param suffix: Suffix to be added to the file name.
550 :param prefix: Prefix to be added to the file name.
551 :param dir: Directory in which the temporary file is created.
552 :param text: Whether the file is opened in text mode.
553 :return: A tuple containing the file descriptor and the file name.
554
555 """
556 return await to_thread.run_sync(tempfile.mkstemp, suffix, prefix, dir, text)
557
558
559@overload
560async def mkdtemp(
561 suffix: str | None = None,
562 prefix: str | None = None,
563 dir: str | None = None,
564) -> str: ...
565
566
567@overload
568async def mkdtemp(
569 suffix: bytes | None = None,
570 prefix: bytes | None = None,
571 dir: bytes | None = None,
572) -> bytes: ...
573
574
575async def mkdtemp(
576 suffix: AnyStr | None = None,
577 prefix: AnyStr | None = None,
578 dir: AnyStr | None = None,
579) -> str | bytes:
580 """
581 Asynchronously create a temporary directory and return its path.
582
583 This function wraps `tempfile.mkdtemp` and executes it in a background thread.
584
585 :param suffix: Suffix to be added to the directory name.
586 :param prefix: Prefix to be added to the directory name.
587 :param dir: Parent directory where the temporary directory is created.
588 :return: The path of the created temporary directory.
589
590 """
591 return await to_thread.run_sync(tempfile.mkdtemp, suffix, prefix, dir)
592
593
594async def gettempdir() -> str:
595 """
596 Asynchronously return the name of the directory used for temporary files.
597
598 This function wraps `tempfile.gettempdir` and executes it in a background thread.
599
600 :return: The path of the temporary directory as a string.
601
602 """
603 return await to_thread.run_sync(tempfile.gettempdir)
604
605
606async def gettempdirb() -> bytes:
607 """
608 Asynchronously return the name of the directory used for temporary files in bytes.
609
610 This function wraps `tempfile.gettempdirb` and executes it in a background thread.
611
612 :return: The path of the temporary directory as bytes.
613
614 """
615 return await to_thread.run_sync(tempfile.gettempdirb)