v6d3music/v6d3music/core/caching.py
2022-12-24 16:37:05 +00:00

68 lines
2.5 KiB
Python

import asyncio
from contextlib import AsyncExitStack
from v6d3music.config import myroot
from v6d3music.utils.tor_prefix import tor_prefix
from ptvp35 import *
from v6d2ctx.lock_for import lock_for
__all__ = ('Caching',)
cache_root = myroot / 'cache'
cache_root.mkdir(exist_ok=True)
class Caching:
async def cache_url(self, hurl: str, rurl: str, override: bool, tor: bool) -> None:
async with lock_for(('cache', hurl), 'cache failed'):
if not override and self.__db.get(f'url:{hurl}', None) is not None:
return
cachable: bool = self.__db.get(f'cachable:{hurl}', False)
if cachable:
print('caching', hurl)
path = cache_root / f'{hurl}.opus'
tmp_path = cache_root / f'{hurl}.tmp.opus'
args = []
if tor:
args.extend(tor_prefix())
args.extend(
[
'ffmpeg', '-hide_banner', '-loglevel', 'warning',
'-reconnect', '1', '-reconnect_at_eof', '0',
'-reconnect_streamed', '1', '-reconnect_delay_max', '10', '-copy_unknown',
'-y', '-i', rurl, '-b:a', '128k', str(tmp_path)
]
)
ap = await asyncio.create_subprocess_exec(*args)
code = await ap.wait()
if code:
print(f'caching {hurl} failed with {code}')
return
await asyncio.to_thread(tmp_path.rename, path)
await self.__db.set(f'url:{hurl}', str(path))
print('cached', hurl)
# await cache_db.set(f'cachable:{hurl}', False)
else:
await self.__db.set(f'cachable:{hurl}', True)
def get(self, hurl: str) -> str | None:
return self.__db.get(f'url:{hurl}', None)
async def __aenter__(self) -> 'Caching':
es = AsyncExitStack()
async with es:
self.__db = await es.enter_async_context(DbFactory(myroot / 'cache.db', kvfactory=KVJson()))
self.__tasks = set()
self.__es = es.pop_all()
return self
async def __aexit__(self, exc_type, exc_val, exc_tb):
async with self.__es:
del self.__es
def schedule_cache(self, hurl: str, rurl: str, override: bool, tor: bool):
task = asyncio.create_task(self.cache_url(hurl, rurl, override, tor))
self.__tasks.add(task)
task.add_done_callback(self.__tasks.discard)