Merge pull request #38 from yagebu/improvements

Improvements to the deoplete source.

1 files changed, 56 insertions, 17 deletions

diff --git a/rplugin/python3/deoplete/sources/beancount.py b/rplugin/python3/deoplete/sources/beancount.py
index c21c057..0194da3 100644
--- a/rplugin/python3/deoplete/sources/beancount.py
+++ b/rplugin/python3/deoplete/sources/beancount.py
@@ -1,16 +1,19 @@
+import collections
 import re
 
 from deoplete.source.base import Base
 
 try:
     from beancount.loader import load_file
-    from beancount.core.data import Open, Transaction
+    from beancount.core import data
     HAS_BEANCOUNT = True
 except ImportError:
     HAS_BEANCOUNT = False
 
-DIRECTIVES = ['open', 'close', 'commodity', 'txn', 'balance', 'pad', 'note',
-              'document', 'price', 'event', 'query', 'custom']
+DIRECTIVES = [
+    'open', 'close', 'commodity', 'txn', 'balance', 'pad', 'note', 'document',
+    'price', 'event', 'query', 'custom'
+]
 
 
 class Source(Base):
@@ -21,7 +24,9 @@ class Source(Base):
         self.name = 'beancount'
         self.mark = '[bc]'
         self.filetypes = ['beancount']
+        self.rank = 500
         self.min_pattern_length = 0
+        self.attributes = collections.defaultdict(list)
 
     def on_init(self, context):
         if not HAS_BEANCOUNT:
@@ -35,41 +40,75 @@ class Source(Base):
         return m.start() if m else -1
 
     def gather_candidates(self, context):
+        attrs = self.attributes
         if re.match(r'^\d{4}[/-]\d\d[/-]\d\d \w*$', context['input']):
             return [{'word': x, 'kind': 'directive'} for x in DIRECTIVES]
+        # line that starts with whitespace (-> accounts)
+        if re.match(r'^(\s)+[\w:]+$', context['input']):
+            return [{'word': x, 'kind': 'account'} for x in attrs['accounts']]
+        # directive followed by account
+        if re.search(
+                r'(balance|document|note|open|close|pad(\s[\w:]+)?)'
+                r'\s[\w:]+$',
+                context['input']):
+            return [{'word': x, 'kind': 'account'} for x in attrs['accounts']]
+        # events
+        if re.search(r'event "[^"]*$', context['input']):
+            return [{
+                'word': '"{}"'.format(x),
+                'kind': 'event'
+            } for x in attrs['events']]
+        # commodity after number
+        if re.search(r'\s([0-9]+|[0-9][0-9,]+[0-9])(\.[0-9]*)?\s\w+$',
+                     context['input']):
+            return [{
+                'word': x,
+                'kind': 'commodity'
+            } for x in attrs['commodities']]
         if not context['complete_str']:
             return []
         first = context['complete_str'][0]
         if first == '#':
-            return [{'word': '#' + w, 'kind': 'tag'} for w in self._tags]
+            return [{'word': '#' + w, 'kind': 'tag'} for w in attrs['tags']]
         elif first == '^':
-            return [{'word': '^' + w, 'kind': 'link'} for w in self._links]
+            return [{'word': '^' + w, 'kind': 'link'} for w in attrs['links']]
         elif first == '"':
-            return [{'word': '"{}"'.format(w), 'kind': 'payee'} for w in self._payees]
-        return [{'word': x, 'kind': 'account'} for x in self._accounts]
+            return [{
+                'word': '"{}"'.format(w),
+                'kind': 'payee'
+            } for w in attrs['payees']]
+        return []
 
     def __make_cache(self, context):
+        if not HAS_BEANCOUNT:
+            return
+
+        entries, _, options = load_file(self.vim.eval("beancount#get_root()"))
+
         accounts = set()
+        events = set()
         links = set()
         payees = set()
         tags = set()
-        if HAS_BEANCOUNT:
-            entries, _, _ = load_file(self.vim.eval("beancount#get_root()"))
-        else:
-            entries = []
 
         for entry in entries:
-            if isinstance(entry, Open):
+            if isinstance(entry, data.Open):
                 accounts.add(entry.account)
-            elif isinstance(entry, Transaction):
+            elif isinstance(entry, data.Transaction):
                 if entry.payee:
                     payees.add(entry.payee)
             if hasattr(entry, 'links') and entry.links:
                 links.update(entry.links)
             if hasattr(entry, 'tags') and entry.tags:
                 tags.update(entry.tags)
+            if isinstance(entry, data.Event):
+                events.add(entry.type)
 
-        self._accounts = sorted(accounts)
-        self._links = sorted(links)
-        self._payees = sorted(payees)
-        self._tags = sorted(tags)
+        self.attributes = {
+            'accounts': sorted(accounts),
+            'events': sorted(events),
+            'commodities': options['commodities'],
+            'links': sorted(links),
+            'payees': sorted(payees),
+            'tags': sorted(tags),
+        }