Skip to content

Fixup typing of dtypes #321

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.

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
merged 2 commits into from
Nov 14, 2023
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
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
63 changes: 27 additions & 36 deletions spec/API_specification/dataframe_api/typing.py
Original file line number Diff line number Diff line change
Expand Up @@ -17,41 +17,6 @@
if TYPE_CHECKING:
from collections.abc import Sequence

from .dtypes import (
Bool,
Date,
Datetime,
Duration,
Float32,
Float64,
Int8,
Int16,
Int32,
Int64,
String,
UInt8,
UInt16,
UInt32,
UInt64,
)

DType = Union[
Bool,
Float64,
Float32,
Int64,
Int32,
Int16,
Int8,
UInt64,
UInt32,
UInt16,
UInt8,
String,
Date,
Datetime,
Duration,
]

# Type alias: Mypy needs Any, but for readability we need to make clear this
# is a Python scalar (i.e., an instance of `bool`, `int`, `float`, `str`, etc.)
Expand Down Expand Up @@ -104,7 +69,14 @@ class Datetime:
def __init__( # noqa: ANN204
self,
time_unit: Literal["ms", "us"],
time_zone: str | None,
time_zone: str | None = None,
):
...

class Duration:
def __init__( # noqa: ANN204
self,
time_unit: Literal["ms", "us"],
):
...

Expand Down Expand Up @@ -155,6 +127,25 @@ def date(self, year: int, month: int, day: int) -> Scalar:
...


DType = Union[
Namespace.Bool,
Namespace.Float64,
Namespace.Float32,
Namespace.Int64,
Namespace.Int32,
Namespace.Int16,
Namespace.Int8,
Namespace.UInt64,
Namespace.UInt32,
Namespace.UInt16,
Namespace.UInt8,
Namespace.String,
Namespace.Date,
Namespace.Datetime,
Namespace.Duration,
]


class SupportsDataFrameAPI(Protocol):
def __dataframe_consortium_standard__(
self,
Expand Down
27 changes: 27 additions & 0 deletions spec/API_specification/examples/04_datatypes.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,27 @@
from __future__ import annotations

from typing import TYPE_CHECKING, Any, Callable

if TYPE_CHECKING:
from dataframe_api.typing import SupportsDataFrameAPI

some_array_function: Callable[[Any], Any]


def main(df_raw: SupportsDataFrameAPI) -> SupportsDataFrameAPI:
df = df_raw.__dataframe_consortium_standard__(api_version="2023-11.beta").persist()
namespace = df.__dataframe_namespace__()
df = df.select(
*[
col_name
for col_name in df.column_names
if isinstance(df.col(col_name).dtype, namespace.Int64)
],
)
arr = df.to_array(namespace.Int64())
arr = some_array_function(arr)
df = namespace.dataframe_from_2d_array(
arr,
schema={"a": df.col("a").dtype, "b": namespace.Float64()},
)
return df.dataframe