Repository: pola-rs/polars
File: py-polars/src/polars/functions/lit.py (Python)

Code before:

    >>> pl.lit(1)  # doctest: +IGNORE_RESULT
    >>> pl.lit(5.5)  # doctest: +IGNORE_RESULT
    >>> pl.lit(None)  # doctest: +IGNORE_RESULT
    >>> pl.lit("foo_bar")  # doctest: +IGNORE_RESULT
    >>> pl.lit(date(2021, 1, 20))  # doctest: +IGNORE_RESULT
    >>> pl.lit(datetime(2023, 3, 31, 10, 30, 45))  # doctest: +IGNORE_RESULT

    Literal list/Series data (1D):

    >>> pl.lit([1, 2, 3])  # doctest: +SKIP
    >>> pl.lit(pl.Series("x", [1, 2, 3]))  # doctest: +IGNORE_RESULT

    Literal list/Series data (2D):

    >>> pl.lit([[1, 2], [3, 4]])  # doctest: +SKIP
    >>> pl.lit(pl.Series("y", [[1, 2], [3, 4]]))  # doctest: +IGNORE_RESULT
    """
    time_unit: TimeUnit

    if isinstance(dtype, BaseExtension):
        return lit(value, dtype.ext_storage()).ext.to(dtype)
    elif isinstance(dtype, type) and issubclass(dtype, BaseExtension):
        msg = f"dtype '{dtype}' is a BaseExtension class, it should be an instance"
        raise TypeError(msg)
    elif dtype and dtype.is_object():
        value_s = pl.Series("literal", [value], dtype=dtype)
        return wrap_expr(plr.lit(value_s._s, allow_object, is_scalar=True))

    if isinstance(value, datetime):
        if dtype == Date:
            return wrap_expr(plr.lit(value.date(), allow_object=False, is_scalar=True))

        # parse time unit
        if dtype is not None and (tu := getattr(dtype, "time_unit", "us")) is not None:
            time_unit = tu  # type: ignore[assignment]
        else:
            time_unit = "us"

        # parse time zone
        dtype_tz = getattr(dtype, "time_zone", None)
        value_tz = value.tzinfo
        if value_tz is None:
            tz = dtype_tz
        else:
            # value has time zone, but dtype does not: keep value time zone
            if dtype_tz is None:
                if isinstance(value_tz, ZoneInfo) or (
                    _check_for_pytz(value_tz)
                    and isinstance(value_tz, pytz.tzinfo.BaseTzInfo)
                    and value_tz.zone is not None

Review comment (suggestion):
We don't know yet that `dtype` actually is a proper `DataType` so it may not have `is_object`.

Code after:

    >>> pl.lit(1)  # doctest: +IGNORE_RESULT
    >>> pl.lit(5.5)  # doctest: +IGNORE_RESULT
    >>> pl.lit(None)  # doctest: +IGNORE_RESULT
    >>> pl.lit("foo_bar")  # doctest: +IGNORE_RESULT
    >>> pl.lit(date(2021, 1, 20))  # doctest: +IGNORE_RESULT
    >>> pl.lit(datetime(2023, 3, 31, 10, 30, 45))  # doctest: +IGNORE_RESULT

    Literal list/Series data (1D):

    >>> pl.lit([1, 2, 3])  # doctest: +SKIP
    >>> pl.lit(pl.Series("x", [1, 2, 3]))  # doctest: +IGNORE_RESULT

    Literal list/Series data (2D):

    >>> pl.lit([[1, 2], [3, 4]])  # doctest: +SKIP
    >>> pl.lit(pl.Series("y", [[1, 2], [3, 4]]))  # doctest: +IGNORE_RESULT
    """
    time_unit: TimeUnit

    if isinstance(dtype, BaseExtension):
        return lit(value, dtype.ext_storage()).ext.to(dtype)
    elif isinstance(dtype, type) and issubclass(dtype, BaseExtension):
        msg = f"dtype '{dtype}' is a BaseExtension class, it should be an instance"
        raise TypeError(msg)
    elif dtype == Object:
        value_s = pl.Series("literal", [value], dtype=dtype)
        return wrap_expr(plr.lit(value_s._s, allow_object, is_scalar=True))

    if isinstance(value, datetime):
        if dtype == Date:
            return wrap_expr(plr.lit(value.date(), allow_object=False, is_scalar=True))

        # parse time unit
        if dtype is not None and (tu := getattr(dtype, "time_unit", "us")) is not None:
            time_unit = tu  # type: ignore[assignment]
        else:
            time_unit = "us"

        # parse time zone
        dtype_tz = getattr(dtype, "time_zone", None)
        value_tz = value.tzinfo
        if value_tz is None:
            tz = dtype_tz
        else:
            # value has time zone, but dtype does not: keep value time zone
            if dtype_tz is None:
                if isinstance(value_tz, ZoneInfo) or (
                    _check_for_pytz(value_tz)
                    and isinstance(value_tz, pytz.tzinfo.BaseTzInfo)
                    and value_tz.zone is not None