Created
June 6, 2013 23:13
-
-
Save jnothman/5725776 to your computer and use it in GitHub Desktop.
Dictionary key completion for IPython
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
from __future__ import print_function | |
import re | |
import sys | |
import six | |
dict_key_re = re.compile(r'''(?x) | |
( # preceding string of '.'-delimited identifiers | |
(?!\d)\w+ | |
(?:\.(?!\d)\w+)* | |
) | |
\[ # open bracket | |
\s* # and optional whitespace | |
([uUbB]? # string prefix (r not handled) | |
(?: # unclosed string | |
'(?:[^']|(?<!\\)\\')* | |
| | |
"(?:[^"]|(?<!\\)\\")* | |
) | |
)? | |
$ | |
''') | |
assert dict_key_re.search('d[').groups() == ('d', None) | |
assert dict_key_re.search('d["').groups() == ('d', '"') | |
assert dict_key_re.search('d["a').groups() == ('d', '"a') | |
assert dict_key_re.search('d["\\"').groups() == ('d', '"\\"') | |
assert dict_key_re.search('c.d["a').groups() == ('c.d', '"a') | |
assert dict_key_re.search('+ c.d["a').groups() == ('c.d', '"a') | |
assert dict_key_re.search('["a') is None | |
assert dict_key_re.search('d["a"') is None | |
assert dict_key_re.search('d["a" +') is None | |
assert dict_key_re.search('d["\\\\"') is None | |
def match_keys(keys, prefix): | |
if not prefix: | |
return [repr(k) for k in keys if isinstance(k, six.string_types)] | |
quote_match = re.search('["\']', prefix) | |
quote = quote_match.group() | |
try: | |
prefix_str = eval(prefix + quote, {}) | |
except Exception: | |
return None | |
token_prefix = re.search('\w*$', prefix).group() | |
matched = [] | |
for key in keys: | |
try: | |
if not key.startswith(prefix_str): | |
continue | |
except (AttributeError, TypeError): | |
# Python 3+ TypeError on b'a'.startswith('a') | |
continue | |
# reformat remainder of key to begin with prefix | |
rem = key[len(prefix_str):] | |
# force repr with ' | |
rem_repr = repr(rem + '"') | |
if rem_repr.startswith('u') and prefix[0] not in 'uU': | |
try: | |
rem_repr = repr(rem.encode('ascii') + '"') | |
except UnicodeDecodeError: | |
continue | |
rem_repr = rem_repr[1 + rem_repr.index("'"):-2] | |
if quote == '"': | |
rem_repr = rem_repr.replace('"', '\\"') | |
# then reinsert prefix from start of token | |
matched.append('%s%s%s' % (token_prefix, rem_repr, quote)) | |
return matched | |
TEST_KEYS = [5, object(), 'abc', u'abd', "a'b", 'a"b', 'a"\'b', 'a\nb', u'a\u05d0', b'abc', 'a b'] | |
TEST_PREFIXES = [None, "'", '"', "u'", 'u"', "'a", "'b"] # "'a ' | |
# Doing all pairs is too hard to document, especially with py2-3 variations on u'' and b'' | |
# test logically rather than exhaustively. | |
# use if sys.version_info >= (3, 0) | |
def get_keys(obj): | |
if not callable(getattr(obj, '__getitem__', None)): | |
return [] | |
if hasattr(obj, 'keys'): | |
try: | |
return list(obj.keys()) | |
except Exception: | |
return [] | |
return getattr(getattr(obj, 'dtype', None), 'names', []) | |
def dict_key_completer(self): | |
match = dict_key_re.search(self.text_until_cursor) | |
if match is None: | |
return [] | |
expr, prefix = match.groups() | |
try: | |
obj = eval(expr, self.namespace) | |
except Exception: | |
try: | |
obj = eval(expr, self.global_namespace) | |
except Exception: | |
return [] | |
keys = get_keys(obj) | |
if not keys: | |
return keys | |
return ['%s]' % k for k in match_keys(keys, prefix)] |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment