Merge pull request #1 from Rioran/main

style improvements

dbcls/clients/base.py

@@ -1,5 +1,8 @@
 import abc
-from dataclasses import dataclass, field
+from dataclasses import (
+    dataclass,
+    field,
+)
 
 
 @dataclass
@@ -12,13 +15,13 @@ def __str__(self) -> str:
         if self.message:
             return self.message
 
-        if not self.data and self.rowcount:
-            return f'{self.rowcount} rows affected'
+        if self.data:
+            return f'{self.rowcount} rows returned'
 
-        if not self.data and not self.rowcount:
-            return 'Empty set'
+        if self.rowcount:
+            return f'{self.rowcount} rows affected'
 
-        return f'{self.rowcount} rows returned'
+        return 'Empty set'
 
 
 class ClientClass(abc.ABC):
@@ -43,7 +46,7 @@ async def change_database(self, database: str):
         old_db = self.dbname
         self.dbname = database
         try:
-            await self.execute(f'SELECT 1')
+            await self.execute('SELECT 1')
             return Result(message=f'You are now connected to database "{database}"')
         except Exception:
             self.dbname = old_db

dbcls/clients/clickhouse.py

@@ -1,8 +1,10 @@
-from aiohttp import ClientSession
 import aiochclient
+from aiohttp import ClientSession
 
-from .base import ClientClass
-from .base import Result
+from .base import (
+    ClientClass,
+    Result,
+)
 
 
 class ClickhouseClient(ClientClass):

dbcls/clients/mysql.py

@@ -1,7 +1,9 @@
 import aiomysql
 
-from .base import ClientClass
-from .base import Result
+from .base import (
+    ClientClass,
+    Result,
+)
 
 
 class MysqlClient(ClientClass):

dbcls/clients/postgres.py

@@ -1,8 +1,11 @@
 import aiopg
-import psycopg2
+from psycopg2 import ProgrammingError
+from psycopg2.extras import RealDictCursor
 
-from .base import ClientClass
-from .base import Result
+from .base import (
+    ClientClass,
+    Result,
+)
 
 
 class PostgresClient(ClientClass):
@@ -14,23 +17,32 @@ def __init__(self, host, username, password, dbname, port='5432'):
             self.port = '5432'
 
     async def get_tables(self) -> Result:
-        sql = "SELECT table_name FROM information_schema.tables WHERE table_schema='public' AND table_type='BASE TABLE';"
+        sql = (
+            "SELECT table_name FROM information_schema.tables "
+            "WHERE table_schema='public' AND table_type='BASE TABLE';"
+        )
         return await self.execute(sql)
 
     async def get_databases(self) -> Result:
         sql = "SELECT datname FROM pg_database;"
         return await self.execute(sql)
 
     async def execute(self, sql) -> Result:
-        if sql.strip().startswith('\\c '):
-            db = sql.strip().split(' ')[1].rstrip(';')
-            return await self.change_database(db)
+        sql_stripped = sql.strip()
+        first_word = sql_stripped.split(' ')[1]
 
-        if sql.strip().startswith('\\d '):
-            sql = f"SELECT column_name, data_type FROM information_schema.columns WHERE table_name = '{sql.strip().split(' ')[1]}'"
-        if sql.strip() == ('\\d'):
+        if sql_stripped.startswith('\\c '):
+            db = first_word.rstrip(';')
+            return await self.change_database(db)
+        if sql_stripped.startswith('\\d '):
+            sql = (
+                "SELECT column_name, data_type"
+                "FROM information_schema.columns"
+                f"WHERE table_name = '{first_word}'"
+            )
+        if sql_stripped == ('\\d'):
             return await self.get_tables()
-        if sql.strip().startswith('\\l'):
+        if sql_stripped.startswith('\\l'):
             return await self.get_databases()
 
         async with aiopg.connect(
@@ -40,14 +52,13 @@ async def execute(self, sql) -> Result:
             password=self.password,
             dbname=self.dbname,
         ) as conn:
-            async with conn.cursor(cursor_factory=psycopg2.extras.RealDictCursor) as cur:
+            async with conn.cursor(cursor_factory=RealDictCursor) as cur:
                 await cur.execute(sql)
                 rowcount = cur.rowcount
                 result = Result(rowcount=rowcount)
                 try:
-
                     result.data = await cur.fetchall()
-                except psycopg2.ProgrammingError:
+                except ProgrammingError:
                     pass
 
-                return result
+                return result

dbcls/dbcls.py

@@ -1,53 +1,58 @@
-from time import time
 import asyncio
+import curses
 from functools import partial
+from time import time
 from typing import Callable
 from typing import Optional
 
 import kaa
 import kaa.cui.main
+import visidata
+from kaa.addon import (
+    alt,
+    backspace,
+    command,
+    ctrl,
+    setup,
+)
+from kaa.cui.editor import TextEditorWindow
 from kaa.cui.keydef import KeyEvent
 from kaa.filetype.default.defaultmode import DefaultMode
 from kaa.options import build_parser
 from kaa.syntax_highlight import DefaultToken
-from kaa.cui.editor import TextEditorWindow
-from kaa.addon import command
-from kaa.addon import setup
-from kaa.addon import alt
-from kaa.addon import ctrl
-from kaa.addon import backspace
-import visidata
-import curses
 from ssh_crypt import E
 
-from .sql_tokenizer import make_tokenizer, sql_editor_themes, CaseInsensitiveKeywords, NonSqlComment, Span
 from .clients.base import Result
-
+from .sql_tokenizer import (
+    CaseInsensitiveKeywords,
+    NonSqlComment,
+    make_tokenizer,
+    sql_editor_themes,
+)
 
 client = None
 
 
 def get_sel(wnd: TextEditorWindow) -> str:
-    if wnd.screen.selection.is_selected():
-        if not wnd.screen.selection.is_rectangular():
-            f, t = wnd.screen.selection.get_selrange()
-            return wnd.document.gettext(f, t)
+    selection = wnd.screen.selection
+    if not selection.is_selected():
+        return
+    if not selection.is_rectangular():
+        selected_from, selected_to = selection.get_selrange()
+        return wnd.document.gettext(selected_from, selected_to)
+    data = []
+    position_from, position_to, column_from, column_to = selection.get_rect_range()
+
+    while position_from < position_to:
+        position_and_col_string = selection.get_col_string(position_from, column_from, column_to)
+        if position_and_col_string:
+            *_, col_string = position_and_col_string
+            data.append(col_string.rstrip('\n'))
         else:
-            s = []
-            (posfrom, posto, colfrom, colto
-             ) = wnd.screen.selection.get_rect_range()
+            data.append('')
+        position_from = wnd.document.geteol(position_from)
 
-            while posfrom < posto:
-                sel = wnd.screen.selection.get_col_string(
-                    posfrom, colfrom, colto)
-                if sel:
-                    f, t, org = sel
-                    s.append(org.rstrip('\n'))
-                else:
-                    s.append('')
-                posfrom = wnd.document.geteol(posfrom)
-
-            return '\n'.join(s)
+    return '\n'.join(data)
 
 
 def get_current_sql_rows_pos(wnd: TextEditorWindow) -> list[int]:
@@ -141,7 +146,10 @@ def get_expression_under_cursor(wnd: TextEditorWindow) -> str:
     return line
 
 
-async def await_and_print_time(wnd: TextEditorWindow, coro: asyncio.coroutines) -> Result:
+async def await_and_print_time(
+        wnd: TextEditorWindow,
+        coro: asyncio.coroutines
+) -> Result:
     start = time()
     task = asyncio.create_task(coro)
 
@@ -247,7 +255,12 @@ def on_keypressed(
     return original_fn(wnd, event, key, commands, candidate)
 
 
-def on_cursor_located(self: DefaultMode, original_fn: Callable, wnd: TextEditorWindow, *args, **kwargs):
+def on_cursor_located(
+        self: DefaultMode,
+        original_fn: Callable,
+        wnd: TextEditorWindow,
+        *args, **kwargs
+):
     wnd.document.highlights = []
     for id, row_pos in enumerate(get_current_sql_rows_pos(wnd)):
         wnd.document.highlights.append(

dbcls/sql_tokenizer.py

@@ -1,6 +1,22 @@
-from kaa.theme import Style
 from kaa import doc_re
-from kaa.syntax_highlight import Span, SingleToken, Tokenizer, Keywords, Token
+from kaa.syntax_highlight import (
+    Keywords,
+    SingleToken,
+    Span,
+    Token,
+    Tokenizer,
+)
+from kaa.theme import Style
+
+
+KEYWORDS = [
+    'SELECT', 'UPDATE', 'DELETE', 'DROP', 'ALTER', 'COLUMN', 'USE',
+    'FROM', 'JOIN', 'OUTER', 'INNER', 'LIMIT', 'ORDER BY', 'AS',
+    'SHOW', 'FROM', 'WHERE', 'DESC', 'TABLES', 'CREATE', 'TABLE',
+    'SET', 'IS', 'NOT', 'NULL', 'ON', 'IN', 'LIKE', 'ILIKE', 'AND',
+    'OR', 'INSERT', 'INTO', 'VALUES', 'INTERVAL', 'GROUP', 'BY',
+    'HAVING', 'GRANT', 'LEFT', 'RIGHT', 'FULL', 'CROSS'
+]
 
 
 sql_editor_themes = {
@@ -31,16 +47,6 @@ class NonSqlComment(Span):
     pass
 
 
-KEYWORDS = [
-    'SELECT', 'UPDATE', 'DELETE', 'DROP', 'ALTER', 'COLUMN', 'USE',
-    'FROM', 'JOIN', 'OUTER', 'INNER', 'LIMIT', 'ORDER BY', 'AS',
-    'SHOW', 'FROM', 'WHERE', 'DESC', 'TABLES', 'CREATE', 'TABLE',
-    'SET', 'IS', 'NOT', 'NULL', 'ON', 'IN', 'LIKE', 'ILIKE', 'AND',
-    'OR', 'INSERT', 'INTO', 'VALUES', 'INTERVAL', 'GROUP', 'BY',
-    'HAVING', 'GRANT', 'LEFT', 'RIGHT', 'FULL', 'CROSS'
-]
-
-
 def sqleditor_tokens() -> list[tuple[str, Token]]:
     return [
         ('comment1', Span('comment', r'--', '$')),

setup.py

@@ -1,5 +1,9 @@
 import os
-from setuptools import setup, find_packages
+
+from setuptools import (
+    find_packages,
+    setup,
+)
 
 
 def get_requirements():