(WIP) Changes to improve data.table completion

pythonx/ncm_r.py

@@ -187,7 +187,6 @@ def get_matches(self, word, pkg=None, pipe=None, data=None):
         :returns: list of ncm matches
         """
 
-        self.get_all_obj_matches()
         obj_m = self._obj_matches
 
         if pipe or data:
@@ -281,8 +280,12 @@ def on_complete(self, ctx):
         if isinquot and func and not re.search('(library|require|data)', func):
             return
 
+        self.get_all_obj_matches()
         pipe = rlang.get_pipe(cur_buffer, lnum, col)
-        data = rlang.get_df_inside_brackets(ctx['typed'])
+        data = rlang.get_df_inside_brackets(ctx['typed'], obj_m=self._obj_matches)
+
+        if pipe and data == '.':
+            data = pipe
 
         self._info('word: "{}", func: "{}", pkg: {}, pipe: {}, data: {}'.format(
             word, func, pkg, pipe, data))
@@ -297,7 +300,7 @@ def on_complete(self, ctx):
 
             matches = self.get_matches(word, pkg=pkg)
 
-        self.complete(ctx, ctx['startccol'], matches)
+        self.complete(ctx, ctx['startccol'], matches, refresh=1)
 
 
 SOURCE = Source(vim)

pythonx/rlang.py

@@ -159,16 +159,35 @@ def get_option(typed=''):
     return None
 
 
-def get_df_inside_brackets(typed=''):
+def get_df_inside_brackets(typed='', obj_m=[]):
     """Return df name when cursor is inside brackets"""
 
     if not typed:
         return ''
 
-    df_brackets = re.compile(r'(\w+)\[[^\[\]]*,[^\[\]]*$')
+    df_brackets = re.compile(r'([\w.]+)\[[^\[\]]*$')
+
     df_match = df_brackets.search(typed)
 
     if df_match:
-        return df_match.group(1)
+        full_match = df_match.group(0)
+        word = df_match.group(1)
+
+        if ',' in full_match:
+            return word
+
+        # Check if the object is a data.table or '.' (indicating it's referring
+        # to a data.frame or data.table being used in a pipe), if so then
+        # return it
+        if word == '.':
+            return word
+
+        try:
+            match_struct = [m.get('struct', '') for m in obj_m if m.get('word', '') == word][0]
+        except IndexError:
+            match_struct = ''
+
+        if match_struct == 'data.table':
+            return word
 
     return ''

test/test_ncmr.py

@@ -222,6 +222,23 @@ def feedkeys(keys):
 feedkeys(['$i'])
 TEST.ask()
 
+# ==== DATA.TABLE ==== #
+send_rcmd("library('data.table'); sleep_dt <- as.data.table(sleep)")
+TEST = TestCase('Is ncm-R suggesting median?',
+                ['sleep_dt[, medi'])
+NVIM.feedkeys('A')
+TEST.ask()
+
+TEST = TestCase('Is ncm-R suggesting the sleep variable "extra"?',
+                ['sleep_dt[ext'])
+NVIM.feedkeys('A')
+TEST.ask()
+
+TEST = TestCase('Is ncm-R suggesting the sleep variable "extra"?',
+                ['sleep_dt %>%', '  .[ext'])
+NVIM.feedkeys(DOWN + 'A')
+TEST.ask()
+
 # ==== IT'S  OVER ==== #
 TEST = TestCase(r'Testing is over \o/')
 TEST.ask()