Skip to content

Commit

Permalink
typing: type hint with standard collections
Browse files Browse the repository at this point in the history
Signed-off-by: ZhengYu, Xu <[email protected]>
  • Loading branch information
zen-xu committed Apr 17, 2024
1 parent fa3164c commit bdcdba4
Showing 1 changed file with 22 additions and 20 deletions.
42 changes: 22 additions & 20 deletions connectorx-python/connectorx/__init__.py
Original file line number Diff line number Diff line change
@@ -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,
Expand Down Expand Up @@ -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.
Expand All @@ -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)
Expand All @@ -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
Expand Down Expand Up @@ -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.
Expand Down Expand Up @@ -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.
Expand Down Expand Up @@ -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
Expand All @@ -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"]
Expand Down

0 comments on commit bdcdba4

Please sign in to comment.