From b38f78eb34b703645452e699298bbca5458bbbec Mon Sep 17 00:00:00 2001 From: "ZhengYu, Xu" Date: Wed, 17 Apr 2024 15:47:36 +0800 Subject: [PATCH] typing: add stubs for connectorx.so Signed-off-by: ZhengYu, Xu --- connectorx-python/connectorx/connectorx.pyi | 34 +++++++++++++++++++++ connectorx-python/connectorx/py.typed | 0 2 files changed, 34 insertions(+) create mode 100644 connectorx-python/connectorx/connectorx.pyi create mode 100644 connectorx-python/connectorx/py.typed diff --git a/connectorx-python/connectorx/connectorx.pyi b/connectorx-python/connectorx/connectorx.pyi new file mode 100644 index 0000000000..b556d918b1 --- /dev/null +++ b/connectorx-python/connectorx/connectorx.pyi @@ -0,0 +1,34 @@ +from __future__ import annotations + +from typing import overload, Literal, Any, TypeAlias +import pandas as pd + +_ArrowArrayPtr: TypeAlias = int +_ArrowSchemaPtr: TypeAlias = int +_Column: TypeAlias = str + +@overload +def read_sql( + conn: str, + return_type: Literal["pandas"], + protocol: str | None, + queries: list[str] | None, + partition_query: dict[str, Any] | None, +) -> pd.DataFrame: ... +@overload +def read_sql( + conn: str, + return_type: Literal["arrow", "arrow2"], + protocol: str | None, + queries: list[str] | None, + partition_query: dict[str, Any] | None, +) -> tuple[list[_Column], list[list[tuple[_ArrowArrayPtr, _ArrowSchemaPtr]]]]: ... +def partition_sql(conn: str, partition_query: dict[str, Any]) -> list[str]: ... +def read_sql2( + sql: str, db_map: dict[str, str] +) -> tuple[list[_Column], list[list[tuple[_ArrowArrayPtr, _ArrowSchemaPtr]]]]: ... +def get_meta( + conn: str, + protocol: Literal["csv", "binary", "cursor", "simple", "text"] | None, + query: str, +) -> dict[str, Any]: ... diff --git a/connectorx-python/connectorx/py.typed b/connectorx-python/connectorx/py.typed new file mode 100644 index 0000000000..e69de29bb2