Improved keyword capitalisation behaviour

src/crate/crash/repl.py

@@ -21,6 +21,7 @@
 # software solely pursuant to the terms of the relevant commercial agreement.
 
 import os
+import re
 
 from pygments.lexers.sql import SqlLexer
 from pygments.style import Style
@@ -172,22 +173,39 @@ def is_multiline():
 
 class Capitalizer:
 
+    KEYWORD_RE = r'(?:"\w+)|(?:\'\w+)|\w+'
+
     def __init__(self, cmd):
         self.cmd = cmd
         self.last_changed = None
 
     def __call__(self, buffer):
         if not self.cmd.should_autocapitalize():
             return
-        last_word = buffer.document.get_word_before_cursor(WORD=True)
-        if not last_word.isupper() and last_word.lower() in SQLCompleter.keywords:
-            buffer.delete_before_cursor(len(last_word))
-            buffer.insert_text(last_word.upper(), fire_event=False)
-            self.last_changed = last_word
-        elif self.last_changed and last_word.startswith(self.last_changed.upper()):
-            buffer.delete_before_cursor(len(last_word))
-            buffer.insert_text(self.last_changed + last_word[-1:], fire_event=False)
 
+        current_line = buffer.document.text
+        cursor_position = buffer.document.cursor_position
+
+        if self.last_changed and self.is_prefix(current_line.lower(), self.last_changed.lower()):
+            diff = len(self.last_changed) - len(current_line)
+            current_line = self.last_changed + current_line[diff:]
+
+        new_line = re.sub(self.KEYWORD_RE, self.keyword_replacer, current_line)
+
+        if new_line != buffer.document.text:
+            buffer.delete_before_cursor(cursor_position)
+            buffer.delete(len(new_line) - cursor_position)
+            buffer.insert_text(new_line, fire_event=False)
+            self.last_changed = current_line
+
+    def keyword_replacer(self, match):
+        if match.group(0).lower() in SQLCompleter.keywords:
+            return match.group(0).upper()
+        else:
+            return match.group(0)
+
+    def is_prefix(self, string, prefix):
+        return string.startswith(prefix) and string != prefix
 
 def loop(cmd, history_file):
     key_binding_manager = KeyBindingManager(

src/crate/crash/test_repl.py

@@ -22,7 +22,7 @@
 from prompt_toolkit.document import Document
 from .repl import SQLCompleter, Capitalizer
 from .command import CrateCmd
-
+import re
 
 class SQLCompleterTest(TestCase):
 
@@ -63,6 +63,27 @@ def test_capitalize(self):
         self.capitalizer(buffer)
         self.assertEqual(u'CREATE TABLE "select', buffer.text)
 
+        text = u'CREATE TABLE \'select\''
+        buffer.set_document(Document(text, len(text)))
+        self.capitalizer(buffer)
+        self.assertEqual(u'CREATE TABLE \'select\'', buffer.text)
+
+        text = u'create table test (x int)'
+        buffer.set_document(Document(text, len(text)))
+        self.capitalizer(buffer)
+        self.assertEqual(u'CREATE TABLE test (x INT)', buffer.text)
+
+        text = u'create table test (a boolean, b string, c integer)'
+        buffer.set_document(Document(text, len(text)))
+        self.capitalizer(buffer)
+        self.assertEqual(u'CREATE TABLE test (a BOOLEAN, b STRING, c INTEGER)', buffer.text)
+
+        text = u'create table test\n(a boolean, b string, c integer)'
+        buffer.set_document(Document(text, len(text)))
+        self.capitalizer(buffer)
+        self.assertEqual(u'CREATE TABLE test\n(a BOOLEAN, b STRING, c INTEGER)', buffer.text)
+
+
     def test_undo_capitalize(self):
         buffer = Buffer()
 
@@ -85,3 +106,21 @@ def test_undo_capitalize(self):
         buffer.set_document(Document(text, len(text)))
         self.capitalizer(buffer)
         self.assertEqual(u'Selection', buffer.text)
+
+    def test_keyword_replacer(self):
+        KEYWORD_RE = r'\w+'
+        text = u'select'
+        match = re.match(KEYWORD_RE, text)
+        self.assertEqual(u'SELECT', self.capitalizer.keyword_replacer(match))
+
+        text = u'definitelyNotAnSQLKeyword'
+        match = re.match(KEYWORD_RE, text)
+        self.assertEqual(u'definitelyNotAnSQLKeyword', self.capitalizer.keyword_replacer(match))
+
+    def test_is_prefix(self):
+        string = u'SELECT * FROM test'
+        prefix = u'SELECT * FROM'
+        self.assertTrue(self.capitalizer.is_prefix(string, prefix))
+
+        prefix = u'SELECT testCol FROM'
+        self.assertFalse(self.capitalizer.is_prefix(string, prefix))