diff --git a/connectorx-python/connectorx/__init__.py b/connectorx-python/connectorx/__init__.py index b58590e80d..2d354c642b 100644 --- a/connectorx-python/connectorx/__init__.py +++ b/connectorx-python/connectorx/__init__.py @@ -1,4 +1,6 @@ -from typing import Optional, Tuple, Union, List, Dict, Any +from __future__ import annotations + +from typing import Any from .connectorx import ( read_sql as _read_sql, @@ -35,7 +37,7 @@ ) -def rewrite_conn(conn: str, protocol: Optional[str] = None): +def rewrite_conn(conn: str, protocol: str | None = None): if not protocol: # note: redshift/clickhouse are not compatible with the 'binary' protocol, and use other database # drivers to connect. set a compatible protocol and masquerade as the appropriate backend. @@ -54,7 +56,7 @@ def rewrite_conn(conn: str, protocol: Optional[str] = None): def get_meta( conn: str, query: str, - protocol: Optional[str] = None, + protocol: str | None = None, ): """ Get metadata (header) of the given query (only for pandas) @@ -81,7 +83,7 @@ def partition_sql( query: str, partition_on: str, partition_num: int, - partition_range: Optional[Tuple[int, int]] = None, + partition_range: tuple[int, int] | None = None, ): """ Partition the sql query @@ -110,13 +112,13 @@ def partition_sql( def read_sql_pandas( - sql: Union[List[str], str], - con: Union[str, Dict[str, str]], - index_col: Optional[str] = None, - protocol: Optional[str] = None, - partition_on: Optional[str] = None, - partition_range: Optional[Tuple[int, int]] = None, - partition_num: Optional[int] = None, + sql: list[str] | str, + con: str | dict[str, str], + index_col: str | None = None, + protocol: str | None = None, + partition_on: str | None = None, + partition_range: tuple[int, int] | None = None, + partition_num: int | None = None, ): """ Run the SQL query, download the data from database into a dataframe. @@ -150,15 +152,15 @@ def read_sql_pandas( def read_sql( - conn: Union[str, Dict[str, str]], - query: Union[List[str], str], + conn: str | dict[str, str], + query: list[str] | str, *, return_type: str = "pandas", - protocol: Optional[str] = None, - partition_on: Optional[str] = None, - partition_range: Optional[Tuple[int, int]] = None, - partition_num: Optional[int] = None, - index_col: Optional[str] = None, + protocol: str | None = None, + partition_on: str | None = None, + partition_range: tuple[int, int] | None = None, + partition_num: int | None = None, + index_col: str | None = None, ): """ Run the SQL query, download the data from database into a dataframe. @@ -327,7 +329,7 @@ def read_sql( return df -def reconstruct_arrow(result: Tuple[List[str], List[List[Tuple[int, int]]]]): +def reconstruct_arrow(result: tuple[list[str], list[list[tuple[int, int]]]]): import pyarrow as pa names, ptrs = result @@ -343,7 +345,7 @@ def reconstruct_arrow(result: Tuple[List[str], List[List[Tuple[int, int]]]]): return pa.Table.from_batches(rbs) -def reconstruct_pandas(df_infos: Dict[str, Any]): +def reconstruct_pandas(df_infos: dict[str, Any]): import pandas as pd data = df_infos["data"]