]>
Commit | Line | Data |
---|---|---|
1 | import contextlib | |
2 | import json | |
3 | import os | |
4 | import re | |
5 | import shutil | |
6 | import traceback | |
7 | import urllib.parse | |
8 | ||
9 | from .utils import expand_path, traverse_obj, version_tuple, write_json_file | |
10 | from .version import __version__ | |
11 | ||
12 | ||
13 | class Cache: | |
14 | def __init__(self, ydl): | |
15 | self._ydl = ydl | |
16 | ||
17 | def _get_root_dir(self): | |
18 | res = self._ydl.params.get('cachedir') | |
19 | if res is None: | |
20 | cache_root = os.getenv('XDG_CACHE_HOME', '~/.cache') | |
21 | res = os.path.join(cache_root, 'yt-dlp') | |
22 | return expand_path(res) | |
23 | ||
24 | def _get_cache_fn(self, section, key, dtype): | |
25 | assert re.match(r'^[\w.-]+$', section), f'invalid section {section!r}' | |
26 | key = urllib.parse.quote(key, safe='').replace('%', ',') # encode non-ascii characters | |
27 | return os.path.join(self._get_root_dir(), section, f'{key}.{dtype}') | |
28 | ||
29 | @property | |
30 | def enabled(self): | |
31 | return self._ydl.params.get('cachedir') is not False | |
32 | ||
33 | def store(self, section, key, data, dtype='json'): | |
34 | assert dtype in ('json',) | |
35 | ||
36 | if not self.enabled: | |
37 | return | |
38 | ||
39 | fn = self._get_cache_fn(section, key, dtype) | |
40 | try: | |
41 | os.makedirs(os.path.dirname(fn), exist_ok=True) | |
42 | self._ydl.write_debug(f'Saving {section}.{key} to cache') | |
43 | write_json_file({'yt-dlp_version': __version__, 'data': data}, fn) | |
44 | except Exception: | |
45 | tb = traceback.format_exc() | |
46 | self._ydl.report_warning(f'Writing cache to {fn!r} failed: {tb}') | |
47 | ||
48 | def _validate(self, data, min_ver): | |
49 | version = traverse_obj(data, 'yt-dlp_version') | |
50 | if not version: # Backward compatibility | |
51 | data, version = {'data': data}, '2022.08.19' | |
52 | if not min_ver or version_tuple(version) >= version_tuple(min_ver): | |
53 | return data['data'] | |
54 | self._ydl.write_debug(f'Discarding old cache from version {version} (needs {min_ver})') | |
55 | ||
56 | def load(self, section, key, dtype='json', default=None, *, min_ver=None): | |
57 | assert dtype in ('json',) | |
58 | ||
59 | if not self.enabled: | |
60 | return default | |
61 | ||
62 | cache_fn = self._get_cache_fn(section, key, dtype) | |
63 | with contextlib.suppress(OSError): | |
64 | try: | |
65 | with open(cache_fn, encoding='utf-8') as cachef: | |
66 | self._ydl.write_debug(f'Loading {section}.{key} from cache') | |
67 | return self._validate(json.load(cachef), min_ver) | |
68 | except (ValueError, KeyError): | |
69 | try: | |
70 | file_size = os.path.getsize(cache_fn) | |
71 | except OSError as oe: | |
72 | file_size = str(oe) | |
73 | self._ydl.report_warning(f'Cache retrieval from {cache_fn} failed ({file_size})') | |
74 | ||
75 | return default | |
76 | ||
77 | def remove(self): | |
78 | if not self.enabled: | |
79 | self._ydl.to_screen('Cache is disabled (Did you combine --no-cache-dir and --rm-cache-dir?)') | |
80 | return | |
81 | ||
82 | cachedir = self._get_root_dir() | |
83 | if not any((term in cachedir) for term in ('cache', 'tmp')): | |
84 | raise Exception('Not removing directory %s - this does not look like a cache dir' % cachedir) | |
85 | ||
86 | self._ydl.to_screen( | |
87 | 'Removing cache dir %s .' % cachedir, skip_eol=True) | |
88 | if os.path.exists(cachedir): | |
89 | self._ydl.to_screen('.', skip_eol=True) | |
90 | shutil.rmtree(cachedir) | |
91 | self._ydl.to_screen('.') |