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 A value of 0 (the default) disables automatic rollover.
231 :param mode: The mode in which the file is opened. Defaults to "w+b".
232 :param buffering: The buffering policy (-1 means the default buffering).
233 :param encoding: The encoding used to decode or encode the file (text mode only).
234 :param newline: Controls how universal newlines mode works (text mode only).
235 :param suffix: The suffix for the temporary file name.
236 :param prefix: The prefix for the temporary file name.
237 :param dir: The directory in which the temporary file is created.
238 :param errors: The error handling scheme used for encoding/decoding errors.
239 """
240
241 _rolled: bool = False
242
243 @overload
244 def __init__(
245 self: SpooledTemporaryFile[bytes],
246 max_size: int = ...,
247 mode: OpenBinaryMode = ...,
248 buffering: int = ...,
249 encoding: str | None = ...,
250 newline: str | None = ...,
251 suffix: str | None = ...,
252 prefix: str | None = ...,
253 dir: str | None = ...,
254 *,
255 errors: str | None = ...,
256 ): ...
257 @overload
258 def __init__(
259 self: SpooledTemporaryFile[str],
260 max_size: int = ...,
261 mode: OpenTextMode = ...,
262 buffering: int = ...,
263 encoding: str | None = ...,
264 newline: str | None = ...,
265 suffix: str | None = ...,
266 prefix: str | None = ...,
267 dir: str | None = ...,
268 *,
269 errors: str | None = ...,
270 ): ...
271
272 def __init__(
273 self,
274 max_size: int = 0,
275 mode: OpenBinaryMode | OpenTextMode = "w+b",
276 buffering: int = -1,
277 encoding: str | None = None,
278 newline: str | None = None,
279 suffix: str | None = None,
280 prefix: str | None = None,
281 dir: str | None = None,
282 *,
283 errors: str | None = None,
284 ) -> None:
285 self._tempfile_params: dict[str, Any] = {
286 "mode": mode,
287 "buffering": buffering,
288 "encoding": encoding,
289 "newline": newline,
290 "suffix": suffix,
291 "prefix": prefix,
292 "dir": dir,
293 "errors": errors,
294 }
295 self._max_size = max_size
296 if "b" in mode:
297 super().__init__(BytesIO()) # type: ignore[arg-type]
298 else:
299 super().__init__(
300 TextIOWrapper( # type: ignore[arg-type]
301 BytesIO(),
302 encoding=encoding,
303 errors=errors,
304 newline=newline,
305 write_through=True,
306 )
307 )
308
309 async def aclose(self) -> None:
310 if not self._rolled:
311 self._fp.close()
312 return
313
314 await super().aclose()
315
316 async def _check(self) -> None:
317 if self._rolled or not self._max_size or self._fp.tell() <= self._max_size:
318 return
319
320 await self.rollover()
321
322 async def rollover(self) -> None:
323 if self._rolled:
324 return
325
326 await checkpoint_if_cancelled()
327 with CancelScope(shield=True):
328 self._rolled = True
329 buffer = self._fp
330 position = buffer.tell()
331 buffer.seek(0)
332 self._fp = await to_thread.run_sync(
333 lambda: tempfile.TemporaryFile(**self._tempfile_params)
334 )
335 await self.write(buffer.read())
336 await self.seek(position)
337 buffer.close()
338
339 @property
340 def closed(self) -> bool:
341 return self._fp.closed
342
343 async def read(self, size: int = -1) -> AnyStr:
344 if not self._rolled:
345 await checkpoint_if_cancelled()
346 return self._fp.read(size)
347
348 return await super().read(size) # type: ignore[return-value]
349
350 async def read1(self: SpooledTemporaryFile[bytes], size: int = -1) -> bytes:
351 if not self._rolled:
352 await checkpoint_if_cancelled()
353 return self._fp.read1(size)
354
355 return await super().read1(size)
356
357 async def readline(self) -> AnyStr:
358 if not self._rolled:
359 await checkpoint_if_cancelled()
360 return self._fp.readline()
361
362 return await super().readline() # type: ignore[return-value]
363
364 async def readlines(self) -> list[AnyStr]:
365 if not self._rolled:
366 await checkpoint_if_cancelled()
367 return self._fp.readlines()
368
369 return await super().readlines() # type: ignore[return-value]
370
371 async def readinto(self: SpooledTemporaryFile[bytes], b: WriteableBuffer) -> int:
372 if not self._rolled:
373 await checkpoint_if_cancelled()
374 return self._fp.readinto(b)
375
376 return await super().readinto(b)
377
378 async def readinto1(self: SpooledTemporaryFile[bytes], b: WriteableBuffer) -> int:
379 if not self._rolled:
380 await checkpoint_if_cancelled()
381 return self._fp.readinto1(b)
382
383 return await super().readinto1(b)
384
385 async def seek(self, offset: int, whence: int | None = os.SEEK_SET) -> int:
386 if not self._rolled:
387 await checkpoint_if_cancelled()
388 return self._fp.seek(offset, whence)
389
390 return await super().seek(offset, whence)
391
392 async def tell(self) -> int:
393 if not self._rolled:
394 await checkpoint_if_cancelled()
395 return self._fp.tell()
396
397 return await super().tell()
398
399 async def truncate(self, size: int | None = None) -> int:
400 if not self._rolled:
401 await checkpoint_if_cancelled()
402 return self._fp.truncate(size)
403
404 return await super().truncate(size)
405
406 @overload
407 async def write(self: SpooledTemporaryFile[bytes], b: ReadableBuffer) -> int: ...
408 @overload
409 async def write(self: SpooledTemporaryFile[str], b: str) -> int: ...
410
411 async def write(self, b: ReadableBuffer | str) -> int:
412 """
413 Asynchronously write data to the spooled temporary file.
414
415 If the file has not yet been rolled over, the data is written synchronously,
416 and a rollover is triggered if the size exceeds the maximum size.
417
418 :param b: The data to write.
419 :return: The number of bytes written.
420 :raises RuntimeError: If the underlying file is not initialized.
421
422 """
423 if not self._rolled:
424 await checkpoint_if_cancelled()
425 result = self._fp.write(b)
426 await self._check()
427 return result
428
429 return await super().write(b) # type: ignore[misc]
430
431 @overload
432 async def writelines(
433 self: SpooledTemporaryFile[bytes], lines: Iterable[ReadableBuffer]
434 ) -> None: ...
435 @overload
436 async def writelines(
437 self: SpooledTemporaryFile[str], lines: Iterable[str]
438 ) -> None: ...
439
440 async def writelines(self, lines: Iterable[str] | Iterable[ReadableBuffer]) -> None:
441 """
442 Asynchronously write a list of lines to the spooled temporary file.
443
444 If the file has not yet been rolled over, the lines are written synchronously,
445 and a rollover is triggered if the size exceeds the maximum size.
446
447 :param lines: An iterable of lines to write.
448 :raises RuntimeError: If the underlying file is not initialized.
449
450 """
451 if not self._rolled:
452 await checkpoint_if_cancelled()
453 result = self._fp.writelines(lines)
454 await self._check()
455 return result
456
457 return await super().writelines(lines) # type: ignore[misc]
458
459
460class TemporaryDirectory(Generic[AnyStr]):
461 """
462 An asynchronous temporary directory that is created and cleaned up automatically.
463
464 This class provides an asynchronous context manager for creating a temporary
465 directory. It wraps Python's standard :class:`~tempfile.TemporaryDirectory` to
466 perform directory creation and cleanup operations in a background thread.
467
468 :param suffix: Suffix to be added to the temporary directory name.
469 :param prefix: Prefix to be added to the temporary directory name.
470 :param dir: The parent directory where the temporary directory is created.
471 :param ignore_cleanup_errors: Whether to ignore errors during cleanup
472 :param delete: Whether to delete the directory upon closing (Python 3.12+).
473 """
474
475 def __init__(
476 self,
477 suffix: AnyStr | None = None,
478 prefix: AnyStr | None = None,
479 dir: AnyStr | None = None,
480 *,
481 ignore_cleanup_errors: bool = False,
482 delete: bool = True,
483 ) -> None:
484 self.suffix: AnyStr | None = suffix
485 self.prefix: AnyStr | None = prefix
486 self.dir: AnyStr | None = dir
487 self.ignore_cleanup_errors = ignore_cleanup_errors
488 self.delete = delete
489
490 self._tempdir: tempfile.TemporaryDirectory | None = None
491
492 async def __aenter__(self) -> str:
493 params: dict[str, Any] = {
494 "suffix": self.suffix,
495 "prefix": self.prefix,
496 "dir": self.dir,
497 "ignore_cleanup_errors": self.ignore_cleanup_errors,
498 }
499 if sys.version_info >= (3, 12):
500 params["delete"] = self.delete
501
502 self._tempdir = await to_thread.run_sync(
503 lambda: tempfile.TemporaryDirectory(**params)
504 )
505 return self._tempdir.name
506
507 async def __aexit__(
508 self,
509 exc_type: type[BaseException] | None,
510 exc_value: BaseException | None,
511 traceback: TracebackType | None,
512 ) -> None:
513 if self._tempdir is not None:
514 with CancelScope(shield=True):
515 await to_thread.run_sync(
516 self._tempdir.__exit__, exc_type, exc_value, traceback
517 )
518
519 async def cleanup(self) -> None:
520 if self._tempdir is not None:
521 with CancelScope(shield=True):
522 await to_thread.run_sync(self._tempdir.cleanup)
523
524 await checkpoint_if_cancelled()
525
526
527@overload
528async def mkstemp(
529 suffix: str | None = None,
530 prefix: str | None = None,
531 dir: str | None = None,
532 text: bool = False,
533) -> tuple[int, str]: ...
534
535
536@overload
537async def mkstemp(
538 suffix: bytes | None = None,
539 prefix: bytes | None = None,
540 dir: bytes | None = None,
541 text: bool = False,
542) -> tuple[int, bytes]: ...
543
544
545async def mkstemp(
546 suffix: AnyStr | None = None,
547 prefix: AnyStr | None = None,
548 dir: AnyStr | None = None,
549 text: bool = False,
550) -> tuple[int, str | bytes]:
551 """
552 Asynchronously create a temporary file and return an OS-level handle and the file
553 name.
554
555 This function wraps `tempfile.mkstemp` and executes it in a background thread.
556
557 :param suffix: Suffix to be added to the file name.
558 :param prefix: Prefix to be added to the file name.
559 :param dir: Directory in which the temporary file is created.
560 :param text: Whether the file is opened in text mode.
561 :return: A tuple containing the file descriptor and the file name.
562
563 """
564 return await to_thread.run_sync(tempfile.mkstemp, suffix, prefix, dir, text)
565
566
567@overload
568async def mkdtemp(
569 suffix: str | None = None,
570 prefix: str | None = None,
571 dir: str | None = None,
572) -> str: ...
573
574
575@overload
576async def mkdtemp(
577 suffix: bytes | None = None,
578 prefix: bytes | None = None,
579 dir: bytes | None = None,
580) -> bytes: ...
581
582
583async def mkdtemp(
584 suffix: AnyStr | None = None,
585 prefix: AnyStr | None = None,
586 dir: AnyStr | None = None,
587) -> str | bytes:
588 """
589 Asynchronously create a temporary directory and return its path.
590
591 This function wraps `tempfile.mkdtemp` and executes it in a background thread.
592
593 :param suffix: Suffix to be added to the directory name.
594 :param prefix: Prefix to be added to the directory name.
595 :param dir: Parent directory where the temporary directory is created.
596 :return: The path of the created temporary directory.
597
598 """
599 return await to_thread.run_sync(tempfile.mkdtemp, suffix, prefix, dir)
600
601
602async def gettempdir() -> str:
603 """
604 Asynchronously return the name of the directory used for temporary files.
605
606 This function wraps `tempfile.gettempdir` and executes it in a background thread.
607
608 :return: The path of the temporary directory as a string.
609
610 """
611 return await to_thread.run_sync(tempfile.gettempdir)
612
613
614async def gettempdirb() -> bytes:
615 """
616 Asynchronously return the name of the directory used for temporary files in bytes.
617
618 This function wraps `tempfile.gettempdirb` and executes it in a background thread.
619
620 :return: The path of the temporary directory as bytes.
621
622 """
623 return await to_thread.run_sync(tempfile.gettempdirb)