Merge branch 'master' into prompt-ansi-escape-sequences

Author: amjith

changelog.md

@@ -7,11 +7,18 @@ Features:
 * Add an option `--init-command` to execute SQL after connecting (Thanks: [KITAGAWA Yasutaka]).
 * Use InputMode.REPLACE_SINGLE
 * Add support for ANSI escape sequences for coloring the prompt.
+* Allow customization of Pygments SQL syntax-highlighting styles.
+* Add a `\clip` special command to copy queries to the system clipboard.
+* Add a special command `\pipe_once` to pipe output to a subprocess.
+
 
 Bug Fixes:
 ----------
 * Fixed compatibility with sqlparse 0.4 (Thanks: [mtorromeo]).
 *  Fixed iPython magic (Thanks: [mwcm]).
+* Send "Connecting to socket" message to the standard error.
+* Respect empty string for prompt_continuation via `prompt_continuation = ''` in `.myclirc`
+* Fix \once -o to overwrite output whole, instead of line-by-line.
 
 1.22.2
 ======

mycli/clistyle.py

@@ -44,6 +44,36 @@
     v: k for k, v in TOKEN_TO_PROMPT_STYLE.items()
 }
 
+# all tokens that the Pygments MySQL lexer can produce
+OVERRIDE_STYLE_TO_TOKEN = {
+    'sql.comment': Token.Comment,
+    'sql.comment.multi-line': Token.Comment.Multiline,
+    'sql.comment.single-line': Token.Comment.Single,
+    'sql.comment.optimizer-hint': Token.Comment.Special,
+    'sql.escape': Token.Error,
+    'sql.keyword': Token.Keyword,
+    'sql.datatype': Token.Keyword.Type,
+    'sql.literal': Token.Literal,
+    'sql.literal.date': Token.Literal.Date,
+    'sql.symbol': Token.Name,
+    'sql.quoted-schema-object': Token.Name.Quoted,
+    'sql.quoted-schema-object.escape': Token.Name.Quoted.Escape,
+    'sql.constant': Token.Name.Constant,
+    'sql.function': Token.Name.Function,
+    'sql.variable': Token.Name.Variable,
+    'sql.number': Token.Number,
+    'sql.number.binary': Token.Number.Bin,
+    'sql.number.float': Token.Number.Float,
+    'sql.number.hex': Token.Number.Hex,
+    'sql.number.integer': Token.Number.Integer,
+    'sql.operator': Token.Operator,
+    'sql.punctuation': Token.Punctuation,
+    'sql.string': Token.String,
+    'sql.string.double-quouted': Token.String.Double,
+    'sql.string.escape': Token.String.Escape,
+    'sql.string.single-quoted': Token.String.Single,
+    'sql.whitespace': Token.Text,
+}
 
 def parse_pygments_style(token_name, style_object, style_dict):
     """Parse token type and style string.
@@ -108,6 +138,9 @@ def style_factory_output(name, cli_style):
         elif token in PROMPT_STYLE_TO_TOKEN:
             token_type = PROMPT_STYLE_TO_TOKEN[token]
             style.update({token_type: cli_style[token]})
+        elif token in OVERRIDE_STYLE_TO_TOKEN:
+            token_type = OVERRIDE_STYLE_TO_TOKEN[token]
+            style.update({token_type: cli_style[token]})
         else:
             # TODO: cli helpers will have to switch to ptk.Style
             logger.error('Unhandled style / class name: %s', token)

mycli/main.py

@@ -439,7 +439,7 @@ def _connect():
             if not WIN and socket:
                 socket_owner = getpwuid(os.stat(socket).st_uid).pw_name
                 self.echo(
-                    f"Connecting to socket {socket}, owned by user {socket_owner}")
+                    f"Connecting to socket {socket}, owned by user {socket_owner}", err=True)
                 try:
                     _connect()
                 except OperationalError as e:
@@ -512,6 +512,24 @@ def handle_editor_command(self, text):
             continue
         return text
 
+    def handle_clip_command(self, text):
+        """A clip command is any query that is prefixed or suffixed by a
+        '\clip'.
+
+        :param text: Document
+        :return: Boolean
+
+        """
+
+        if special.clip_command(text):
+            query = (special.get_clip_query(text) or
+                     self.get_last_query())
+            message = special.copy_query_to_clipboard(sql=query)
+            if message:
+                raise RuntimeError(message)
+            return True
+        return False
+
     def run_cli(self):
         iterations = 0
         sqlexecute = self.sqlexecute
@@ -553,7 +571,9 @@ def get_message():
             return ANSI(prompt)
 
         def get_continuation(width, *_):
-            if self.multiline_continuation_char:
+            if self.multiline_continuation_char == '':
+                continuation = ''
+            elif self.multiline_continuation_char:
                 left_padding = width - len(self.multiline_continuation_char)
                 continuation = " " * \
                     max((left_padding - 1), 0) + \
@@ -582,6 +602,15 @@ def one_iteration(text=None):
                     self.echo(str(e), err=True, fg='red')
                     return
 
+                try:
+                    if self.handle_clip_command(text):
+                        return
+                except RuntimeError as e:
+                    logger.error("sql: %r, error: %r", text, e)
+                    logger.error("traceback: %r", traceback.format_exc())
+                    self.echo(str(e), err=True, fg='red')
+                    return
+
             if not text.strip():
                 return
 
@@ -656,6 +685,7 @@ def one_iteration(text=None):
                     result_count += 1
                     mutating = mutating or destroy or is_mutating(status)
                 special.unset_once_if_written()
+                special.unset_pipe_once_if_written()
             except EOFError as e:
                 raise e
             except KeyboardInterrupt:
@@ -816,6 +846,7 @@ def output(self, output, status=None):
                 self.log_output(line)
                 special.write_tee(line)
                 special.write_once(line)
+                special.write_pipe_once(line)
 
                 if fits or output_via_pager:
                     # buffering
@@ -1293,7 +1324,7 @@ def is_select(status):
 def thanks_picker(files=()):
     contents = []
     for line in fileinput.input(files=files):
-        m = re.match('^ *\* (.*)', line)
+        m = re.match(r'^ *\* (.*)', line)
         if m:
             contents.append(m.group(1))
     return choice(contents)

mycli/myclirc

@@ -41,6 +41,7 @@ table_format = ascii
 # friendly, monokai, paraiso, colorful, murphy, bw, pastie, paraiso, trac, default,
 # fruity.
 # Screenshots at http://mycli.net/syntax
+# Can be further modified in [colors]
 syntax_style = default
 
 # Keybindings: Possible values: emacs, vi.
@@ -114,6 +115,35 @@ output.odd-row = ""
 output.even-row = ""
 output.null = "#808080"
 
+# SQL syntax highlighting overrides
+# sql.comment = 'italic #408080'
+# sql.comment.multi-line = ''
+# sql.comment.single-line = ''
+# sql.comment.optimizer-hint = ''
+# sql.escape = 'border:#FF0000'
+# sql.keyword = 'bold #008000'
+# sql.datatype = 'nobold #B00040'
+# sql.literal = ''
+# sql.literal.date = ''
+# sql.symbol = ''
+# sql.quoted-schema-object = ''
+# sql.quoted-schema-object.escape = ''
+# sql.constant = '#880000'
+# sql.function = '#0000FF'
+# sql.variable = '#19177C'
+# sql.number = '#666666'
+# sql.number.binary = ''
+# sql.number.float = ''
+# sql.number.hex = ''
+# sql.number.integer = ''
+# sql.operator = '#666666'
+# sql.punctuation = ''
+# sql.string = '#BA2121'
+# sql.string.double-quouted = ''
+# sql.string.escape = 'bold #BB6622'
+# sql.string.single-quoted = ''
+# sql.whitespace = ''
+
 # Favorite queries.
 [favorite_queries]
 

mycli/packages/parseutils.py

@@ -11,7 +11,7 @@
         # This matches everything except spaces, parens, colon, comma, and period
         'most_punctuations': re.compile(r'([^\.():,\s]+)$'),
         # This matches everything except a space.
-        'all_punctuations': re.compile('([^\s]+)$'),
+        'all_punctuations': re.compile(r'([^\s]+)$'),
         }
 
 def last_word(text, include='alphanum_underscore'):

mycli/packages/special/iocommands.py

@@ -8,6 +8,7 @@
 from time import sleep
 
 import click
+import pyperclip
 import sqlparse
 
 from . import export
@@ -23,6 +24,8 @@
 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('(^\\\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)"""
@@ -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',

mycli/packages/special/main.py

@@ -112,6 +112,8 @@ def quit(*_args):
 
 @special_command('\\e', '\\e', 'Edit command with editor (uses $EDITOR).',
                  arg_type=NO_QUERY, case_sensitive=True)
+@special_command('\\clip', '\\clip', 'Copy query to the system clipboard.',
+                 arg_type=NO_QUERY, case_sensitive=True)
 @special_command('\\G', '\\G', 'Display current query results vertically.',
                  arg_type=NO_QUERY, case_sensitive=True)
 def stub():

setup.py

@@ -25,6 +25,7 @@
     'configobj >= 5.0.5',
     'cryptography >= 1.0.0',
     'cli_helpers[styles] >= 2.0.1',
+    'pyperclip >= 1.8.1'
 ]