Skip to content

Commit fd36834

Browse files
vbaltrusaitis-reefvbaltrusaitis-reef
authored
Cache argument parser for quicker autocomplete suggestions (Backblaze#965)
* Cache autocomplete suggestions * Add town crier update * Fix imports * from __future__ import annotations * Autocomplete cache improvements * Invalidate cache when VERSION changes instead of looking at the file hashes. * Use platformdirs for determining cache directory. * Move some modules to under internal _cli module. * Relax platformdirs dependency * Move autocomplete cache tests to unit * Remove unnecessary deepcopy * Skip autocomplete cache tests on Windows * Fix skipping forked tests on Windows * Run formatter * Use tmp_path instead of tmpdir * Remove unused import * Move common fixtures to /test/unit
1 parent 26a7e51 commit fd36834

16 files changed

+639
-171
lines changed

‎b2/_cli/arg_parser_types.py‎

Lines changed: 78 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,78 @@
1+
######################################################################
2+
#
3+
# File: b2/_cli/arg_parser_types.py
4+
#
5+
# Copyright 2020 Backblaze Inc. All Rights Reserved.
6+
#
7+
# License https://www.backblaze.com/using_b2_code.html
8+
#
9+
######################################################################
10+
11+
import argparse
12+
import functools
13+
import re
14+
15+
import arrow
16+
from b2sdk.v2 import RetentionPeriod
17+
18+
_arrow_version = tuple(int(p) for p in arrow.__version__.split("."))
19+
20+
21+
def parse_comma_separated_list(s):
22+
"""
23+
Parse comma-separated list.
24+
"""
25+
return [word.strip() for word in s.split(",")]
26+
27+
28+
def parse_millis_from_float_timestamp(s):
29+
"""
30+
Parse timestamp, e.g. 1367900664 or 1367900664.152
31+
"""
32+
parsed = arrow.get(float(s))
33+
if _arrow_version < (1, 0, 0):
34+
return int(parsed.format("XSSS"))
35+
else:
36+
return int(parsed.format("x")[:13])
37+
38+
39+
def parse_range(s):
40+
"""
41+
Parse optional integer range
42+
"""
43+
bytes_range = None
44+
if s is not None:
45+
bytes_range = s.split(',')
46+
if len(bytes_range) != 2:
47+
raise argparse.ArgumentTypeError('the range must have 2 values: start,end')
48+
bytes_range = (
49+
int(bytes_range[0]),
50+
int(bytes_range[1]),
51+
)
52+
53+
return bytes_range
54+
55+
56+
def parse_default_retention_period(s):
57+
unit_part = '(' + ')|('.join(RetentionPeriod.KNOWN_UNITS) + ')'
58+
m = re.match(r'^(?P<duration>\d+) (?P<unit>%s)$' % (unit_part), s)
59+
if not m:
60+
raise argparse.ArgumentTypeError(
61+
'default retention period must be in the form of "X days|years "'
62+
)
63+
return RetentionPeriod(**{m.group('unit'): int(m.group('duration'))})
64+
65+
66+
def wrap_with_argument_type_error(func, translator=str, exc_type=ValueError):
67+
"""
68+
Wrap function that may raise an exception into a function that raises ArgumentTypeError error.
69+
"""
70+
71+
@functools.wraps(func)
72+
def wrapper(*args, **kwargs):
73+
try:
74+
return func(*args, **kwargs)
75+
except exc_type as e:
76+
raise argparse.ArgumentTypeError(translator(e))
77+
78+
return wrapper

‎b2/_cli/argcompleters.py‎

Lines changed: 23 additions & 25 deletions
Original file line numberDiff line numberDiff line change
@@ -7,40 +7,32 @@
77
# License https://www.backblaze.com/using_b2_code.html
88
#
99
######################################################################
10-
from functools import wraps
11-
from itertools import islice
12-
13-
from b2sdk.v2 import LIST_FILE_NAMES_MAX_LIMIT
14-
from b2sdk.v2.api import B2Api
15-
16-
from b2._cli.b2api import _get_b2api_for_profile
17-
from b2._utils.python_compat import removeprefix
18-
from b2._utils.uri import parse_b2_uri
19-
20-
21-
def _with_api(func):
22-
"""Decorator to inject B2Api instance into argcompleter function."""
2310

24-
@wraps(func)
25-
def wrapper(prefix, parsed_args, **kwargs):
26-
api = _get_b2api_for_profile(parsed_args.profile)
27-
return func(prefix=prefix, parsed_args=parsed_args, api=api, **kwargs)
11+
# We import all the necessary modules lazily in completers in order
12+
# to avoid upfront cost of the imports when argcompleter is used for
13+
# autocompletions.
2814

29-
return wrapper
15+
from itertools import islice
3016

3117

32-
@_with_api
33-
def bucket_name_completer(api: B2Api, **kwargs):
34-
return [bucket.name for bucket in api.list_buckets(use_cache=True)]
18+
def bucket_name_completer(prefix, parsed_args, **kwargs):
19+
from b2._cli.b2api import _get_b2api_for_profile
20+
api = _get_b2api_for_profile(getattr(parsed_args, 'profile', None))
21+
res = [bucket.name for bucket in api.list_buckets(use_cache=True)]
22+
return res
3523

3624

37-
@_with_api
38-
def file_name_completer(api: B2Api, parsed_args, **kwargs):
25+
def file_name_completer(prefix, parsed_args, **kwargs):
3926
"""
4027
Completes file names in a bucket.
4128
4229
To limit delay & cost only lists files returned from by single call to b2_list_file_names
4330
"""
31+
from b2sdk.v2 import LIST_FILE_NAMES_MAX_LIMIT
32+
33+
from b2._cli.b2api import _get_b2api_for_profile
34+
35+
api = _get_b2api_for_profile(parsed_args.profile)
4436
bucket = api.get_bucket_by_name(parsed_args.bucketName)
4537
file_versions = bucket.ls(
4638
getattr(parsed_args, 'folderName', None) or '',
@@ -54,11 +46,17 @@ def file_name_completer(api: B2Api, parsed_args, **kwargs):
5446
]
5547

5648

57-
@_with_api
58-
def b2uri_file_completer(api: B2Api, prefix: str, **kwargs):
49+
def b2uri_file_completer(prefix: str, parsed_args, **kwargs):
5950
"""
6051
Complete B2 URI pointing to a file-like object in a bucket.
6152
"""
53+
from b2sdk.v2 import LIST_FILE_NAMES_MAX_LIMIT
54+
55+
from b2._cli.b2api import _get_b2api_for_profile
56+
from b2._utils.python_compat import removeprefix
57+
from b2._utils.uri import parse_b2_uri
58+
59+
api = _get_b2api_for_profile(getattr(parsed_args, 'profile', None))
6260
if prefix.startswith('b2://'):
6361
prefix_without_scheme = removeprefix(prefix, 'b2://')
6462
if '/' not in prefix_without_scheme:

‎b2/_cli/autocomplete_cache.py‎

Lines changed: 146 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,146 @@
1+
######################################################################
2+
#
3+
# File: b2/_cli/autocomplete_cache.py
4+
#
5+
# Copyright 2020 Backblaze Inc. All Rights Reserved.
6+
#
7+
# License https://www.backblaze.com/using_b2_code.html
8+
#
9+
######################################################################
10+
from __future__ import annotations
11+
12+
import abc
13+
import argparse
14+
import os
15+
import pathlib
16+
import pickle
17+
from typing import Callable
18+
19+
import argcomplete
20+
import platformdirs
21+
22+
from b2.version import VERSION
23+
24+
25+
def identity(x):
26+
return x
27+
28+
29+
class StateTracker(abc.ABC):
30+
@abc.abstractmethod
31+
def current_state_identifier(self) -> str:
32+
raise NotImplementedError()
33+
34+
35+
class PickleStore(abc.ABC):
36+
@abc.abstractmethod
37+
def get_pickle(self, identifier: str) -> bytes | None:
38+
raise NotImplementedError()
39+
40+
@abc.abstractmethod
41+
def set_pickle(self, identifier: str, data: bytes) -> None:
42+
raise NotImplementedError()
43+
44+
45+
class VersionTracker(StateTracker):
46+
def current_state_identifier(self) -> str:
47+
return VERSION
48+
49+
50+
class HomeCachePickleStore(PickleStore):
51+
_dir: pathlib.Path
52+
53+
def __init__(self, dir: pathlib.Path | None = None) -> None:
54+
self._dir = dir
55+
56+
def _cache_dir(self) -> pathlib.Path:
57+
if self._dir:
58+
return self._dir
59+
self._dir = pathlib.Path(
60+
platformdirs.user_cache_dir(appname='b2', appauthor='backblaze')
61+
) / 'autocomplete'
62+
return self._dir
63+
64+
def _fname(self, identifier: str) -> str:
65+
return f"b2-autocomplete-cache-{identifier}.pickle"
66+
67+
def get_pickle(self, identifier: str) -> bytes | None:
68+
path = self._cache_dir() / self._fname(identifier)
69+
if path.exists():
70+
with open(path, 'rb') as f:
71+
return f.read()
72+
73+
def set_pickle(self, identifier: str, data: bytes) -> None:
74+
"""Sets the pickle for identifier if it doesn't exist.
75+
When a new pickle is added, old ones are removed."""
76+
77+
dir = self._cache_dir()
78+
os.makedirs(dir, exist_ok=True)
79+
path = dir / self._fname(identifier)
80+
for file in dir.glob('b2-autocomplete-cache-*.pickle'):
81+
file.unlink()
82+
with open(path, 'wb') as f:
83+
f.write(data)
84+
85+
86+
class AutocompleteCache:
87+
_tracker: StateTracker
88+
_store: PickleStore
89+
_unpickle: Callable[[bytes], argparse.ArgumentParser]
90+
91+
def __init__(
92+
self,
93+
tracker: StateTracker,
94+
store: PickleStore,
95+
unpickle: Callable[[bytes], argparse.ArgumentParser] | None = None
96+
):
97+
self._tracker = tracker
98+
self._store = store
99+
self._unpickle = unpickle or pickle.loads
100+
101+
def _is_autocomplete_run(self) -> bool:
102+
return '_ARGCOMPLETE' in os.environ
103+
104+
def autocomplete_from_cache(self, uncached_args: dict | None = None) -> None:
105+
if not self._is_autocomplete_run():
106+
return
107+
108+
try:
109+
identifier = self._tracker.current_state_identifier()
110+
pickle_data = self._store.get_pickle(identifier)
111+
if pickle_data:
112+
parser = self._unpickle(pickle_data)
113+
argcomplete.autocomplete(parser, **(uncached_args or {}))
114+
except Exception:
115+
# Autocomplete from cache failed but maybe we can autocomplete from scratch
116+
return
117+
118+
def _clean_parser(self, parser: argparse.ArgumentParser) -> None:
119+
parser.register('type', None, identity)
120+
for action in parser._actions:
121+
if action.type not in [str, int]:
122+
action.type = None
123+
for action in parser._action_groups:
124+
for key in parser._defaults:
125+
action.set_defaults(**{key: None})
126+
parser.description = None
127+
if parser._subparsers:
128+
for group_action in parser._subparsers._group_actions:
129+
for parser in group_action.choices.values():
130+
self._clean_parser(parser)
131+
132+
def cache_and_autocomplete(
133+
self, parser: argparse.ArgumentParser, uncached_args: dict | None = None
134+
) -> None:
135+
if not self._is_autocomplete_run():
136+
return
137+
138+
try:
139+
identifier = self._tracker.current_state_identifier()
140+
self._clean_parser(parser)
141+
self._store.set_pickle(identifier, pickle.dumps(parser))
142+
finally:
143+
argcomplete.autocomplete(parser, **(uncached_args or {}))
144+
145+
146+
AUTOCOMPLETE = AutocompleteCache(tracker=VersionTracker(), store=HomeCachePickleStore())

‎b2/_cli/b2args.py‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -12,9 +12,9 @@
1212
"""
1313
import argparse
1414

15+
from b2._cli.arg_parser_types import wrap_with_argument_type_error
1516
from b2._cli.argcompleters import b2uri_file_completer
1617
from b2._utils.uri import B2URI, B2URIBase, parse_b2_uri
17-
from b2.arg_parser import wrap_with_argument_type_error
1818

1919

2020
def b2_file_uri(value: str) -> B2URIBase:

‎b2/arg_parser.py‎

Lines changed: 0 additions & 65 deletions
Original file line numberDiff line numberDiff line change
@@ -12,17 +12,12 @@
1212
import argparse
1313
import functools
1414
import locale
15-
import re
1615
import sys
1716
import textwrap
1817
import unittest.mock
1918

20-
import arrow
21-
from b2sdk.v2 import RetentionPeriod
2219
from rst2ansi import rst2ansi
2320

24-
_arrow_version = tuple(int(p) for p in arrow.__version__.split("."))
25-
2621

2722
class B2RawTextHelpFormatter(argparse.RawTextHelpFormatter):
2823
"""
@@ -152,63 +147,3 @@ def print_help(self, *args, show_all: bool = False, **kwargs):
152147
self, 'formatter_class', functools.partial(B2RawTextHelpFormatter, show_all=show_all)
153148
):
154149
super().print_help(*args, **kwargs)
155-
156-
157-
def parse_comma_separated_list(s):
158-
"""
159-
Parse comma-separated list.
160-
"""
161-
return [word.strip() for word in s.split(",")]
162-
163-
164-
def parse_millis_from_float_timestamp(s):
165-
"""
166-
Parse timestamp, e.g. 1367900664 or 1367900664.152
167-
"""
168-
parsed = arrow.get(float(s))
169-
if _arrow_version < (1, 0, 0):
170-
return int(parsed.format("XSSS"))
171-
else:
172-
return int(parsed.format("x")[:13])
173-
174-
175-
def parse_range(s):
176-
"""
177-
Parse optional integer range
178-
"""
179-
bytes_range = None
180-
if s is not None:
181-
bytes_range = s.split(',')
182-
if len(bytes_range) != 2:
183-
raise argparse.ArgumentTypeError('the range must have 2 values: start,end')
184-
bytes_range = (
185-
int(bytes_range[0]),
186-
int(bytes_range[1]),
187-
)
188-
189-
return bytes_range
190-
191-
192-
def parse_default_retention_period(s):
193-
unit_part = '(' + ')|('.join(RetentionPeriod.KNOWN_UNITS) + ')'
194-
m = re.match(r'^(?P<duration>\d+) (?P<unit>%s)$' % (unit_part), s)
195-
if not m:
196-
raise argparse.ArgumentTypeError(
197-
'default retention period must be in the form of "X days|years "'
198-
)
199-
return RetentionPeriod(**{m.group('unit'): int(m.group('duration'))})
200-
201-
202-
def wrap_with_argument_type_error(func, translator=str, exc_type=ValueError):
203-
"""
204-
Wrap function that may raise an exception into a function that raises ArgumentTypeError error.
205-
"""
206-
207-
@functools.wraps(func)
208-
def wrapper(*args, **kwargs):
209-
try:
210-
return func(*args, **kwargs)
211-
except exc_type as e:
212-
raise argparse.ArgumentTypeError(translator(e))
213-
214-
return wrapper

0 commit comments

Comments
 (0)