Extracted handling of the editor command into method.

1 files changed, 27 insertions, 17 deletions

diff --git a/pgcli/main.py b/pgcli/main.py
index 35664cdd..77a3fa08 100755
--- a/pgcli/main.py
+++ b/pgcli/main.py
@@ -161,6 +161,32 @@ class PGCli(object):
 
         self.pgexecute = pgexecute
 
+    def handle_editor_command(self, cli, document):
+        """
+        Editor command is any query that is prefixed or suffixed
+        by a '\e'. The reason for a while loop is because a user
+        might edit a query multiple times.
+        For eg:
+        "select * from \e"<enter> to edit it in vim, then come
+        back to the prompt with the edited query "select * from
+        blah where q = 'abc'\e" to edit it again.
+        :param cli: CommandLineInterface
+        :param document: Document
+        :return: Document
+        """
+        while iospecial.editor_command(document.text):
+            filename = iospecial.get_filename(document.text)
+            sql, message = iospecial.open_external_editor(filename,
+                                                          sql=document.text)
+            if message:
+                # Something went wrong. Raise an exception and bail.
+                raise RuntimeError(message)
+            document = cli.read_input(
+                initial_document=Document(sql, cursor_position=len(sql)),
+                on_exit=AbortAction.RAISE_EXCEPTION)
+            continue
+        return document
+
     def run_cli(self):
         pgexecute = self.pgexecute
         prompt = '%s> ' % pgexecute.dbname
@@ -198,23 +224,7 @@ class PGCli(object):
                     raise Exit
 
                 try:
-                    # Editor command is any query that is prefixed or suffixed
-                    # by a '\e'. The reason for a while loop is because a user
-                    # might edit a query multiple times.
-                    # For eg:
-                    # "select * from \e"<enter> to edit it in vim, then come
-                    # back to the prompt with the edited query "select * from
-                    # blah where q = 'abc'\e" to edit it again.
-                    while iospecial.editor_command(document.text):
-                        filename = iospecial.get_filename(document.text)
-                        sql, message = iospecial.open_external_editor(filename, sql=document.text)
-                        if message:
-                            # Something went wrong. Raise an exception and bail.
-                            raise RuntimeError(message)
-                        document = cli.read_input(
-                                initial_document=Document(sql, cursor_position=len(sql)),
-                                on_exit=AbortAction.RAISE_EXCEPTION)
-                        continue
+                    document = self.handle_editor_command(cli, document)
                 except RuntimeError as e:
                     logger.error("sql: %r, error: %r", document.text, e)
                     logger.error("traceback: %r", traceback.format_exc())