Bump fastdatatable for tz fix and key bindings (#384)

* fix: add pytz for older pythons

* fix: bump fastdatatable for tz fix and key bindings

* fix: update help screen snapshot

* feat: download tzdata on windows

* fix: blacken conftest

.pre-commit-config.yaml

@@ -24,6 +24,7 @@ repos:
           - rich-click>=1.7.1
           - questionary
           - tomlkit
+          - pandas-stubs
         args:
           - "--disallow-untyped-calls"
           - "--disallow-untyped-defs"

.python-version

@@ -1 +1 @@
-3.10.6
+3.8.10

CHANGELOG.md

@@ -4,6 +4,16 @@ All notable changes to this project will be documented in this file.
 
 ## [Unreleased]
 
+### Features
+
+-   Improves keyboard navigation of the Results Viewer by adding key bindings, including <kbd>ctrl+right/left/up/down/home/end</kbd>, <kbd>tab</kbd>, and <kbd>ctrl+a</kbd>.
+-   The Trino adapter is now installable as an extra; use `pip install harlequin[trino]`.
+-   Harlequin will automatically download a missing timezone database on Windows. Prevent this behavior with `--no-download-tzdata`.
+
+### Bug Fixes
+
+-   Fixes a crash when selecting data from a timestamptz field ([#382](https://github.com/tconbeer/harlequin/issues/382)) (or another field with an invalid Arrow data type).
+
 ## [1.8.0] - 2023-12-21
 
 ### Features

pyproject.toml

@@ -22,7 +22,7 @@ python = ">=3.8.1,<4.0.0"
 
 # textual and component libraries
 textual = "==0.46.0"
-textual-fastdatatable = "==0.5.0"
+textual-fastdatatable = "==0.6.2"
 textual-textarea = "==0.9.5"
 
 # click
@@ -55,6 +55,7 @@ black = "^23.3.0"
 ruff = ">=0.0.285"
 mypy = "^1.2.0"
 types-pygments = "^2.16.0.0"
+pandas-stubs = "^1"
 
 [tool.poetry.group.test.dependencies]
 pytest = "^7.3.1"

src/harlequin/cli.py

@@ -11,8 +11,13 @@
 from harlequin.colors import GREEN, PINK, PURPLE, YELLOW
 from harlequin.config import get_config_for_profile
 from harlequin.config_wizard import wizard
-from harlequin.exception import HarlequinConfigError, pretty_print_error
+from harlequin.exception import (
+    HarlequinConfigError,
+    HarlequinTzDataError,
+    pretty_print_error,
+)
 from harlequin.plugins import load_plugins
+from harlequin.windows_timezone import check_and_install_tzdata
 
 if sys.version_info < (3, 10):
     from importlib_metadata import entry_points, version
@@ -67,6 +72,7 @@
                 "--theme",
                 "--limit",
                 "--config",
+                "--no-download-tzdata",
                 "--version",
                 "--help",
             ],
@@ -190,6 +196,14 @@ def build_cli() -> click.Command:
         expose_value=True,
         is_eager=True,
     )
+    @click.option(
+        "--no-download-tzdata",
+        help=(
+            "(Windows Only) Prevent Harlequin from downloading an IANA timezone "
+            "database, even if one is missing. May cause undesired behavior."
+        ),
+        is_flag=True,
+    )
     @click.pass_context
     def inner_cli(
         ctx: click.Context,
@@ -226,6 +240,14 @@ def inner_cli(
         # merge the config and the cli options
         config.update(kwargs)
 
+        # detect and install (if necessary) a tzdatabase on Windows
+        if sys.platform == "win32" and not config.pop("no_download_tzdata", None):
+            try:
+                check_and_install_tzdata()
+            except HarlequinTzDataError as e:
+                pretty_print_error(e)
+                ctx.exit(2)
+
         # load and instantiate the adapter
         adapter = config.pop("adapter", DEFAULT_ADAPTER)
         conn_str: Sequence[str] = config.pop("conn_str", tuple())  # type: ignore

src/harlequin/components/help_screen.md

@@ -70,8 +70,14 @@
 #### Moving the Cursor
 
 - up,down,left,right: Move the cursor one cell.
-- home: Move the cursor to the top of the current column.
-- end: Move the cursor to the bottom of the current column.
+- tab: Move the cursor to the next cell in the table.
+- shift+tab: Move the cursor to the previous cell in the table.
+- home, ctrl+up: Move the cursor to the top of the current column.
+- end, ctrl+down: Move the cursor to the bottom of the current column.
+- ctrl+left: Move the cursor to the start of the current row.
+- ctrl+right: Move the cursor to the end of the current row.
+- ctrl+home: Move the cursor to the start of the table.
+- ctrl+end: Move the cursor to the end of the table.
 - PgUp: Move the cursor up one screen.
 - PgDn: Move the cursor down one screen.
 - shift+[any]: Select range while moving the cursor.

src/harlequin/exception.py

@@ -33,6 +33,10 @@ class HarlequinWizardError(HarlequinError):
     pass
 
 
+class HarlequinTzDataError(HarlequinError):
+    pass
+
+
 def pretty_print_error(error: HarlequinError) -> None:
     from rich import print
     from rich.panel import Panel

src/harlequin/windows_timezone.py

@@ -0,0 +1,80 @@
+import tarfile
+from datetime import datetime
+from pathlib import Path
+from tempfile import TemporaryDirectory
+from urllib.request import urlopen
+
+import platformdirs
+import pyarrow as pa
+import pyarrow.compute as pc
+import pyarrow.lib as pl
+from rich import print
+from rich.panel import Panel
+
+from harlequin.colors import GREEN
+from harlequin.exception import HarlequinTzDataError
+
+HARLEQUIN_TZ_DATA_PATH = platformdirs.user_data_path(appname="harlequin") / "tzdata"
+
+
+def check_and_install_tzdata() -> None:
+    """
+    On Windows, Arrow expects to find a timezone database in the User's
+    Downloads folder. We check to see if it can find one there,
+    and if not, we override the tz db search path to a Harlequin-specific
+    location. If it still can't find one, we download the DB to that
+    harlequin-specific location.
+    """
+    try:
+        pc.assume_timezone(datetime(2024, 1, 1), "America/New_York")
+    except pl.ArrowInvalid:
+        # no tz database in the default location; try the harlequin location
+        try:
+            pa.set_timezone_db_path(str(HARLEQUIN_TZ_DATA_PATH))
+            pc.assume_timezone(datetime(2024, 1, 1), "America/New_York")
+        except (OSError, pl.ArrowInvalid):
+            message = (
+                "Harlequin could not find a timezone database, which it needs "
+                "to support Apache Arrow's timestamptz features. It is downloading one "
+                "now (it will only do this once). For more info, see "
+                "[link]https://harlequin.sh/docs/troubleshooting/timezone-windows[/]"
+            )
+            print(
+                Panel.fit(
+                    message,
+                    title="Harlequin Timezone Support",
+                    title_align="left",
+                    border_style=GREEN,
+                )
+            )
+            try:
+                # tz database is missing. We need to install it.
+                response = urlopen(
+                    "https://www.iana.org/time-zones/repository/tzdata-latest.tar.gz"
+                )
+                with TemporaryDirectory() as tmpdir:
+                    tar_path = Path(tmpdir) / "tzdata.tar.gz"
+                    with tar_path.open("wb") as f:
+                        f.write(response.read())
+                    tarfile.open(tar_path).extractall(HARLEQUIN_TZ_DATA_PATH)
+                zone_response = urlopen(
+                    "https://raw.githubusercontent.com/unicode-org/cldr/main/common/"
+                    "supplemental/windowsZones.xml"
+                )
+                zone_target = HARLEQUIN_TZ_DATA_PATH / "windowsZones.xml"
+                with zone_target.open("wb") as f:
+                    f.write(zone_response.read())
+                pa.set_timezone_db_path(str(HARLEQUIN_TZ_DATA_PATH))
+            except Exception as e:
+                err_msg = (
+                    "Harlequin was not able to download a timezone database. Without "
+                    "a timezone database, Harlequin may crash if you attempt to load "
+                    "timestamptz values into the results viewer. To use Harlequin "
+                    "anyway, set the --no-download-tzdata option.\n"
+                    "For more info, see "
+                    "[link]https://harlequin.sh/docs/troubleshooting/timezone-windows[/]"
+                    f"\nDownload failed with the following exception:\n{e}"
+                )
+                raise HarlequinTzDataError(
+                    msg=err_msg, title="Harlequin Timezone Error"
+                ) from e

stubs/pyarrow/__init__.pyi

@@ -1,10 +1,39 @@
 from __future__ import annotations
 
-from typing import Any, Iterable, Iterator, Literal, Mapping, Type, TypeVar
+from typing import Any, Iterable, Iterator, Literal, Mapping, Sequence, Type, TypeVar
+
+import pandas as pd
 
 from .compute import CastOptions
-from .types import DataType as DataType
-from .types import string as string
+
+class DataType: ...
+
+def string() -> DataType: ...
+def null() -> DataType: ...
+def bool_() -> DataType: ...
+def int8() -> DataType: ...
+def int16() -> DataType: ...
+def int32() -> DataType: ...
+def int64() -> DataType: ...
+def uint8() -> DataType: ...
+def uint16() -> DataType: ...
+def uint32() -> DataType: ...
+def uint64() -> DataType: ...
+def float16() -> DataType: ...
+def float32() -> DataType: ...
+def float64() -> DataType: ...
+def date32() -> DataType: ...
+def date64() -> DataType: ...
+def binary(length: int = -1) -> DataType: ...
+def large_binary() -> DataType: ...
+def large_string() -> DataType: ...
+def month_day_nano_interval() -> DataType: ...
+def time32(unit: Literal["s", "ms", "us", "ns"]) -> DataType: ...
+def time64(unit: Literal["s", "ms", "us", "ns"]) -> DataType: ...
+def timestamp(
+    unit: Literal["s", "ms", "us", "ns"], tz: str | None = None
+) -> DataType: ...
+def duration(unit: Literal["s", "ms", "us", "ns"]) -> DataType: ...
 
 class MemoryPool: ...
 class Schema: ...
@@ -26,6 +55,7 @@ class _PandasConvertible:
         options: CastOptions | None = None,
     ) -> A: ...
     def __getitem__(self, index: int) -> Scalar: ...
+    def __iter__(self) -> Any: ...
     def to_pylist(self) -> list[Any]: ...
     def fill_null(self: A, fill_value: Any) -> A: ...
 
@@ -69,20 +99,23 @@ class _Tabular:
         self: T, field_: str | Field, column: Array | ChunkedArray
     ) -> T: ...
     def column(self, i: int | str) -> _PandasConvertible: ...
+    def equals(self: T, other: T, check_metadata: bool = False) -> bool: ...
     def itercolumns(self) -> Iterator[_PandasConvertible]: ...
+    def rename_columns(self: T, names: list[str]) -> T: ...
+    def select(self: T, columns: Sequence[str | int]) -> T: ...
     def set_column(
         self: T, i: int, field_: str | Field, column: Array | ChunkedArray
     ) -> T: ...
-    def sort_by(
-        self: T,
-        sorting: str | list[tuple[str, Literal["ascending", "descending"]]],
-        **kwargs: Any,
-    ) -> T: ...
     def slice(  # noqa: A003
         self: T,
         offset: int = 0,
         length: int | None = None,
     ) -> T: ...
+    def sort_by(
+        self: T,
+        sorting: str | list[tuple[str, Literal["ascending", "descending"]]],
+        **kwargs: Any,
+    ) -> T: ...
     def to_pylist(self) -> list[dict[str, Any]]: ...
 
 class RecordBatch(_Tabular): ...
@@ -105,11 +138,21 @@ def array(
     safe: bool = True,
     memory_pool: MemoryPool | None = None,
 ) -> Array | ChunkedArray: ...
+def concat_arrays(
+    arrays: Iterable[Array], memory_pool: MemoryPool | None = None
+) -> Array: ...
 def nulls(
     size: int,
     type: DataType | None = None,  # noqa: A002
     memory_pool: MemoryPool | None = None,
 ) -> Array: ...
-def concat_arrays(
-    arrays: Iterable[Array], memory_pool: MemoryPool | None = None
-) -> Array: ...
+def table(
+    data: pd.DataFrame
+    | Mapping[str, _PandasConvertible | list]
+    | list[_PandasConvertible],
+    names: list[str] | None = None,
+    schema: Schema | None = None,
+    metadata: Mapping | None = None,
+    nthreads: int | None = None,
+) -> Table: ...
+def set_timezone_db_path(path: str) -> None: ...

stubs/pyarrow/compute.pyi

@@ -1,7 +1,9 @@
 from __future__ import annotations
 
-from . import MemoryPool, Scalar, _PandasConvertible
-from .types import DataType
+from datetime import datetime
+from typing import Any, Callable, Literal
+
+from . import DataType, MemoryPool, Scalar, _PandasConvertible
 
 class Expression: ...
 class ScalarAggregateOptions: ...
@@ -30,3 +32,24 @@ def max(  # noqa: A001
 def utf8_length(
     strings: _PandasConvertible, /, *, memory_pool: MemoryPool | None = None
 ) -> _PandasConvertible: ...
+def register_scalar_function(
+    func: Callable,
+    function_name: str,
+    function_doc: dict[Literal["summary", "description"], str],
+    in_types: dict[str, DataType],
+    out_type: DataType,
+    func_registry: Any | None = None,
+) -> None: ...
+def call_function(
+    function_name: str, target: list[_PandasConvertible]
+) -> _PandasConvertible: ...
+def assume_timezone(
+    timestamps: _PandasConvertible | Scalar | datetime,
+    /,
+    timezone: str,
+    *,
+    ambiguous: Literal["raise", "earliest", "latest"] = "raise",
+    nonexistent: Literal["raise", "earliest", "latest"] = "raise",
+    options: Any | None = None,
+    memory_pool: MemoryPool | None = None,
+) -> _PandasConvertible: ...

stubs/pyarrow/lib.pyi

@@ -1 +1,3 @@
 class ArrowNotImplementedError(Exception): ...
+class ArrowInvalid(Exception): ...
+class ArrowKeyError(Exception): ...

tests/conftest.py

@@ -9,13 +9,20 @@
 import pytest
 from harlequin import Harlequin
 from harlequin.adapter import HarlequinAdapter
+from harlequin.windows_timezone import check_and_install_tzdata
 
 if sys.version_info < (3, 10):
     from importlib_metadata import entry_points
 else:
     from importlib.metadata import entry_points
 
 
+@pytest.fixture(scope="session", autouse=True)
+def install_tzdata() -> None:
+    if sys.platform == "win32":
+        check_and_install_tzdata()
+
+
 @pytest.fixture
 def data_dir() -> Path:
     here = Path(__file__)