Repository navigation
Expand file tree
/
Copy path_module_completer.py
More file actions
388 lines (331 loc) · 13.3 KB
/
Copy path_module_completer.py
File metadata and controls
388 lines (331 loc) · 13.3 KB
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
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
from __future__ import annotations
import pkgutil
import sys
import token
import tokenize
from io import StringIO
from contextlib import contextmanager
from dataclasses import dataclass
from itertools import chain
from tokenize import TokenInfo
TYPE_CHECKING = False
if TYPE_CHECKING:
from typing import Any, Iterable, Iterator, Mapping
def make_default_module_completer() -> ModuleCompleter:
# Inside pyrepl, __package__ is set to None by default
return ModuleCompleter(namespace={'__package__': None})
class ModuleCompleter:
"""A completer for Python import statements.
Examples:
- import <tab>
- import foo<tab>
- import foo.<tab>
- import foo as bar, baz<tab>
- from <tab>
- from foo<tab>
- from foo import <tab>
- from foo import bar<tab>
- from foo import (bar as baz, qux<tab>
"""
def __init__(self, namespace: Mapping[str, Any] | None = None) -> None:
self.namespace = namespace or {}
self._global_cache: list[pkgutil.ModuleInfo] = []
self._curr_sys_path: list[str] = sys.path[:]
def get_completions(self, line: str) -> list[str] | None:
"""Return the next possible import completions for 'line'."""
result = ImportParser(line).parse()
if not result:
return None
try:
return self.complete(*result)
except Exception:
# Some unexpected error occurred, make it look like
# no completions are available
return []
def complete(self, from_name: str | None, name: str | None) -> list[str]:
if from_name is None:
# import x.y.z<tab>
assert name is not None
path, prefix = self.get_path_and_prefix(name)
modules = self.find_modules(path, prefix)
return [self.format_completion(path, module) for module in modules]
if name is None:
# from x.y.z<tab>
path, prefix = self.get_path_and_prefix(from_name)
modules = self.find_modules(path, prefix)
return [self.format_completion(path, module) for module in modules]
# from x.y import z<tab>
return self.find_modules(from_name, name)
def find_modules(self, path: str, prefix: str) -> list[str]:
"""Find all modules under 'path' that start with 'prefix'."""
modules = self._find_modules(path, prefix)
# Filter out invalid module names
# (for example those containing dashes that cannot be imported with 'import')
return [mod for mod in modules if mod.isidentifier()]
def _find_modules(self, path: str, prefix: str) -> list[str]:
if not path:
# Top-level import (e.g. `import foo<tab>`` or `from foo<tab>`)`
builtin_modules = [name for name in sys.builtin_module_names
if self.is_suggestion_match(name, prefix)]
third_party_modules = [module.name for module in self.global_cache
if self.is_suggestion_match(module.name, prefix)]
return sorted(builtin_modules + third_party_modules)
if path.startswith('.'):
# Convert relative path to absolute path
package = self.namespace.get('__package__', '')
path = self.resolve_relative_name(path, package) # type: ignore[assignment]
if path is None:
return []
modules: Iterable[pkgutil.ModuleInfo] = self.global_cache
for segment in path.split('.'):
modules = [mod_info for mod_info in modules
if mod_info.ispkg and mod_info.name == segment]
modules = self.iter_submodules(modules)
return [module.name for module in modules
if self.is_suggestion_match(module.name, prefix)]
def is_suggestion_match(self, module_name: str, prefix: str) -> bool:
if prefix:
return module_name.startswith(prefix)
# For consistency with attribute completion, which
# does not suggest private attributes unless requested.
return not module_name.startswith("_")
def iter_submodules(self, parent_modules: list[pkgutil.ModuleInfo]) -> Iterator[pkgutil.ModuleInfo]:
"""Iterate over all submodules of the given parent modules."""
specs = [info.module_finder.find_spec(info.name, None)
for info in parent_modules if info.ispkg]
search_locations = set(chain.from_iterable(
getattr(spec, 'submodule_search_locations', [])
for spec in specs if spec
))
return pkgutil.iter_modules(search_locations)
def get_path_and_prefix(self, dotted_name: str) -> tuple[str, str]:
"""
Split a dotted name into an import path and a
final prefix that is to be completed.
Examples:
'foo.bar' -> 'foo', 'bar'
'foo.' -> 'foo', ''
'.foo' -> '.', 'foo'
"""
if '.' not in dotted_name:
return '', dotted_name
if dotted_name.startswith('.'):
stripped = dotted_name.lstrip('.')
dots = '.' * (len(dotted_name) - len(stripped))
if '.' not in stripped:
return dots, stripped
path, prefix = stripped.rsplit('.', 1)
return dots + path, prefix
path, prefix = dotted_name.rsplit('.', 1)
return path, prefix
def format_completion(self, path: str, module: str) -> str:
if path == '' or path.endswith('.'):
return f'{path}{module}'
return f'{path}.{module}'
def resolve_relative_name(self, name: str, package: str) -> str | None:
"""Resolve a relative module name to an absolute name.
Example: resolve_relative_name('.foo', 'bar') -> 'bar.foo'
"""
# taken from importlib._bootstrap
level = 0
for character in name:
if character != '.':
break
level += 1
bits = package.rsplit('.', level - 1)
if len(bits) < level:
return None
base = bits[0]
name = name[level:]
return f'{base}.{name}' if name else base
@property
def global_cache(self) -> list[pkgutil.ModuleInfo]:
"""Global module cache"""
if not self._global_cache or self._curr_sys_path != sys.path:
self._curr_sys_path = sys.path[:]
# print('getting packages')
self._global_cache = list(pkgutil.iter_modules())
return self._global_cache
class ImportParser:
"""
Parses incomplete import statements that are
suitable for autocomplete suggestions.
Examples:
- import foo -> Result(from_name=None, name='foo')
- import foo. -> Result(from_name=None, name='foo.')
- from foo -> Result(from_name='foo', name=None)
- from foo import bar -> Result(from_name='foo', name='bar')
- from .foo import ( -> Result(from_name='.foo', name='')
Note that the parser works in reverse order, starting from the
last token in the input string. This makes the parser more robust
when parsing multiple statements.
"""
_ignored_tokens = {
token.INDENT, token.DEDENT, token.COMMENT,
token.NL, token.NEWLINE, token.ENDMARKER
}
_keywords = {'import', 'from', 'as'}
def __init__(self, code: str) -> None:
self.code = code
tokens = []
try:
for t in tokenize.generate_tokens(StringIO(code).readline):
if t.type not in self._ignored_tokens:
tokens.append(t)
except tokenize.TokenError as e:
if 'unexpected EOF' not in str(e):
# unexpected EOF is fine, since we're parsing an
# incomplete statement, but other errors are not
# because we may not have all the tokens so it's
# safer to bail out
tokens = []
except SyntaxError:
tokens = []
self.tokens = TokenQueue(tokens[::-1])
def parse(self) -> tuple[str | None, str | None] | None:
if not (res := self._parse()):
return None
return res.from_name, res.name
def _parse(self) -> Result | None:
with self.tokens.save_state():
return self.parse_from_import()
with self.tokens.save_state():
return self.parse_import()
def parse_import(self) -> Result:
if self.code.rstrip().endswith('import') and self.code.endswith(' '):
return Result(name='')
if self.tokens.peek_string(','):
name = ''
else:
if self.code.endswith(' '):
raise ParseError('parse_import')
name = self.parse_dotted_name()
if name.startswith('.'):
raise ParseError('parse_import')
while self.tokens.peek_string(','):
self.tokens.pop()
self.parse_dotted_as_name()
if self.tokens.peek_string('import'):
return Result(name=name)
raise ParseError('parse_import')
def parse_from_import(self) -> Result:
stripped = self.code.rstrip()
if stripped.endswith('import') and self.code.endswith(' '):
return Result(from_name=self.parse_empty_from_import(), name='')
if stripped.endswith('from') and self.code.endswith(' '):
return Result(from_name='')
if self.tokens.peek_string('(') or self.tokens.peek_string(','):
return Result(from_name=self.parse_empty_from_import(), name='')
if self.code.endswith(' '):
raise ParseError('parse_from_import')
name = self.parse_dotted_name()
if '.' in name:
self.tokens.pop_string('from')
return Result(from_name=name)
if self.tokens.peek_string('from'):
return Result(from_name=name)
from_name = self.parse_empty_from_import()
return Result(from_name=from_name, name=name)
def parse_empty_from_import(self) -> str:
if self.tokens.peek_string(','):
self.tokens.pop()
self.parse_as_names()
if self.tokens.peek_string('('):
self.tokens.pop()
self.tokens.pop_string('import')
return self.parse_from()
def parse_from(self) -> str:
from_name = self.parse_dotted_name()
self.tokens.pop_string('from')
return from_name
def parse_dotted_as_name(self) -> str:
self.tokens.pop_name()
if self.tokens.peek_string('as'):
self.tokens.pop()
with self.tokens.save_state():
return self.parse_dotted_name()
def parse_dotted_name(self) -> str:
name = []
if self.tokens.peek_string('.'):
name.append('.')
self.tokens.pop()
if (self.tokens.peek_name()
and (tok := self.tokens.peek())
and tok.string not in self._keywords):
name.append(self.tokens.pop_name())
if not name:
raise ParseError('parse_dotted_name')
while self.tokens.peek_string('.'):
name.append('.')
self.tokens.pop()
if (self.tokens.peek_name()
and (tok := self.tokens.peek())
and tok.string not in self._keywords):
name.append(self.tokens.pop_name())
else:
break
while self.tokens.peek_string('.'):
name.append('.')
self.tokens.pop()
return ''.join(name[::-1])
def parse_as_names(self) -> None:
self.parse_as_name()
while self.tokens.peek_string(','):
self.tokens.pop()
self.parse_as_name()
def parse_as_name(self) -> None:
self.tokens.pop_name()
if self.tokens.peek_string('as'):
self.tokens.pop()
self.tokens.pop_name()
class ParseError(Exception):
pass
@dataclass(frozen=True)
class Result:
from_name: str | None = None
name: str | None = None
class TokenQueue:
"""Provides helper functions for working with a sequence of tokens."""
def __init__(self, tokens: list[TokenInfo]) -> None:
self.tokens: list[TokenInfo] = tokens
self.index: int = 0
self.stack: list[int] = []
@contextmanager
def save_state(self) -> Any:
try:
self.stack.append(self.index)
yield
except ParseError:
self.index = self.stack.pop()
else:
self.stack.pop()
def __bool__(self) -> bool:
return self.index < len(self.tokens)
def peek(self) -> TokenInfo | None:
if not self:
return None
return self.tokens[self.index]
def peek_name(self) -> bool:
if not (tok := self.peek()):
return False
return tok.type == token.NAME
def pop_name(self) -> str:
tok = self.pop()
if tok.type != token.NAME:
raise ParseError('pop_name')
return tok.string
def peek_string(self, string: str) -> bool:
if not (tok := self.peek()):
return False
return tok.string == string
def pop_string(self, string: str) -> str:
tok = self.pop()
if tok.string != string:
raise ParseError('pop_string')
return tok.string
def pop(self) -> TokenInfo:
if not self:
raise ParseError('pop')
tok = self.tokens[self.index]
self.index += 1
return tok