Add ZonedDateTime data type to GQLAlchemy

memgraph · Jul 1, 2024 · 109c774 · 109c774
1 parent 9bf07f6
commit 109c774
Show file tree

Hide file tree

Showing 3 changed files with 37 additions and 2 deletions.
diff --git a/gqlalchemy/models.py b/gqlalchemy/models.py
@@ -37,13 +37,15 @@ class DatetimeKeywords(Enum):
     LOCALTIME = "localTime"
     LOCALDATETIME = "localDateTime"
     DATE = "date"
+    ZONEDDATETIME = "datetime"
 
 
 datetimeKwMapping = {
     timedelta: DatetimeKeywords.DURATION.value,
     time: DatetimeKeywords.LOCALTIME.value,
     datetime: DatetimeKeywords.LOCALDATETIME.value,
     date: DatetimeKeywords.DATE.value,
+    datetime: DatetimeKeywords.ZONEDDATETIME.value  
 }
 
 
@@ -378,7 +380,18 @@ def escape_value(
         elif value_type == dict:
             return "{" + ", ".join(f"{key}: {self.escape_value(val)}" for key, val in value.items()) + "}"
         if isinstance(value, (timedelta, time, datetime, date)):
-            return f"{datetimeKwMapping[value_type]}('{_format_timedelta(value) if isinstance(value, timedelta) else value.isoformat()}')"
+            if isinstance(value, timedelta):
+                formatted_value = _format_timedelta(value)
+            else:
+                formatted_value = value.isoformat()
+
+            if isinstance(value, datetime) and value.tzinfo is not None:
+                tz_offset = value.strftime('%z')
+                tz_name = value.tzinfo.zone
+                return f"datetime('{value.strftime('%Y-%m-%dT%H:%M:%S')}{tz_offset}[{tz_name}]')"
+            else:
+                keyword = datetimeKwMapping[value_type]
+                return f"{keyword}('{formatted_value}')"
         else:
             raise GQLAlchemyError(
                 f"Unsupported value data type: {type(value)}."

diff --git a/gqlalchemy/utilities.py b/gqlalchemy/utilities.py
@@ -17,6 +17,7 @@
 from enum import Enum
 import inspect
 import math
+import pytz
 from typing import Any, Dict, List, Optional, Tuple, Union
 
 import numpy as np
@@ -34,13 +35,15 @@ class DatetimeKeywords(Enum):
     LOCALTIME = "localTime"
     LOCALDATETIME = "localDateTime"
     DATE = "date"
+    ZONEDDATETIME = "datetime"
 
 
 datetimeKwMapping = {
     timedelta: DatetimeKeywords.DURATION.value,
     time: DatetimeKeywords.LOCALTIME.value,
     datetime: DatetimeKeywords.LOCALDATETIME.value,
     date: DatetimeKeywords.DATE.value,
+    datetime: DatetimeKeywords.ZONEDDATETIME.value 
 }
 
 
@@ -106,7 +109,20 @@ def to_cypher_value(value: Any, config: NetworkXCypherConfig = None) -> str:
         return str(value)
 
     if isinstance(value, (timedelta, time, datetime, date)):
-        return f"{datetimeKwMapping[value_type]}('{_format_timedelta(value) if isinstance(value, timedelta) else value.isoformat()}')"
+        if isinstance(value, datetime):
+            if value.tzinfo == pytz.UTC:
+                formatted_date = value.strftime(f"%Y-%m-%dT%H:%M:%SZ")
+                return f"{DatetimeKeywords.ZONEDDATETIME.value}('{formatted_date}')"
+            elif value.tzinfo is not None:
+                tz = value.strftime('%z')
+                tz = f"{tz[:3]}:{tz[3:]}" 
+                tz_name = value.tzinfo.zone
+                formatted_date = value.strftime(f"%Y-%m-%dT%H:%M:%S{tz}")
+                return f"{DatetimeKeywords.ZONEDDATETIME.value}('{formatted_date}[{tz_name}]')"
+            else:
+                return f"{DatetimeKeywords.LOCALDATETIME.value}('{value.isoformat()}')"
+        else:
+            return f"{datetimeKwMapping[value_type]}('{_format_timedelta(value) if isinstance(value, timedelta) else value.isoformat()}')"
 
     if value_type == str and value.lower() in ["true", "false", "null"]:
         return value

diff --git a/tests/test_utilities.py b/tests/test_utilities.py
@@ -16,6 +16,7 @@
 import math
 import numpy as np
 import pytest
+import pytz
 
 from gqlalchemy.utilities import (
     NanException,
@@ -68,11 +69,16 @@ def test_to_cypher_datetime():
     localtime = datetime.time(12, 12, 12)
     localdatetime = datetime.datetime(1999, 12, 12, 12, 12, 12)
     duration = datetime.timedelta(days=1, hours=5, minutes=16, seconds=12)
+    zoned_dt_naive = datetime.datetime(2024, 4, 21, 14, 15, 0)
+    zoned_dt = pytz.timezone("America/Los_Angeles").localize(zoned_dt_naive)
+    zoned_dt_utc = datetime.datetime(2021, 4, 21, 14, 15, tzinfo=pytz.UTC)
 
     assert to_cypher_value(date) == "date('1970-01-19')"
     assert to_cypher_value(localtime) == "localTime('12:12:12')"
     assert to_cypher_value(localdatetime) == "localDateTime('1999-12-12T12:12:12')"
     assert to_cypher_value(duration) == "duration('P1DT5H16M12.0S')"
+    assert to_cypher_value(zoned_dt) == "datetime('2024-04-21T14:15:00-07:00[America/Los_Angeles]')"
+    assert to_cypher_value(zoned_dt_utc) == "datetime('2021-04-21T14:15:00Z')"
 
 
 def test_to_cypher_labels_single_label():