From 32042a52d16b19cab809e7504bacf8f2740143e9 Mon Sep 17 00:00:00 2001
From: Kamil Piechowiak <32928185+KamilPiechowiak@users.noreply.github.com>
Date: Mon, 29 Jan 2024 10:53:03 +0100
Subject: [PATCH] fix conversion to rust from pw.Duration (#5519)
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

* fix conversion to rust from duration

* tests

* changelog

* Update public/pathway/CHANGELOG.md

Co-authored-by: Michał Bartoszkiewicz <embe@pathway.com>

---------

Co-authored-by: Michał Bartoszkiewicz <embe@pathway.com>
GitOrigin-RevId: a5eec69638f424e08da143bd190b114fd7a39b77
---
 CHANGELOG.md                                  |  3 ++
 .../tests/expressions/test_datetimes.py       | 53 +++++++++++++++++++
 src/python_api.rs                             |  4 +-
 3 files changed, 58 insertions(+), 2 deletions(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 45b53ba2..49266dcf 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -5,6 +5,9 @@ All notable changes to this project will be documented in this file.
 This project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.html).
 ## [Unreleased]
 
+### Fixed
+- Returning `pw.Duration` from UDFs or using them as constant values no longer results in errors.
+
 ## [0.7.10] - 2024-01-26
 
 ### Added
diff --git a/python/pathway/tests/expressions/test_datetimes.py b/python/pathway/tests/expressions/test_datetimes.py
index 4931d5db..37554c50 100644
--- a/python/pathway/tests/expressions/test_datetimes.py
+++ b/python/pathway/tests/expressions/test_datetimes.py
@@ -1014,3 +1014,56 @@ def test_weekday(is_naive: bool) -> None:
     table_pw = table.select(txt=table.ts.dt.weekday())
     table_pd = table_from_pandas(df_new)
     assert_table_equality(table_pw, table_pd)
+
+
+def test_pathway_duration():
+    t = table_from_markdown(
+        """
+        value
+        1
+    """
+    )
+
+    @pw.udf
+    def to_duration(a) -> pw.Duration:
+        return pw.Duration(days=a)
+
+    result = t.select(value=to_duration(pw.this.value))
+    assert_table_equality(
+        result, table_from_pandas(pd.DataFrame({"value": [pd.Timedelta(days=1)]}))
+    )
+
+
+def test_pathway_datetimes():
+    @pw.udf
+    def to_naive(year, month, day) -> pw.DateTimeNaive:
+        return pw.DateTimeNaive(year=year, month=month, day=day)
+
+    @pw.udf
+    def to_utc(year, month, day) -> pw.DateTimeUtc:
+        return pw.DateTimeUtc(year=year, month=month, day=day, tz=tz.UTC)
+
+    t = table_from_markdown(
+        """
+        year | month | day
+        2023 |   8   |  12
+    """
+    )
+
+    result = t.select(value=to_naive(pw.this.year, pw.this.month, pw.this.day))
+    assert_table_equality(
+        result,
+        table_from_pandas(
+            pd.DataFrame({"value": [pd.Timestamp(year=2023, month=8, day=12)]})
+        ),
+    )
+
+    result = t.select(value=to_utc(pw.this.year, pw.this.month, pw.this.day))
+    assert_table_equality(
+        result,
+        table_from_pandas(
+            pd.DataFrame(
+                {"value": [pd.Timestamp(year=2023, month=8, day=12, tz=tz.UTC)]}
+            )
+        ),
+    )
diff --git a/src/python_api.rs b/src/python_api.rs
index 8d45aa7a..0b16fbe3 100644
--- a/src/python_api.rs
+++ b/src/python_api.rs
@@ -308,9 +308,9 @@ impl<'source> FromPyObject<'source> for Value {
                 return value_from_python_datetime(ob);
             } else if type_name == "timedelta" {
                 return value_from_python_timedelta(ob);
-            } else if type_name == "Timestamp" {
+            } else if matches!(type_name, "Timestamp" | "DateTimeNaive" | "DateTimeUtc") {
                 return value_from_pandas_timestamp(ob);
-            } else if type_name == "Timedelta" {
+            } else if matches!(type_name, "Timedelta" | "Duration") {
                 return value_from_pandas_timedelta(ob);
             } else if type_name == "Json" {
                 return value_json_from_py_any(ob.getattr("value")?);