mirror of
				https://github.com/python/cpython.git
				synced 2025-10-26 03:04:41 +00:00 
			
		
		
		
	gh-133447: Add basic color to sqlite3 CLI (#133461)
				
					
				
			This commit is contained in:
		
							parent
							
								
									116a9f9b37
								
							
						
					
					
						commit
						30b1d8f11d
					
				
					 3 changed files with 36 additions and 13 deletions
				
			
		|  | @ -10,9 +10,10 @@ | |||
| from argparse import ArgumentParser | ||||
| from code import InteractiveConsole | ||||
| from textwrap import dedent | ||||
| from _colorize import get_theme, theme_no_color | ||||
| 
 | ||||
| 
 | ||||
| def execute(c, sql, suppress_errors=True): | ||||
| def execute(c, sql, suppress_errors=True, theme=theme_no_color): | ||||
|     """Helper that wraps execution of SQL code. | ||||
| 
 | ||||
|     This is used both by the REPL and by direct execution from the CLI. | ||||
|  | @ -25,11 +26,15 @@ def execute(c, sql, suppress_errors=True): | |||
|         for row in c.execute(sql): | ||||
|             print(row) | ||||
|     except sqlite3.Error as e: | ||||
|         t = theme.traceback | ||||
|         tp = type(e).__name__ | ||||
|         try: | ||||
|             print(f"{tp} ({e.sqlite_errorname}): {e}", file=sys.stderr) | ||||
|             tp += f" ({e.sqlite_errorname})" | ||||
|         except AttributeError: | ||||
|             print(f"{tp}: {e}", file=sys.stderr) | ||||
|             pass | ||||
|         print( | ||||
|             f"{t.type}{tp}{t.reset}: {t.message}{e}{t.reset}", file=sys.stderr | ||||
|         ) | ||||
|         if not suppress_errors: | ||||
|             sys.exit(1) | ||||
| 
 | ||||
|  | @ -37,10 +42,11 @@ def execute(c, sql, suppress_errors=True): | |||
| class SqliteInteractiveConsole(InteractiveConsole): | ||||
|     """A simple SQLite REPL.""" | ||||
| 
 | ||||
|     def __init__(self, connection): | ||||
|     def __init__(self, connection, use_color=False): | ||||
|         super().__init__() | ||||
|         self._con = connection | ||||
|         self._cur = connection.cursor() | ||||
|         self._use_color = use_color | ||||
| 
 | ||||
|     def runsource(self, source, filename="<input>", symbol="single"): | ||||
|         """Override runsource, the core of the InteractiveConsole REPL. | ||||
|  | @ -48,6 +54,8 @@ def runsource(self, source, filename="<input>", symbol="single"): | |||
|         Return True if more input is needed; buffering is done automatically. | ||||
|         Return False if input is a complete statement ready for execution. | ||||
|         """ | ||||
|         theme = get_theme(force_no_color=not self._use_color) | ||||
| 
 | ||||
|         if not source or source.isspace(): | ||||
|             return False | ||||
|         if source[0] == ".": | ||||
|  | @ -61,12 +69,13 @@ def runsource(self, source, filename="<input>", symbol="single"): | |||
|                 case "": | ||||
|                     pass | ||||
|                 case _ as unknown: | ||||
|                     self.write("Error: unknown command or invalid arguments:" | ||||
|                                f'  "{unknown}".\n') | ||||
|                     t = theme.traceback | ||||
|                     self.write(f'{t.type}Error{t.reset}:{t.message} unknown' | ||||
|                                f'command or invalid arguments:  "{unknown}".\n{t.reset}') | ||||
|         else: | ||||
|             if not sqlite3.complete_statement(source): | ||||
|                 return True | ||||
|             execute(self._cur, source) | ||||
|             execute(self._cur, source, theme=theme) | ||||
|         return False | ||||
| 
 | ||||
| 
 | ||||
|  | @ -113,17 +122,21 @@ def main(*args): | |||
|         Each command will be run using execute() on the cursor. | ||||
|         Type ".help" for more information; type ".quit" or {eofkey} to quit. | ||||
|     """).strip() | ||||
|     sys.ps1 = "sqlite> " | ||||
|     sys.ps2 = "    ... " | ||||
| 
 | ||||
|     theme = get_theme() | ||||
|     s = theme.syntax | ||||
| 
 | ||||
|     sys.ps1 = f"{s.prompt}sqlite> {s.reset}" | ||||
|     sys.ps2 = f"{s.prompt}    ... {s.reset}" | ||||
| 
 | ||||
|     con = sqlite3.connect(args.filename, isolation_level=None) | ||||
|     try: | ||||
|         if args.sql: | ||||
|             # SQL statement provided on the command-line; execute it directly. | ||||
|             execute(con, args.sql, suppress_errors=False) | ||||
|             execute(con, args.sql, suppress_errors=False, theme=theme) | ||||
|         else: | ||||
|             # No SQL provided; start the REPL. | ||||
|             console = SqliteInteractiveConsole(con) | ||||
|             console = SqliteInteractiveConsole(con, use_color=True) | ||||
|             try: | ||||
|                 import readline  # noqa: F401 | ||||
|             except ImportError: | ||||
|  |  | |||
|  | @ -8,10 +8,11 @@ | |||
|     captured_stdout, | ||||
|     captured_stderr, | ||||
|     captured_stdin, | ||||
|     force_not_colorized, | ||||
|     force_not_colorized_test_class, | ||||
| ) | ||||
| 
 | ||||
| 
 | ||||
| @force_not_colorized_test_class | ||||
| class CommandLineInterface(unittest.TestCase): | ||||
| 
 | ||||
|     def _do_test(self, *args, expect_success=True): | ||||
|  | @ -37,7 +38,6 @@ def expect_failure(self, *args): | |||
|         self.assertEqual(out, "") | ||||
|         return err | ||||
| 
 | ||||
|     @force_not_colorized | ||||
|     def test_cli_help(self): | ||||
|         out = self.expect_success("-h") | ||||
|         self.assertIn("usage: ", out) | ||||
|  | @ -69,6 +69,7 @@ def test_cli_on_disk_db(self): | |||
|         self.assertIn("(0,)", out) | ||||
| 
 | ||||
| 
 | ||||
| @force_not_colorized_test_class | ||||
| class InteractiveSession(unittest.TestCase): | ||||
|     MEMORY_DB_MSG = "Connected to a transient in-memory database" | ||||
|     PS1 = "sqlite> " | ||||
|  | @ -190,6 +191,14 @@ def test_interact_on_disk_file(self): | |||
|         out, _ = self.run_cli(TESTFN, commands=("SELECT count(t) FROM t;",)) | ||||
|         self.assertIn("(0,)\n", out) | ||||
| 
 | ||||
|     def test_color(self): | ||||
|         with unittest.mock.patch("_colorize.can_colorize", return_value=True): | ||||
|             out, err = self.run_cli(commands="TEXT\n") | ||||
|             self.assertIn("\x1b[1;35msqlite> \x1b[0m", out) | ||||
|             self.assertIn("\x1b[1;35m    ... \x1b[0m\x1b", out) | ||||
|             out, err = self.run_cli(commands=("sel;",)) | ||||
|             self.assertIn('\x1b[1;35mOperationalError (SQLITE_ERROR)\x1b[0m: ' | ||||
|                           '\x1b[35mnear "sel": syntax error\x1b[0m', err) | ||||
| 
 | ||||
| if __name__ == "__main__": | ||||
|     unittest.main() | ||||
|  |  | |||
|  | @ -0,0 +1 @@ | |||
| Add basic color to :mod:`sqlite3` CLI interface. | ||||
		Loading…
	
	Add table
		Add a link
		
	
		Reference in a new issue
	
	 Stan Ulbrych
						Stan Ulbrych