From: Daniele Varrazzo <daniele.varrazzo@gmail.com>
Date: Fri, 22 Jan 2021 21:30:08 +0000 (+0100)
Subject: Dump time and timestamp naive and tz-aware with the respective oids
X-Git-Tag: 3.0.dev0~132
X-Git-Url: http://git.ipfire.org/cgi-bin/gitweb.cgi?a=commitdiff_plain;h=26f965f9c8411f5a1c79de263d6e343780880c2a;p=thirdparty%2Fpsycopg.git

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 e89a41b5e..828c36421 100644
--- 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 40714928a..ad918dbd0 100644
--- 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 71f57b8a1..0042f475c 100644
--- 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
 #