1 files changed, 104 insertions, 14 deletions
diff --git a/mycli/packages/special/iocommands.py b/mycli/packages/special/iocommands.py
index 11dca8d..58066b8 100644
--- a/mycli/packages/special/iocommands.py
+++ b/mycli/packages/special/iocommands.py
@@ -8,6 +8,7 @@ from io import open
 from time import sleep
 
 import click
+import pyperclip
 import sqlparse
 
 from . import export
@@ -23,6 +24,8 @@ PAGER_ENABLED = True
 tee_file = None
 once_file = None
 written_to_once_file = False
+pipe_once_process = None
+written_to_pipe_once_process = False
 delimiter_command = DelimiterCommand()
 
 
@@ -115,7 +118,7 @@ def get_editor_query(sql):
     # The reason we can't simply do .strip('\e') is that it strips characters,
     # not a substring. So it'll strip "e" in the end of the sql also!
     # Ex: "select * from style\e" -> "select * from styl".
-    pattern = re.compile('(^\\\e|\\\e$)')
+    pattern = re.compile(r'(^\\e|\\e$)')
     while pattern.search(sql):
         sql = pattern.sub('', sql)
 
@@ -159,6 +162,47 @@ def open_external_editor(filename=None, sql=None):
     return (query, message)
 
 
+@export
+def clip_command(command):
+    """Is this a clip command?
+
+    :param command: string
+
+    """
+    # It is possible to have `\clip` or `SELECT * FROM \clip`. So we check
+    # for both conditions.
+    return command.strip().endswith('\\clip') or command.strip().startswith('\\clip')
+
+
+@export
+def get_clip_query(sql):
+    """Get the query part of a clip command."""
+    sql = sql.strip()
+
+    # The reason we can't simply do .strip('\clip') is that it strips characters,
+    # not a substring. So it'll strip "c" in the end of the sql also!
+    pattern = re.compile(r'(^\\clip|\\clip$)')
+    while pattern.search(sql):
+        sql = pattern.sub('', sql)
+
+    return sql
+
+
+@export
+def copy_query_to_clipboard(sql=None):
+    """Send query to the clipboard."""
+
+    sql = sql or ''
+    message = None
+
+    try:
+        pyperclip.copy(u'{sql}'.format(sql=sql))
+    except RuntimeError as e:
+        message = 'Error clipping query: %s.' % e.strerror
+
+    return message
+
+
 @special_command('\\f', '\\f [name [args..]]', 'List or execute favorite queries.', arg_type=PARSED_QUERY, case_sensitive=True)
 def execute_favorite_query(cur, arg, **_):
     """Returns (title, rows, headers, status)"""
@@ -213,7 +257,7 @@ def subst_favorite_query_args(query, args):
 
         query = query.replace(subst_var, val)
 
-    match = re.search('\\$\d+', query)
+    match = re.search(r'\$\d+', query)
     if match:
         return[None, 'missing substitution for ' + match.group(0) + ' in query:\n  ' + query]
 
@@ -337,7 +381,11 @@ def write_tee(output):
 def set_once(arg, **_):
     global once_file, written_to_once_file
 
-    once_file = parseargfile(arg)
+    try:
+        once_file = open(**parseargfile(arg))
+    except (IOError, OSError) as e:
+        raise OSError("Cannot write to file '{}': {}".format(
+            e.filename, e.strerror))
     written_to_once_file = False
 
     return [(None, None, None, "")]
@@ -347,26 +395,68 @@ def set_once(arg, **_):
 def write_once(output):
     global once_file, written_to_once_file
     if output and once_file:
-        try:
-            f = open(**once_file)
-        except (IOError, OSError) as e:
-            once_file = None
-            raise OSError("Cannot write to file '{}': {}".format(
-                e.filename, e.strerror))
-        with f:
-            click.echo(output, file=f, nl=False)
-            click.echo(u"\n", file=f, nl=False)
+        click.echo(output, file=once_file, nl=False)
+        click.echo(u"\n", file=once_file, nl=False)
+        once_file.flush()
         written_to_once_file = True
 
 
 @export
 def unset_once_if_written():
     """Unset the once file, if it has been written to."""
-    global once_file
-    if written_to_once_file:
+    global once_file, written_to_once_file
+    if written_to_once_file and once_file:
+        once_file.close()
         once_file = None
 
 
+@special_command('\\pipe_once', '\\| command',
+                 'Send next result to a subprocess.',
+                 aliases=('\\|', ))
+def set_pipe_once(arg, **_):
+    global pipe_once_process, written_to_pipe_once_process
+    pipe_once_cmd = shlex.split(arg)
+    if len(pipe_once_cmd) == 0:
+        raise OSError("pipe_once requires a command")
+    written_to_pipe_once_process = False
+    pipe_once_process = subprocess.Popen(pipe_once_cmd,
+                                         stdin=subprocess.PIPE,
+                                         stdout=subprocess.PIPE,
+                                         stderr=subprocess.PIPE,
+                                         bufsize=1,
+                                         encoding='UTF-8',
+                                         universal_newlines=True)
+    return [(None, None, None, "")]
+
+
+@export
+def write_pipe_once(output):
+    global pipe_once_process, written_to_pipe_once_process
+    if output and pipe_once_process:
+        try:
+            click.echo(output, file=pipe_once_process.stdin, nl=False)
+            click.echo(u"\n", file=pipe_once_process.stdin, nl=False)
+        except (IOError, OSError) as e:
+            pipe_once_process.terminate()
+            raise OSError(
+                "Failed writing to pipe_once subprocess: {}".format(e.strerror))
+        written_to_pipe_once_process = True
+
+
+@export
+def unset_pipe_once_if_written():
+    """Unset the pipe_once cmd, if it has been written to."""
+    global pipe_once_process, written_to_pipe_once_process
+    if written_to_pipe_once_process:
+        (stdout_data, stderr_data) = pipe_once_process.communicate()
+        if len(stdout_data) > 0:
+            print(stdout_data.rstrip(u"\n"))
+        if len(stderr_data) > 0:
+            print(stderr_data.rstrip(u"\n"))
+        pipe_once_process = None
+        written_to_pipe_once_process = False
+
+
 @special_command(
     'watch',
     'watch [seconds] [-c] query',