1 files changed, 37 insertions, 34 deletions
diff --git a/pgcli/pgtoolbar.py b/pgcli/pgtoolbar.py
index b1a3596a..2fd978c1 100644
--- a/pgcli/pgtoolbar.py
+++ b/pgcli/pgtoolbar.py
@@ -1,59 +1,62 @@
-from pygments.token import Token
-from prompt_toolkit.enums import DEFAULT_BUFFER
+from __future__ import unicode_literals
+
 from prompt_toolkit.key_binding.vi_state import InputMode
+from prompt_toolkit.application import get_app
 
 
-def _get_vi_mode(cli):
+def _get_vi_mode():
     return {
         InputMode.INSERT: 'I',
         InputMode.NAVIGATION: 'N',
         InputMode.REPLACE: 'R',
         InputMode.INSERT_MULTIPLE: 'M',
-    }[cli.vi_state.input_mode]
-
-
-def create_toolbar_tokens_func(get_vi_mode_enabled, get_is_refreshing,
-                               failed_transaction, valid_transaction):
-    """
-    Return a function that generates the toolbar tokens.
-    """
-    assert callable(get_vi_mode_enabled)
+    }[get_app().vi_state.input_mode]
 
-    token = Token.Toolbar
 
-    def get_toolbar_tokens(cli):
+def create_toolbar_tokens_func(pgcli):
+    """Return a function that generates the toolbar tokens."""
+    def get_toolbar_tokens():
         result = []
-        result.append((token, ' '))
+        result.append(('class:bottom-toolbar', ' '))
 
-        if cli.buffers[DEFAULT_BUFFER].completer.smart_completion:
-            result.append((token.On, '[F2] Smart Completion: ON  '))
+        if pgcli.completer.smart_completion:
+            result.append(('class:bottom-toolbar.on',
+                           '[F2] Smart Completion: ON  '))
         else:
-            result.append((token.Off, '[F2] Smart Completion: OFF  '))
+            result.append(('class:bottom-toolbar.off',
+                           '[F2] Smart Completion: OFF  '))
 
-        if cli.buffers[DEFAULT_BUFFER].always_multiline:
-            result.append((token.On, '[F3] Multiline: ON  '))
+        if pgcli.multi_line:
+            result.append(('class:bottom-toolbar.on', '[F3] Multiline: ON  '))
         else:
-            result.append((token.Off, '[F3] Multiline: OFF  '))
+            result.append(('class:bottom-toolbar.off',
+                           '[F3] Multiline: OFF  '))
 
-        if cli.buffers[DEFAULT_BUFFER].always_multiline:
-            if cli.buffers[DEFAULT_BUFFER].multiline_mode == 'safe':
-                result.append((token,' ([Esc] [Enter] to execute]) '))
+        if pgcli.multi_line:
+            if pgcli.multiline_mode == 'safe':
+                result.append(
+                    ('class:bottom-toolbar', ' ([Esc] [Enter] to execute]) '))
             else:
-                result.append((token,' (Semi-colon [;] will end the line) '))
+                result.append(
+                    ('class:bottom-toolbar', ' (Semi-colon [;] will end the line) '))
 
-        if get_vi_mode_enabled():
-            result.append((token.On, '[F4] Vi-mode (' + _get_vi_mode(cli) + ')'))
+        if pgcli.vi_mode:
+            result.append(
+                ('class:bottom-toolbar', '[F4] Vi-mode (' + _get_vi_mode() + ')'))
         else:
-            result.append((token.On, '[F4] Emacs-mode'))
+            result.append(('class:bottom-toolbar', '[F4] Emacs-mode'))
 
-        if failed_transaction():
-            result.append((token.Transaction.Failed, '     Failed transaction'))
+        if pgcli.pgexecute.failed_transaction():
+            result.append(('class:bottom-toolbar.transaction.failed',
+                           '     Failed transaction'))
 
-        if valid_transaction():
-            result.append((token.Transaction.Valid, '     Transaction'))
+        if pgcli.pgexecute.valid_transaction():
+            result.append(
+                ('class:bottom-toolbar.transaction.valid', '     Transaction'))
 
-        if get_is_refreshing():
-            result.append((token, '     Refreshing completions...'))
+        if pgcli.completion_refresher.is_refreshing():
+            result.append(
+                ('class:bottom-toolbar', '     Refreshing completions...'))
 
         return result
     return get_toolbar_tokens