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