Merge pull request #137 from darikg/api

Minor PGCli API improvements

pgcli/main.py

@@ -8,6 +8,7 @@ import logging
 import click
 
 from prompt_toolkit import CommandLineInterface, AbortAction, Exit
+from prompt_toolkit.document import Document
 from prompt_toolkit.layout import Layout
 from prompt_toolkit.layout.prompt import DefaultPrompt
 from prompt_toolkit.layout.menus import CompletionsMenu
@@ -56,7 +57,6 @@ class PGCli(object):
 
         # Load config.
         c = self.config = load_config('~/.pgclirc', default_config)
-        self.smart_completion = c.getboolean('main', 'smart_completion')
         self.multi_line = c.getboolean('main', 'multi_line')
 
         self.logger = logging.getLogger(__name__)
@@ -64,6 +64,13 @@ class PGCli(object):
 
         self.query_history = []
 
+        # Initialize completer
+        smart_completion = c.getboolean('main', 'smart_completion')
+        completer = PGCompleter(smart_completion)
+        completer.extend_special_commands(CASE_SENSITIVE_COMMANDS.keys())
+        completer.extend_special_commands(NON_CASE_SENSITIVE_COMMANDS.keys())
+        self.completer = completer
+
     def initialize_logging(self):
 
         log_file = self.config.get('main', 'log_file')
@@ -152,10 +159,9 @@ class PGCli(object):
             menus=[CompletionsMenu(max_height=10)],
             lexer=SqlLexer, bottom_toolbars=[PGToolbar()])
 
-        completer = PGCompleter(self.smart_completion)
-        completer.extend_special_commands(CASE_SENSITIVE_COMMANDS.keys())
-        completer.extend_special_commands(NON_CASE_SENSITIVE_COMMANDS.keys())
-        refresh_completions(pgexecute, completer)
+        completer = self.completer
+        self.refresh_completions()
+
         buf = PGBuffer(always_multiline=self.multi_line, completer=completer,
                 history=FileHistory(os.path.expanduser('~/.pgcli-history')))
         cli = CommandLineInterface(style=PGStyle, layout=layout, buffer=buf,
@@ -231,8 +237,7 @@ class PGCli(object):
                 # Refresh the table names and column names if necessary.
                 if need_completion_refresh(document.text):
                     prompt = '%s> ' % pgexecute.dbname
-                    completer.reset_completions()
-                    refresh_completions(pgexecute, completer)
+                    self.refresh_completions()
 
                 query = Query(document.text, successful, mutating)
                 self.query_history.append(query)
@@ -256,6 +261,19 @@ class PGCli(object):
 
         return less_opts
 
+    def refresh_completions(self):
+        self.completer.reset_completions()
+        tables, columns = self.pgexecute.tables()
+        self.completer.extend_table_names(tables)
+        for table in tables:
+            table = table[1:-1] if table[0] == '"' and table[-1] == '"' else table
+            self.completer.extend_column_names(table, columns[table])
+        self.completer.extend_database_names(self.pgexecute.databases())
+
+    def get_completions(self, text, cursor_positition):
+        return self.completer.get_completions(
+            Document(text=text, cursor_position=cursor_positition), None)
+
 @click.command()
 # Default host is '' so psycopg2 can default to either localhost or unix socket
 @click.option('-h', '--host', default='', envvar='PGHOST',
@@ -331,13 +349,6 @@ def quit_command(sql):
             or sql.strip() == '\q'
             or sql.strip() == ':q')
 
-def refresh_completions(pgexecute, completer):
-    tables, columns = pgexecute.tables()
-    completer.extend_table_names(tables)
-    for table in tables:
-        table = table[1:-1] if table[0] == '"' and table[-1] == '"' else table
-        completer.extend_column_names(table, columns[table])
-    completer.extend_database_names(pgexecute.databases())
 
 if __name__ == "__main__":
     cli()