fix: Allow for date/datetime subclasses (e.g. pd.Timestamp, FreezeGun…

…) in pl.lit (#18497)

crates/polars-python/src/functions/lazy.rs

@@ -436,18 +436,31 @@ pub fn lit(value: &Bound<'_, PyAny>, allow_object: bool) -> PyResult<PyExpr> {
     ) {
         let av = py_object_to_any_value(value, true)?;
         Ok(Expr::Literal(LiteralValue::try_from(av).unwrap()).into())
-    } else if allow_object {
-        let s = Python::with_gil(|py| {
-            PySeries::new_object(py, "", vec![ObjectValue::from(value.into_py(py))], false).series
-        });
-        Ok(dsl::lit(s).into())
     } else {
-        Err(PyTypeError::new_err(format!(
-            "cannot create expression literal for value of type {}: {}\
-            \n\nHint: Pass `allow_object=True` to accept any value and create a literal of type Object.",
-            value.get_type().qualname()?,
-            value.repr()?
-        )))
+        Python::with_gil(|py| {
+            // One final attempt before erroring. Do we have a date/datetime subclass?
+            // E.g. pd.Timestamp, or Freezegun.
+            let datetime_module = PyModule::import_bound(py, "datetime")?;
+            let datetime_class = datetime_module.getattr("datetime")?;
+            let date_class = datetime_module.getattr("date")?;
+            if value.is_instance(&datetime_class)? || value.is_instance(&date_class)? {
+                let av = py_object_to_any_value(value, true)?;
+                Ok(Expr::Literal(LiteralValue::try_from(av).unwrap()).into())
+            } else if allow_object {
+                let s = Python::with_gil(|py| {
+                    PySeries::new_object(py, "", vec![ObjectValue::from(value.into_py(py))], false)
+                        .series
+                });
+                Ok(dsl::lit(s).into())
+            } else {
+                Err(PyTypeError::new_err(format!(
+                    "cannot create expression literal for value of type {}: {}\
+                    \n\nHint: Pass `allow_object=True` to accept any value and create a literal of type Object.",
+                    value.get_type().qualname()?,
+                    value.repr()?
+                )))
+            }
+        })
     }
 }
 

py-polars/tests/unit/functions/range/test_date_range.py

@@ -310,3 +310,17 @@ def test_date_ranges_datetime_input() -> None:
         "literal", [[date(2022, 1, 1), date(2022, 1, 2), date(2022, 1, 3)]]
     )
     assert_series_equal(result, expected)
+
+
+def test_date_range_with_subclass_18470_18447() -> None:
+    class MyAmazingDate(date):
+        pass
+
+    class MyAmazingDatetime(datetime):
+        pass
+
+    result = pl.datetime_range(
+        MyAmazingDate(2020, 1, 1), MyAmazingDatetime(2020, 1, 2), eager=True
+    )
+    expected = pl.Series("literal", [datetime(2020, 1, 1), datetime(2020, 1, 2)])
+    assert_series_equal(result, expected)

py-polars/tests/unit/functions/test_lit.py

@@ -1,7 +1,7 @@
 from __future__ import annotations
 
 import enum
-from datetime import datetime, timedelta
+from datetime import date, datetime, timedelta
 from decimal import Decimal
 from typing import TYPE_CHECKING, Any
 
@@ -195,3 +195,27 @@ def test_lit_decimal_parametric(s: pl.Series) -> None:
 
     assert df.dtypes[0] == pl.Decimal(None, scale)
     assert result == value
+
+
+def test_lit_datetime_subclass_w_allow_object() -> None:
+    class MyAmazingDate(date):
+        pass
+
+    class MyAmazingDatetime(datetime):
+        pass
+
+    result = pl.select(
+        a=pl.lit(MyAmazingDatetime(2020, 1, 1)),
+        b=pl.lit(MyAmazingDate(2020, 1, 1)),
+        c=pl.lit(MyAmazingDatetime(2020, 1, 1), allow_object=True),
+        d=pl.lit(MyAmazingDate(2020, 1, 1), allow_object=True),
+    )
+    expected = pl.DataFrame(
+        {
+            "a": [datetime(2020, 1, 1)],
+            "b": [date(2020, 1, 1)],
+            "c": [datetime(2020, 1, 1)],
+            "d": [date(2020, 1, 1)],
+        }
+    )
+    assert_frame_equal(result, expected)