diff --git a/README.md b/README.md index a8e909c69..a312b109b 100644 --- a/README.md +++ b/README.md @@ -90,9 +90,11 @@ pip install opteryx Example usage, filtering one of the internal example datasets and saving the results as a CSV. ~~~bash -python -m opteryx --o 'planets.csv' "SELECT * FROM \$planets" +python -m opteryx "SELECT * FROM \$astronauts LIMIT 10;" ~~~ +![Opteryx](https://github.com/mabel-dev/opteryx.dev/raw/main/assets/cli.png) + **Query Data (Python)** Example usage, querying one of the internal example datasets. diff --git a/opteryx/__main__.py b/opteryx/__main__.py index d24af5492..958e83116 100644 --- a/opteryx/__main__.py +++ b/opteryx/__main__.py @@ -45,7 +45,7 @@ def main( cur.execute(sql) if o == "console": - print(display.ascii_table(cur.arrow(), display_width=True)) + print(display.ascii_table(cur.arrow(), limit=-1, display_width=True)) return else: ext = o.lower().split(".")[-1] diff --git a/opteryx/third_party/pyarrow_ops/ops.py b/opteryx/third_party/pyarrow_ops/ops.py index 3bae3297b..e60ae0efe 100644 --- a/opteryx/third_party/pyarrow_ops/ops.py +++ b/opteryx/third_party/pyarrow_ops/ops.py @@ -286,7 +286,7 @@ def drop_duplicates(table, columns=None): # Show for easier printing -def head(table, n=5, max_width=100, colorize: bool = False): +def head(table, n=5, max_width=100): # Updated to yield rather than print for Opteryx if table == set(): # pragma: no cover yield "No data in table" @@ -295,8 +295,6 @@ def head(table, n=5, max_width=100, colorize: bool = False): yield "No data in table" return - from opteryx.utils.colors import colorize as add_color - # Extract head data if n > 0: t = table.slice(length=n) @@ -317,6 +315,4 @@ def head(table, n=5, max_width=100, colorize: bool = False): w.ljust(max(cw, dw) + 2) for w, cw, dw in zip(data[i], col_width, data_width) ] - yield add_color( - "{UNDERLINE_WHITE}Row " if i == 0 else str(i - 1).ljust(5) - ) + "".join(adjust)[:max_width] + add_color("{OFF}") + yield "Row " if i == 0 else str(i - 1).ljust(5) + "".join(adjust)[:max_width] diff --git a/opteryx/utils/colors.py b/opteryx/utils/colors.py index de6a9987a..d5132f80d 100644 --- a/opteryx/utils/colors.py +++ b/opteryx/utils/colors.py @@ -4,10 +4,10 @@ "{BLACK}": "\033[0;30m", # Black "{RED}": "\033[0;31m", # Red "{GREEN}": "\033[0;32m", # Green - "{YELLOW}": "\033[0;33m", # Yellow + "{YELLOW}": "\033[38;5;228m", # "\033[0;33m", # Yellow "{BLUE}": "\033[0;34m", # Blue "{PURPLE}": "\033[0;35m", # Purple - "{CYAN}": "\033[0;36m", # Cyan + "{CYAN}": "\033[38;5;117m", # "\033[0;36m", # Cyan "{WHITE}": "\033[0;37m", # White # Bold "{BOLD_BLACK}": "\033[1;30m", # Black @@ -27,6 +27,15 @@ "{UNDERLINE_PURPLE}": "\033[4;35m", # Purple "{UNDERLINE_CYAN}": "\033[4;36m", # Cyan "{UNDERLINE_WHITE}": "\033[4;37m", # White + # Background + "{BACKGROUND_BLACK}": "\033[40m", # Black + "{BACKGROUND_RED}": "\033[41m", # Red + "{BACKGROUND_GREEN}": "\033[42m", # Green + "{BACKGROUND_YELLOW}": "\033[43m", # Yellow + "{BACKGROUND_BLUE}": "\033[44m", # Blue + "{BACKGROUND_PURPLE}": "\033[45m", # Purple + "{BACKGROUND_CYAN}": "\033[46m", # Cyan + "{BACKGROUND_WHITE}": "\033[47m", # White } diff --git a/opteryx/utils/display.py b/opteryx/utils/display.py index 9e258cb22..b9fa61730 100644 --- a/opteryx/utils/display.py +++ b/opteryx/utils/display.py @@ -12,6 +12,8 @@ from typing import Any, Dict, Iterable, Union +import datetime + def html_table(dictset: Iterable[dict], limit: int = 5): # pragma: no cover """ @@ -90,8 +92,9 @@ def ascii_table( table: Iterable[Dict[Any, Any]], limit: int = 5, display_width: Union[bool, int] = True, + max_column_width: int = 30, colorize: bool = True, -): # pragma: no cover +): """ Render the dictset as a ASCII table. @@ -110,8 +113,13 @@ def ascii_table( Returns: string (ASCII table) """ - from opteryx.third_party.pyarrow_ops import ops + if table == set(): # pragma: no cover + return "No data in table" + if table.num_rows == 0: + return "No data in table" + + # get the width of the display if isinstance(display_width, bool): if not display_width: display_width = 5000 @@ -120,4 +128,96 @@ def ascii_table( display_width = shutil.get_terminal_size((80, 20))[0] - 5 - return "\n".join(ops.head(table, limit, display_width, colorize)) + # Extract head data + if limit > 0: + t = table.slice(length=limit) + else: + t = table + + def type_formatter(value, width): + + punc: str = "\033[38;5;240m" + + if value is None or isinstance(value, bool): + return "{CYAN}" + str(value).ljust(width)[:width] + "{OFF}" + if isinstance(value, int): + return "\033[38;5;203m" + str(value).rjust(width)[:width] + "\033[0m" + if isinstance(value, float): + return "\033[38;5;203m" + str(value).rjust(width)[:width] + "\033[0m" + if isinstance(value, str): + return "{YELLOW}" + str(value).ljust(width)[:width] + "{OFF}" + if isinstance(value, datetime.datetime): + return f"\033[38;5;72m{value.strftime('%Y-%m-%d')} \033[38;5;150m{value.strftime('%H:%M:%S')}\033[0m" + if isinstance(value, list): + value = ( + punc + + "['\033[38;5;26m" + + f"{punc}', '\033[38;5;26m".join(value) + + punc + + "']\033[0m" + ) + return trunc_printable(value, width) + if isinstance(value, dict): + value = ( + punc + + "{" + + f"{punc}, ".join( + f"'\033[38;5;26m{k}{punc}':'\033[38;5;170m{v}{punc}'" + for k, v in value.items() + ) + + punc + + "}\033[0m" + ) + return trunc_printable(value, width) + return str(value).ljust(width)[:width] + + def trunc_printable(value, width): + + offset = 0 + emit = "" + ignoring = False + + for char in value: + emit += char + if char == "\033": + ignoring = True + if not ignoring: + offset += 1 + if ignoring and char == "m": + ignoring = False + if not ignoring and offset >= width: + return emit + "\033[0m" + return emit + "\033[0m" + " " * (width - offset) + + def _inner(): + + head = t.to_pydict() + + # Calculate width + col_width = list(map(len, head.keys())) + data_width = [max(map(len, map(str, h))) for h in head.values()] + + col_width = [ + min(max(cw, dw), max_column_width) for cw, dw in zip(col_width, data_width) + ] + + # Print data + data = [[head[c][i] for c in head.keys()] for i in range(t.num_rows)] + yield ("┌───────┬─" + "─┬─".join("─" * cw for cw in col_width) + "─┐") + yield ( + "│ Row │ " + + " │ ".join(v.ljust(w) for v, w in zip(head.keys(), col_width)) + + " │" + ) + yield ("╞═══════╪═" + "═╪═".join("═" * cw for cw in col_width) + "═╡") + for i in range(len(data)): + formatted = [type_formatter(v, w) for v, w in zip(data[i], col_width)] + yield ("│ " + str(i).ljust(5) + " │ " + " │ ".join(formatted) + " │") + yield ("└───────┴─" + "─┴─".join("─" * cw for cw in col_width) + "─┘") + + from opteryx.utils import colors + + return "\n".join( + trunc_printable(colors.colorize(line), display_width) + colors.colorize("{OFF}") + for line in _inner() + )