1
0
Fork 0
mirror of https://github.com/yt-dlp/yt-dlp.git synced 2025-03-09 12:50:23 -05:00

temp file wrapper

This commit is contained in:
c-basalt 2024-08-11 16:23:58 -04:00
parent 8e39820adc
commit 567fb56a1c

View file

@ -47,39 +47,45 @@ def cookie_jar_to_list(cookie_jar):
return [cookie_to_dict(cookie) for cookie in cookie_jar] return [cookie_to_dict(cookie) for cookie in cookie_jar]
@contextlib.contextmanager class TempFileWrapper:
def _temp_file(content, *, mode='wt', encoding='utf-8', suffix=None, close=True): def __init__(self, content=None, text=True, encoding='utf-8', suffix=None):
if 'b' in mode: self.encoding = None if not text else encoding
encoding = None self.text = text
temp_file_handle = tempfile.NamedTemporaryFile(mode, encoding=encoding, suffix=suffix, delete=False) self._file = tempfile.NamedTemporaryFile('wb', suffix=suffix, delete=False)
try: self._file.close()
temp_file_handle.write(content) if content:
if close: self.write(content)
temp_file_handle.close()
yield temp_file_handle
finally:
with contextlib.suppress(OSError):
os.remove(temp_file_handle.name)
@property
def name(self):
return self._file.name
@contextlib.contextmanager @contextlib.contextmanager
def _tempfile_context(): def opened_file(self, mode, *, seek=None, seek_whence=0):
handles = [] mode = mode if (self.text or 'b' in mode) else mode + 'b'
with open(self._file.name, mode, encoding=self.encoding) as f:
if seek is not None:
self._file.seek(seek, seek_whence)
yield f
def _creater(content, *, mode='wt', encoding='utf-8', suffix=None, close=True): def write(self, s, seek=None, seek_whence=0):
encoding = None if 'b' in mode else encoding with self.opened_file('w', seek=seek, seek_whence=seek_whence) as f:
handle = tempfile.NamedTemporaryFile(mode, encoding=encoding, suffix=suffix, delete=False) f.write(s)
handles.append(handle)
handle.write(content) def append_write(self, s, seek=None, seek_whence=0):
if close: with self.opened_file('a', seek=seek, seek_whence=seek_whence) as f:
handle.close() f.write(s)
return handle
try: def read(self, n=-1, seek=None, seek_whence=0):
yield _creater with self.opened_file('r', seek=seek, seek_whence=seek_whence) as f:
finally: return f.read(n)
for handle in handles:
def cleanup(self):
with contextlib.suppress(OSError): with contextlib.suppress(OSError):
os.remove(handle.name) os.remove(self._file.name)
def __del__(self):
self.cleanup()
class ExternalJSI: class ExternalJSI:
@ -117,7 +123,7 @@ def __init__(self, extractor: InfoExtractor, required_version=None, timeout=1000
@classmethod @classmethod
def _execute(cls, jscode, extractor=None, video_id=None, note='', flags=[], timeout=10000): def _execute(cls, jscode, extractor=None, video_id=None, note='', flags=[], timeout=10000):
with _temp_file(jscode, suffix='.js') as js_file: js_file = TempFileWrapper(jscode)
if note and extractor: if note and extractor:
extractor.to_screen(f'{format_field(video_id, None, "%s: ")}{note}') extractor.to_screen(f'{format_field(video_id, None, "%s: ")}{note}')
cmd = [cls.exe, 'run', *flags, js_file.name] cmd = [cls.exe, 'run', *flags, js_file.name]