Dump time and timestamp naive and tz-aware with the respective oids

Using the tz version of the data types is ok enough for most uses,
because Postgres can downgrade the type to naive, but ranges don't have
an implicit cast and the exact type of the range subtype is necessary.

diff --git a/psycopg3/psycopg3/types/__init__.py b/psycopg3/psycopg3/types/__init__.py
index e89a41b5e8c8f11a6f6aa6a6ddb28cbc2e34f012..828c3642144fe01964313c7cafab0e816737da1c 100644 (file)
--- a/psycopg3/psycopg3/types/__init__.py
+++ b/psycopg3/psycopg3/types/__init__.py
@@ -66,6 +66,8 @@ from .singletons import (
 from .date import (
     DateDumper,
     TimeDumper,
+    TimeTzDumper,
+    DateTimeTzDumper,
     DateTimeDumper,
     TimeDeltaDumper,
     DateLoader,
@@ -178,7 +180,7 @@ def register_default_globals(ctx: AdaptContext) -> None:
 
     DateDumper.register("datetime.date", ctx)
     TimeDumper.register("datetime.time", ctx)
-    DateTimeDumper.register("datetime.datetime", ctx)
+    DateTimeTzDumper.register("datetime.datetime", ctx)
     TimeDeltaDumper.register("datetime.timedelta", ctx)
     DateLoader.register("date", ctx)
     TimeLoader.register("time", ctx)

diff --git a/psycopg3/psycopg3/types/date.py b/psycopg3/psycopg3/types/date.py
index 40714928af7a6cbc1a81732345e4f756d14615c8..ad918dbd0af2f27180089496907d2db93e2f6588 100644 (file)
--- a/psycopg3/psycopg3/types/date.py
+++ b/psycopg3/psycopg3/types/date.py
@@ -7,11 +7,11 @@ Adapters for date/time types.
 import re
 import sys
 from datetime import date, datetime, time, timedelta
-from typing import cast, Optional
+from typing import cast, Optional, Tuple, Union
 
 from ..pq import Format
 from ..oids import builtins
-from ..adapt import Buffer, Dumper, Loader
+from ..adapt import Buffer, Dumper, Loader, Format as Pg3Format
 from ..proto import AdaptContext
 from ..errors import InterfaceError, DataError
 
@@ -30,22 +30,67 @@ class DateDumper(Dumper):
 class TimeDumper(Dumper):
 
     format = Format.TEXT
-    _oid = builtins["timetz"].oid
+
+    # Can change to timetz type if the object dumped is naive
+    _oid = builtins["time"].oid
 
     def dump(self, obj: time) -> bytes:
         return str(obj).encode("utf8")
 
+    def get_key(
+        self, obj: time, format: Pg3Format
+    ) -> Union[type, Tuple[type]]:
+        # Use (cls,) to report the need to upgrade  to a dumper for timetz (the
+        # Frankenstein of the data types).
+        if not obj.tzinfo:
+            return self.cls
+        else:
+            return (self.cls,)
+
+    def upgrade(self, obj: time, format: Pg3Format) -> "Dumper":
+        if not obj.tzinfo:
+            return self
+        else:
+            return TimeTzDumper(self.cls)
+
+
+class TimeTzDumper(TimeDumper):
+
+    _oid = builtins["timetz"].oid
+
 
-class DateTimeDumper(Dumper):
+class DateTimeTzDumper(Dumper):
 
     format = Format.TEXT
+
+    # Can change to timestamp type if the object dumped is naive
     _oid = builtins["timestamptz"].oid
 
-    def dump(self, obj: date) -> bytes:
+    def dump(self, obj: datetime) -> bytes:
         # NOTE: whatever the PostgreSQL DateStyle input format (DMY, MDY, YMD)
         # the YYYY-MM-DD is always understood correctly.
         return str(obj).encode("utf8")
 
+    def get_key(
+        self, obj: datetime, format: Pg3Format
+    ) -> Union[type, Tuple[type]]:
+        # Use (cls,) to report the need to upgrade (downgrade, actually) to a
+        # dumper for naive timestamp.
+        if obj.tzinfo:
+            return self.cls
+        else:
+            return (self.cls,)
+
+    def upgrade(self, obj: datetime, format: Pg3Format) -> "Dumper":
+        if obj.tzinfo:
+            return self
+        else:
+            return DateTimeDumper(self.cls)
+
+
+class DateTimeDumper(DateTimeTzDumper):
+    _oid = builtins["timestamp"].oid
+
 
 class TimeDeltaDumper(Dumper):
 

diff --git a/tests/types/test_date.py b/tests/types/test_date.py
index 71f57b8a1d7b4a9781379c6924047e2040f1e0a4..0042f475c6cd7df4a8283b7c4d7bfd65f4118e5e 100644 (file)
--- a/tests/types/test_date.py
+++ b/tests/types/test_date.py
@@ -270,6 +270,29 @@ def test_load_datetimetz_tzname(conn, val, expr, datestyle_in, datestyle_out):
     assert cur.fetchone()[0] == as_dt(val)
 
 
+@pytest.mark.parametrize(
+    "val, type",
+    [
+        ("2000,1,2,3,4,5,6", "timestamp"),
+        ("2000,1,2,3,4,5,6~0", "timestamptz"),
+        ("2000,1,2,3,4,5,6~2", "timestamptz"),
+    ],
+)
+@pytest.mark.parametrize("fmt_in", [Format.AUTO, Format.TEXT, Format.BINARY])
+def test_dump_datetime_tz_or_not_tz(conn, val, type, fmt_in):
+    if fmt_in == Format.BINARY:
+        pytest.xfail("binary datetime not implemented")
+    val = as_dt(val)
+    cur = conn.cursor()
+    cur.execute(
+        f"select pg_typeof(%{fmt_in}) = %s::regtype, %{fmt_in}",
+        [val, type, val],
+    )
+    rec = cur.fetchone()
+    assert rec[0] is True, type
+    assert rec[1] == val
+
+
 #
 # time tests
 #
@@ -395,6 +418,29 @@ def test_load_timetz_24(conn):
         cur.fetchone()[0]
 
 
+@pytest.mark.parametrize(
+    "val, type",
+    [
+        ("3,4,5,6", "time"),
+        ("3,4,5,6~0", "timetz"),
+        ("3,4,5,6~2", "timetz"),
+    ],
+)
+@pytest.mark.parametrize("fmt_in", [Format.AUTO, Format.TEXT, Format.BINARY])
+def test_dump_time_tz_or_not_tz(conn, val, type, fmt_in):
+    if fmt_in == Format.BINARY:
+        pytest.xfail("binary time not implemented")
+    val = as_time(val)
+    cur = conn.cursor()
+    cur.execute(
+        f"select pg_typeof(%{fmt_in}) = %s::regtype, %{fmt_in}",
+        [val, type, val],
+    )
+    rec = cur.fetchone()
+    assert rec[0] is True, type
+    assert rec[1] == val
+
+
 #
 # Interval
 #