Skip to content

Commit

Permalink
feat(types): add typing overloads to common .cast() dtypes (#10682)
Browse files Browse the repository at this point in the history
  • Loading branch information
NickCrews authored Jan 17, 2025
1 parent c909db3 commit 1394a2a
Showing 1 changed file with 46 additions and 1 deletion.
47 changes: 46 additions & 1 deletion ibis/expr/types/generic.py
Original file line number Diff line number Diff line change
@@ -1,7 +1,7 @@
from __future__ import annotations

from collections.abc import Sequence
from typing import TYPE_CHECKING, Any
from typing import TYPE_CHECKING, Any, Literal, overload

from public import public

Expand All @@ -17,6 +17,9 @@
from ibis.util import deprecated, promote_list

if TYPE_CHECKING:
import datetime
import uuid

import pandas as pd
import polars as pl
import pyarrow as pa
Expand Down Expand Up @@ -135,6 +138,48 @@ def hash(self) -> ir.IntegerValue:
"""
return ops.Hash(self).to_expr()

@overload
def cast(
self, target_type: Literal["string", "str"] | type[str]
) -> ir.StringValue: ...
@overload
def cast(
self,
target_type: Literal[
"int",
"uint",
"int8",
"uint8",
"int16",
"uint16",
"int32",
"uint32",
"int64",
"uint64",
]
| type[int],
) -> ir.IntegerValue: ...
@overload
def cast(
self,
target_type: Literal["float", "float8", "float16", "float32", "float64"]
| type[float],
) -> ir.FloatingValue: ...
@overload
def cast(
self, target_type: Literal["bool", "boolean"] | type[bool]
) -> ir.BooleanValue: ...
@overload
def cast(self, target_type: Literal["date"]) -> ir.DateValue: ...
@overload
def cast(
self, target_type: Literal["datetime", "timestamp"] | type[datetime.datetime]
) -> ir.TimestampValue: ...
@overload
def cast(self, target_type: Literal["time"]) -> ir.TimeValue: ...
@overload
def cast(self, target_type: Literal["uuid"] | type[uuid.UUID]) -> ir.UUIDValue: ...

def cast(self, target_type: Any) -> Value:
"""Cast expression to indicated data type.
Expand Down

0 comments on commit 1394a2a

Please sign in to comment.