Skip to content

Commit 30b1d8f

Browse files
gh-133447: Add basic color to sqlite3 CLI (#133461)
1 parent 116a9f9 commit 30b1d8f

File tree

3 files changed

+36
-13
lines changed

3 files changed

+36
-13
lines changed

Lib/sqlite3/__main__.py

+24-11
Original file line numberDiff line numberDiff line change
@@ -10,9 +10,10 @@
1010
from argparse import ArgumentParser
1111
from code import InteractiveConsole
1212
from textwrap import dedent
13+
from _colorize import get_theme, theme_no_color
1314

1415

15-
def execute(c, sql, suppress_errors=True):
16+
def execute(c, sql, suppress_errors=True, theme=theme_no_color):
1617
"""Helper that wraps execution of SQL code.
1718
1819
This is used both by the REPL and by direct execution from the CLI.
@@ -25,29 +26,36 @@ def execute(c, sql, suppress_errors=True):
2526
for row in c.execute(sql):
2627
print(row)
2728
except sqlite3.Error as e:
29+
t = theme.traceback
2830
tp = type(e).__name__
2931
try:
30-
print(f"{tp} ({e.sqlite_errorname}): {e}", file=sys.stderr)
32+
tp += f" ({e.sqlite_errorname})"
3133
except AttributeError:
32-
print(f"{tp}: {e}", file=sys.stderr)
34+
pass
35+
print(
36+
f"{t.type}{tp}{t.reset}: {t.message}{e}{t.reset}", file=sys.stderr
37+
)
3338
if not suppress_errors:
3439
sys.exit(1)
3540

3641

3742
class SqliteInteractiveConsole(InteractiveConsole):
3843
"""A simple SQLite REPL."""
3944

40-
def __init__(self, connection):
45+
def __init__(self, connection, use_color=False):
4146
super().__init__()
4247
self._con = connection
4348
self._cur = connection.cursor()
49+
self._use_color = use_color
4450

4551
def runsource(self, source, filename="<input>", symbol="single"):
4652
"""Override runsource, the core of the InteractiveConsole REPL.
4753
4854
Return True if more input is needed; buffering is done automatically.
4955
Return False if input is a complete statement ready for execution.
5056
"""
57+
theme = get_theme(force_no_color=not self._use_color)
58+
5159
if not source or source.isspace():
5260
return False
5361
if source[0] == ".":
@@ -61,12 +69,13 @@ def runsource(self, source, filename="<input>", symbol="single"):
6169
case "":
6270
pass
6371
case _ as unknown:
64-
self.write("Error: unknown command or invalid arguments:"
65-
f' "{unknown}".\n')
72+
t = theme.traceback
73+
self.write(f'{t.type}Error{t.reset}:{t.message} unknown'
74+
f'command or invalid arguments: "{unknown}".\n{t.reset}')
6675
else:
6776
if not sqlite3.complete_statement(source):
6877
return True
69-
execute(self._cur, source)
78+
execute(self._cur, source, theme=theme)
7079
return False
7180

7281

@@ -113,17 +122,21 @@ def main(*args):
113122
Each command will be run using execute() on the cursor.
114123
Type ".help" for more information; type ".quit" or {eofkey} to quit.
115124
""").strip()
116-
sys.ps1 = "sqlite> "
117-
sys.ps2 = " ... "
125+
126+
theme = get_theme()
127+
s = theme.syntax
128+
129+
sys.ps1 = f"{s.prompt}sqlite> {s.reset}"
130+
sys.ps2 = f"{s.prompt} ... {s.reset}"
118131

119132
con = sqlite3.connect(args.filename, isolation_level=None)
120133
try:
121134
if args.sql:
122135
# SQL statement provided on the command-line; execute it directly.
123-
execute(con, args.sql, suppress_errors=False)
136+
execute(con, args.sql, suppress_errors=False, theme=theme)
124137
else:
125138
# No SQL provided; start the REPL.
126-
console = SqliteInteractiveConsole(con)
139+
console = SqliteInteractiveConsole(con, use_color=True)
127140
try:
128141
import readline # noqa: F401
129142
except ImportError:

Lib/test/test_sqlite3/test_cli.py

+11-2
Original file line numberDiff line numberDiff line change
@@ -8,10 +8,11 @@
88
captured_stdout,
99
captured_stderr,
1010
captured_stdin,
11-
force_not_colorized,
11+
force_not_colorized_test_class,
1212
)
1313

1414

15+
@force_not_colorized_test_class
1516
class CommandLineInterface(unittest.TestCase):
1617

1718
def _do_test(self, *args, expect_success=True):
@@ -37,7 +38,6 @@ def expect_failure(self, *args):
3738
self.assertEqual(out, "")
3839
return err
3940

40-
@force_not_colorized
4141
def test_cli_help(self):
4242
out = self.expect_success("-h")
4343
self.assertIn("usage: ", out)
@@ -69,6 +69,7 @@ def test_cli_on_disk_db(self):
6969
self.assertIn("(0,)", out)
7070

7171

72+
@force_not_colorized_test_class
7273
class InteractiveSession(unittest.TestCase):
7374
MEMORY_DB_MSG = "Connected to a transient in-memory database"
7475
PS1 = "sqlite> "
@@ -190,6 +191,14 @@ def test_interact_on_disk_file(self):
190191
out, _ = self.run_cli(TESTFN, commands=("SELECT count(t) FROM t;",))
191192
self.assertIn("(0,)\n", out)
192193

194+
def test_color(self):
195+
with unittest.mock.patch("_colorize.can_colorize", return_value=True):
196+
out, err = self.run_cli(commands="TEXT\n")
197+
self.assertIn("\x1b[1;35msqlite> \x1b[0m", out)
198+
self.assertIn("\x1b[1;35m ... \x1b[0m\x1b", out)
199+
out, err = self.run_cli(commands=("sel;",))
200+
self.assertIn('\x1b[1;35mOperationalError (SQLITE_ERROR)\x1b[0m: '
201+
'\x1b[35mnear "sel": syntax error\x1b[0m', err)
193202

194203
if __name__ == "__main__":
195204
unittest.main()
Original file line numberDiff line numberDiff line change
@@ -0,0 +1 @@
1+
Add basic color to :mod:`sqlite3` CLI interface.

0 commit comments

Comments
 (0)