Add dict key completion in IPCompleter

IPython/core/completer.py

@@ -426,6 +426,50 b' def get__all__entries(obj):'
     return [w for w in words if isinstance(w, string_types)]
 
 
+def match_dict_keys(keys, prefix):
+    """Used by dict_key_matches, matching the prefix to a list of keys"""
+    if not prefix:
+        return None, [repr(k) for k in keys
+                      if isinstance(k, (string_types, bytes))]
+    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()
+
+    # TODO: support bytes in Py3k
+    matched = []
+    for key in keys:
+        try:
+            if not key.startswith(prefix_str):
+                continue
+        except (AttributeError, TypeError):
+            # Python 3+ TypeError on b'a'.startswith('a') or vice-versa
+            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 UnicodeEncodeError:
+                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' % (token_prefix, rem_repr))
+    return quote, matched
+
+
+
 class IPCompleter(Completer):
     """Extension of the completer class with IPython-specific features"""
 
@@ -538,6 +582,7 b' class IPCompleter(Completer):'
                          self.file_matches,
                          self.magic_matches,
                          self.python_func_kw_matches,
+                         self.dict_key_matches,
                          ]
 
     def all_completions(self, text):
@@ -804,6 +849,74 b' class IPCompleter(Completer):'
                     argMatches.append("%s=" %namedArg)
         return argMatches
 
+    def dict_key_matches(self, text):
+        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', [])
+
+        try:
+            regexps = self.__dict_key_regexps
+        except AttributeError:
+            dict_key_re_fmt = r'''(?x)
+            (  # match dict-referring expression wrt greedy setting
+                %s
+            )
+            \[   # open bracket
+            \s*  # and optional whitespace
+            ([uUbB]?  # string prefix (r not handled)
+                (?:   # unclosed string
+                    '(?:[^']|(?<!\\)\\')*
+                |
+                    "(?:[^"]|(?<!\\)\\")*
+                )
+            )?
+            $
+            '''
+            regexps = self.__dict_key_regexps = {
+                False: re.compile(dict_key_re_fmt % '''
+                                  # identifiers separated by .
+                                  (?!\d)\w+
+                                  (?:\.(?!\d)\w+)*
+                                  '''),
+                True: re.compile(dict_key_re_fmt % '''
+                                 .+
+                                 ''')
+            }
+
+        match = regexps[self.greedy].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
+        closing_quote, matches = match_dict_keys(keys, prefix)
+
+        # append closing quote and bracket as appropriate
+        continuation = self.line_buffer[len(self.text_until_cursor):]
+        if closing_quote and continuation.startswith(closing_quote):
+            suf = ''
+        elif continuation.startswith(']'):
+            suf = closing_quote or ''
+        else:
+            suf = (closing_quote or '') + ']'
+        return [k + suf for k in matches]
+
     def dispatch_custom_completer(self, text):
         #io.rprint("Custom! '%s' %s" % (text, self.custom_completers)) # dbg
         line = self.line_buffer

IPython/core/tests/test_completer.py

@@ -20,6 +20,7 b' from IPython.utils.tempdir import TemporaryDirectory'
 from IPython.utils.generics import complete_object
 from IPython.utils import py3compat
 from IPython.utils.py3compat import string_types, unicode_type
+from IPython.testing import decorators as dec
 
 #-----------------------------------------------------------------------------
 # Test functions
@@ -392,3 +393,206 b' def test_magic_completion_order():'
     text, matches = c.complete('timeit')
     nt.assert_equal(matches, ["timeit", "%timeit","%%timeit"])
 
+
+
+def test_dict_key_completion_string():
+    """Test dictionary key completion for string keys"""
+    ip = get_ipython()
+    complete = ip.Completer.complete
+
+    ip.user_ns['d'] = {'abc': None}
+
+    # check completion at different stages
+    _, matches = complete(line_buffer="d[")
+    nt.assert_in("'abc']", matches)
+
+    _, matches = complete(line_buffer="d['")
+    nt.assert_in("abc']", matches)
+
+    _, matches = complete(line_buffer="d['a")
+    nt.assert_in("abc']", matches)
+
+    # check use of different quoting
+    _, matches = complete(line_buffer="d[\"")
+    nt.assert_in("abc\"]", matches)
+
+    _, matches = complete(line_buffer="d[\"a")
+    nt.assert_in("abc\"]", matches)
+
+    # check sensitivity to following context
+    _, matches = complete(line_buffer="d[]", cursor_pos=2)
+    nt.assert_in("'abc'", matches)
+
+    _, matches = complete(line_buffer="d['']", cursor_pos=3)
+    nt.assert_in("abc", matches)
+
+    # check multiple solutions are correctly returned and that noise is not
+    ip.user_ns['d'] = {'abc': None, 'abd': None, 'bad': None, object(): None,
+                       5: None}
+
+    _, matches = complete(line_buffer="d['a")
+    nt.assert_in("abc']", matches)
+    nt.assert_in("abd']", matches)
+    nt.assert_not_in("bad']", matches)
+
+    # check escaping and whitespace
+    ip.user_ns['d'] = {'a\nb': None, 'a\'b': None, 'a"b': None, 'a word': None}
+    _, matches = complete(line_buffer="d['a")
+    nt.assert_in("a\\nb']", matches)
+    nt.assert_in("a\\'b']", matches)
+    nt.assert_in("a\"b']", matches)
+    nt.assert_in("a word']", matches)
+
+    # - can complete on non-initial word of the string
+    _, matches = complete(line_buffer="d['a w")
+    nt.assert_in("word']", matches)
+
+    # - understands quote escaping
+    _, matches = complete(line_buffer="d['a\\'")
+    nt.assert_in("b']", matches)
+
+    # - default quoting should work like repr
+    _, matches = complete(line_buffer="d[")
+    nt.assert_in("\"a'b\"]", matches)
+
+    # - when opening quote with ", possible to match with unescaped apostrophe
+    _, matches = complete(line_buffer="d[\"a'")
+    nt.assert_in("b\"]", matches)
+
+
+def test_dict_key_completion_contexts():
+    """Test expression contexts in which dict key completion occurs"""
+    ip = get_ipython()
+    complete = ip.Completer.complete
+    d = {'abc': None}
+    ip.user_ns['d'] = d
+
+    class C:
+        data = d
+    ip.user_ns['C'] = C
+    ip.user_ns['get'] = lambda: d
+
+    def assert_no_completion(**kwargs):
+        _, matches = complete(**kwargs)
+        nt.assert_not_in('abc', matches)
+        nt.assert_not_in('abc\'', matches)
+        nt.assert_not_in('abc\']', matches)
+        nt.assert_not_in('\'abc\'', matches)
+        nt.assert_not_in('\'abc\']', matches)
+
+    def assert_completion(**kwargs):
+        _, matches = complete(**kwargs)
+        nt.assert_in("'abc']", matches)
+
+    # no completion after string closed, even if reopened
+    ip.Completer.greedy = False
+    assert_no_completion(line_buffer="d['a'")
+    assert_no_completion(line_buffer="d[\"a\"")
+    assert_no_completion(line_buffer="d['a' + ")
+    assert_no_completion(line_buffer="d['a' + '")
+
+    # completion in non-trivial expressions
+    assert_completion(line_buffer="+ d[")
+    assert_completion(line_buffer="(d[")
+    assert_completion(line_buffer="C.data[")
+
+    # greedy flag
+    assert_no_completion(line_buffer="get()[")
+    ip.Completer.greedy = True
+    assert_completion(line_buffer="get()[")
+
+
+
[e-mail](sys.version_info[0] >= 3, 'This test only applies in Py>=3')
+def test_dict_key_completion_bytes():
+    """Test handling of bytes in dict key completion"""
+    ip = get_ipython()
+    complete = ip.Completer.complete
+
+    ip.user_ns['d'] = {'abc': None, b'abd': None}
+
+    _, matches = complete(line_buffer="d[")
+    nt.assert_in("'abc']", matches)
+    nt.assert_in("b'abd']", matches)
+
+    if False:  # not currently implemented
+        _, matches = complete(line_buffer="d[b")
+        nt.assert_in("b'abd']", matches)
+        nt.assert_not_in("b'abc']", matches)
+
+        _, matches = complete(line_buffer="d[b'")
+        nt.assert_in("abd']", matches)
+        nt.assert_not_in("abc']", matches)
+
+        _, matches = complete(line_buffer="d[B'")
+        nt.assert_in("abd']", matches)
+        nt.assert_not_in("abc']", matches)
+
+        _, matches = complete(line_buffer="d['")
+        nt.assert_in("abc']", matches)
+        nt.assert_not_in("abd']", matches)
+
+
[e-mail](sys.version_info[0] < 3, 'This test only applies in Py<3')
+def test_dict_key_completion_unicode():
+    """Test handling of unicode in dict key completion"""
+    ip = get_ipython()
+    complete = ip.Completer.complete
+
+    ip.user_ns['d'] = {unicode_type('abc'): None,
+                       unicode_type('a\xd7\x90', 'utf8'): None}
+
+    _, matches = complete(line_buffer="d[")
+    nt.assert_in("u'abc']", matches)
+    nt.assert_in("u'a\\u05d0']", matches)
+
+    _, matches = complete(line_buffer="d['a")
+    nt.assert_in("abc']", matches)
+    nt.assert_not_in("a\\u05d0']", matches)
+
+    _, matches = complete(line_buffer="d[u'a")
+    nt.assert_in("abc']", matches)
+    nt.assert_in("a\\u05d0']", matches)
+
+    _, matches = complete(line_buffer="d[U'a")
+    nt.assert_in("abc']", matches)
+    nt.assert_in("a\\u05d0']", matches)
+
+
+def test_dict_like_key_completion():
+    """Test dict key completion applies where __getitem__ and keys exist"""
+    class D(object):
+        def __getitem__(self):
+            pass
+        def keys(self):
+            return iter(['hello', 'world'])
+    ip = get_ipython()
+    complete = ip.Completer.complete
+    ip.user_ns['d'] = D()
+    _, matches = complete(line_buffer="d['")
+    nt.assert_in("hello']", matches)
+    nt.assert_in("world']", matches)
+
+
[e-mail]_without('numpy')
+def test_struct_array_key_completion():
+    """Test dict key completion applies to numpy struct arrays"""
+    import numpy
+    ip = get_ipython()
+    complete = ip.Completer.complete
+    ip.user_ns['d'] = numpy.array([], dtype=[('hello', 'f'), ('world', 'f')])
+    _, matches = complete(line_buffer="d['")
+    nt.assert_in("hello']", matches)
+    nt.assert_in("world']", matches)
+
+
[e-mail]_without('pandas')
+def test_dataframe_key_completion():
+    """Test dict key completion applies to pandas DataFrames"""
+    import pandas
+    ip = get_ipython()
+    complete = ip.Completer.complete
+    ip.user_ns['d'] = pandas.DataFrame({'hello': [1], 'world': [2]})
+    _, matches = complete(line_buffer="d['")
+    nt.assert_in("hello']", matches)
+    nt.assert_in("world']", matches)