Skip to content

typing: support type hint with standard collections #605

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
wangxiaoying merged 1 commit into from
Apr 17, 2024
Merged

typing: support type hint with standard collections #605

Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
42 changes: 22 additions & 20 deletions connectorx-python/connectorx/__init__.py
View file
Open in desktop
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