feat(timestamp): allow passing nullable flag to ibis.timestamp

Author: NickCrews

ibis/common/temporal.py

@@ -208,7 +208,7 @@ def normalize_timezone(tz):
 
 
 @lazy_singledispatch
-def normalize_datetime(value):
+def normalize_datetime(value) -> datetime.datetime:
     raise TypeError(f"Unable to normalize {type(value)} to timestamp")
 
 

ibis/expr/api.py

@@ -804,11 +804,17 @@ def timestamp(
     second: int | ir.IntegerValue | Deferred,
     /,
     timezone: str | None = None,
+    nullable: bool = True,
 ) -> TimestampValue: ...
 
 
 @overload
-def timestamp(value_or_year: Any, /, timezone: str | None = None) -> TimestampValue: ...
+def timestamp(
+    value_or_year: str | datetime.datetime,
+    /,
+    timezone: str | None = None,
+    nullable: bool = True,
+) -> TimestampValue: ...
 
 
 @deferrable
@@ -821,6 +827,7 @@ def timestamp(
     second=None,
     /,
     timezone=None,
+    nullable: bool = True,
 ):
     """Construct a timestamp scalar or column.
 
@@ -841,6 +848,8 @@ def timestamp(
         The timestamp second component; required if `value_or_year` is a year.
     timezone
         The timezone name, or none for a timezone-naive timestamp.
+    nullable
+        Whether the resulting timestamp should be nullable. Defaults to True.
 
     Returns
     -------
@@ -888,16 +897,17 @@ def timestamp(
             raise NotImplementedError(
                 "Timezone currently not supported when creating a timestamp from components"
             )
-        return ops.TimestampFromYMDHMS(*args).to_expr()
+        return ops.TimestampFromYMDHMS(*args, nullable=nullable).to_expr()
     elif isinstance(value_or_year, (numbers.Real, ir.IntegerValue)):
         raise TypeError("Use ibis.literal(...).as_timestamp() instead")
     elif isinstance(value_or_year, ir.Expr):
-        return value_or_year.cast(dt.Timestamp(timezone=timezone))
+        return value_or_year.cast(dt.Timestamp(timezone=timezone, nullable=nullable))
     else:
         value = normalize_datetime(value_or_year)
         tzinfo = normalize_timezone(timezone or value.tzinfo)
         timezone = tzinfo.tzname(value) if tzinfo is not None else None
-        return literal(value, type=dt.Timestamp(timezone=timezone))
+        dtype = dt.Timestamp.from_datetime(value, timezone=timezone, nullable=nullable)
+        return literal(value, type=dtype)
 
 
 @overload

ibis/expr/datatypes/cast.py

@@ -142,6 +142,10 @@ def higher_precedence(left: dt.DataType, right: dt.DataType) -> dt.DataType:
 @public
 def highest_precedence(dtypes: Iterator[dt.DataType]) -> dt.DataType:
     """Compute the highest precedence of `dtypes`."""
+    # TODO: currently,
+    # highest_precedence([dt.Timestamp(scale=3), dt.Timestamp(timezone="UTC")])
+    # returns dt.Timestamp(timezone="UTC").
+    # Perhaps it should return dt.Timestamp(scale=3, timezone="UTC") instead.
     if collected := list(dtypes):
         return functools.reduce(higher_precedence, collected)
     else:

ibis/expr/datatypes/core.py

@@ -33,6 +33,7 @@
 
 if TYPE_CHECKING:
     import numpy as np
+    import pandas as pd
     import polars as pl
     import pyarrow as pa
     from pandas.api.extensions import ExtensionDtype
@@ -83,6 +84,62 @@ def dtype(
 ) -> DataType: ...
 
 
+if TYPE_CHECKING:
+    import numpy as np
+    import polars as pl
+    import pyarrow as pa
+    from pandas.api.extensions import ExtensionDtype
+
+
+@overload
+def dtype(value: type[int] | Literal["int"], nullable: bool = True) -> Int64: ...
+@overload
+def dtype(
+    value: type[str] | Literal["str", "string"], nullable: bool = True
+) -> String: ...
+@overload
+def dtype(
+    value: type[bool] | Literal["bool", "boolean"], nullable: bool = True
+) -> Boolean: ...
+@overload
+def dtype(value: type[bytes] | Literal["bytes"], nullable: bool = True) -> Binary: ...
+@overload
+def dtype(value: type[Real] | Literal["float"], nullable: bool = True) -> Float64: ...
+@overload
+def dtype(
+    value: type[pydecimal.Decimal] | Literal["decimal"], nullable: bool = True
+) -> Decimal: ...
+@overload
+def dtype(
+    value: type[pydatetime.datetime] | Literal["timestamp"], nullable: bool = True
+) -> Timestamp: ...
+@overload
+def dtype(
+    value: type[pydatetime.date] | Literal["date"], nullable: bool = True
+) -> Date: ...
+@overload
+def dtype(
+    value: type[pydatetime.time] | Literal["time"], nullable: bool = True
+) -> Time: ...
+@overload
+def dtype(
+    value: type[pydatetime.timedelta] | Literal["interval"], nullable: bool = True
+) -> Interval: ...
+@overload
+def dtype(
+    value: type[pyuuid.UUID] | Literal["uuid"], nullable: bool = True
+) -> UUID: ...
+@overload
+def dtype(
+    value: DataType | str | np.dtype | ExtensionDtype | pl.DataType | pa.DataType,
+    nullable: bool = True,
+) -> DataType: ...
+
+
+if TYPE_CHECKING:
+    import pandas as pd
+
+
 @lazy_singledispatch
 def dtype(value, nullable=True) -> DataType:
     """Create a DataType object.
@@ -697,6 +754,34 @@ def from_unit(cls, unit, timezone=None, nullable=True) -> Self:
             raise ValueError(f"Invalid unit {unit}")
         return cls(scale=scale, timezone=timezone, nullable=nullable)
 
+    @classmethod
+    def from_datetime(
+        cls, dt: pydatetime.datetime, timezone: str | None = None, nullable: bool = True
+    ):
+        """Infer from a python datetime.datetime object."""
+        if dt.microsecond:
+            scale = 6
+        else:
+            scale = 0
+        if timezone is None and dt.tzinfo is not None:
+            timezone = str(dt.tzinfo)
+        return cls(scale=scale, timezone=timezone, nullable=nullable)
+
+    @classmethod
+    def from_pandas(
+        cls, value: pd.Timestamp, timezone: str | None = None, nullable: bool = True
+    ):
+        """Infer from a pandas.Timestamp."""
+        if value.nanosecond:
+            scale = 9
+        elif value.microsecond:
+            scale = 6
+        else:
+            scale = 0
+        if timezone is None and value.tz is not None:
+            timezone = str(value.tz)
+        return cls(timezone=timezone, scale=scale, nullable=nullable)
+
     @property
     def unit(self) -> str:
         """Return the unit of the timestamp."""

ibis/expr/datatypes/value.py

@@ -10,7 +10,7 @@
 from collections.abc import Mapping, Sequence
 from functools import partial
 from operator import attrgetter
-from typing import Any
+from typing import TYPE_CHECKING, Any
 
 import toolz
 from public import public
@@ -28,6 +28,9 @@
 )
 from ibis.expr.datatypes.cast import highest_precedence
 
+if TYPE_CHECKING:
+    import pandas as pd
+
 
 @lazy_singledispatch
 def infer(value: Any) -> dt.DataType:
@@ -82,10 +85,7 @@ def infer_date(value: datetime.date) -> dt.Date:
 
 @infer.register(datetime.datetime)
 def infer_timestamp(value: datetime.datetime) -> dt.Timestamp:
-    if value.tzinfo:
-        return dt.Timestamp(timezone=str(value.tzinfo))
-    else:
-        return dt.timestamp
+    return dt.Timestamp.from_datetime(value)
 
 
 @infer.register(datetime.timedelta)
@@ -171,11 +171,8 @@ def infer_numpy_scalar(value):
 
 
 @infer.register("pandas.Timestamp")
-def infer_pandas_timestamp(value):
-    if value.tz is not None:
-        return dt.Timestamp(timezone=str(value.tz))
-    else:
-        return dt.timestamp
+def infer_pandas_timestamp(value: pd.Timestamp) -> dt.Timestamp:
+    return dt.Timestamp.from_pandas(value)
 
 
 @infer.register("pandas.Timedelta")

ibis/expr/operations/temporal.py

@@ -258,9 +258,17 @@ class TimestampFromYMDHMS(Value):
     hours: Value[dt.Integer]
     minutes: Value[dt.Integer]
     seconds: Value[dt.Integer]
+    nullable: bool = True
 
-    dtype = dt.timestamp
-    shape = rlz.shape_like("args")
+    shape = rlz.shape_like(["year", "month", "day", "hours", "minutes", "seconds"])
+
+    @attribute
+    def dtype(self):
+        return dt.Timestamp(
+            nullable=self.nullable,
+            timezone=None,
+            scale=0,  # second precision
+        )
 
 
 @public

ibis/expr/rules.py

@@ -1,10 +1,11 @@
 from __future__ import annotations
 
 from itertools import product, starmap
-from typing import Optional
+from typing import TYPE_CHECKING, Any, Optional
 
 from public import public
 
+import ibis.expr.datashape as ds
 import ibis.expr.datatypes as dt
 import ibis.expr.operations as ops
 from ibis import util
@@ -13,14 +14,17 @@
 from ibis.common.patterns import CoercionError, NoMatch, Pattern
 from ibis.common.temporal import IntervalUnit
 
+if TYPE_CHECKING:
+    from collections.abc import Iterable
+
 
 @public
-def highest_precedence_shape(nodes):
+def highest_precedence_shape(nodes: Iterable[ops.Value]) -> ds.DataShape:
     return max(node.shape for node in nodes)
 
 
 @public
-def highest_precedence_dtype(nodes):
+def highest_precedence_dtype(nodes: Iterable[ops.Value]) -> dt.DataType:
     """Return the highest precedence type from the passed expressions.
 
     Also verifies that there are valid implicit casts between any of the types
@@ -29,8 +33,8 @@ def highest_precedence_dtype(nodes):
 
     Parameters
     ----------
-    nodes : Iterable[ops.Value]
-      A sequence of Expressions
+    nodes
+      An Iterable of Expressions
 
     Returns
     -------
@@ -42,7 +46,7 @@ def highest_precedence_dtype(nodes):
 
 
 @public
-def castable(source, target):
+def castable(source: ops.Value, target: ops.Value) -> bool:
     """Return whether source ir type is implicitly castable to target.
 
     Based on the underlying datatypes and the value in case of Literals
@@ -52,7 +56,7 @@ def castable(source, target):
 
 
 @public
-def comparable(left, right):
+def comparable(left: ops.Value, right: ops.Value) -> bool:
     return castable(left, right) or castable(right, left)
 
 
@@ -61,28 +65,31 @@ def comparable(left, right):
 
 
 @public
-def dtype_like(name):
+def dtype_like(names: str | Iterable[str]):
     @attribute
-    def dtype(self):
-        args = getattr(self, name)
-        args = args if util.is_iterable(args) else [args]
-        return highest_precedence_dtype(args)
+    def dtype(self) -> dt.DataType:
+        return highest_precedence_dtype(_attributes(self, names))
 
     return dtype
 
 
 @public
-def shape_like(name):
+def shape_like(name: str | Iterable[str]):
     @attribute
-    def shape(self):
-        args = getattr(self, name)
-        args = args if util.is_iterable(args) else [args]
+    def shape(self) -> ds.DataShape:
+        args = _attributes(self, name)
         args = [a for a in args if a is not None]
         return highest_precedence_shape(args)
 
     return shape
 
 
+def _attributes(obj: Any, names: str | Iterable[str]) -> tuple:
+    if isinstance(names, str):
+        return util.promote_tuple(getattr(obj, names))
+    return tuple(getattr(obj, name) for name in names)
+
+
 # TODO(kszucs): might just use bounds instead of actual literal values
 # that could simplify interval binop output_type methods
 # TODO(kszucs): pre-generate mapping?