Add PostgreSQL-only AtTimeZone database function

Author: Alex Hill

naivedatetimefield/__init__.py

@@ -5,7 +5,7 @@
 import django
 import pytz
 from django.core import exceptions, checks
-from django.db.models import DateTimeField
+from django.db.models import DateTimeField, Func, Value
 from django.utils import timezone
 from django.utils.dateparse import parse_date, parse_datetime
 from django.utils.translation import gettext_lazy as _
@@ -175,6 +175,38 @@ def as_sql(self, compiler, connection):
         return super(NaiveAsSQLMixin, self).as_sql(compiler, connection)
 
 
+class AtTimeZone(Func):
+    """
+    This implements PostgreSQL's AT TIME ZONE construct, which returns a naive
+    datetime if used with a timezone-aware datetime, and vice versa.
+
+    See https://www.postgresql.org/docs/9.6/functions-datetime.html#FUNCTIONS-DATETIME-ZONECONVERT # noqa
+    """
+    def __init__(self, value, tz):
+        super(AtTimeZone, self).__init__(
+            value,
+            tz,
+            template='%(expressions)s',
+            arg_joiner=' AT TIME ZONE ',
+        )
+
+    def _resolve_output_field(self):
+        if getattr(self, '_output_field', None) is None:
+            value_field, _ = super(AtTimeZone, self).get_source_fields()
+            if isinstance(value_field, NaiveDateTimeField):
+                self._output_field = DateTimeField()
+            elif isinstance(value_field, DateTimeField):
+                self._output_field = NaiveDateTimeField()
+            elif value_field is None:
+                value_expr = self.get_source_expressions()[0]
+                if isinstance(value_expr, Value):
+                    if timezone.is_naive(value_expr.value):
+                        self._output_field = DateTimeField()
+                    else:
+                        self._output_field = NaiveDateTimeField()
+        return getattr(self, '_output_field', None)
+
+
 _monkeypatching = False
 
 

tests/models.py

@@ -6,6 +6,10 @@
 class NaiveDateTimeTestModel(models.Model):
     aware = models.DateTimeField()
     naive = NaiveDateTimeField()
+    timezone = models.CharField(max_length=100, blank=True, default='utc')
+
+    class Meta:
+        ordering = ['pk']
 
 
 class NaiveDateTimeAutoNowAddModel(models.Model):

tests/tests.py

@@ -5,11 +5,13 @@
 import pytz
 from django import db
 from django.contrib.auth.models import User
-from django.db.models import functions
+from django.db import connection
+from django.db.models import functions, Value
 from django.test import TestCase, override_settings
 from django.utils import timezone
 
 import naivedatetimefield
+from naivedatetimefield import AtTimeZone
 from .models import (
     NaiveDateTimeTestModel,
     NaiveDateTimeAutoNowAddModel,
@@ -344,3 +346,138 @@ def test_select_by_naive(self):
         ).count()
 
         self.assertTrue(find_with_naive_in_utc == 1)
+
+
+def identity(v):
+    return v
+
+
+@skipIf(connection.vendor != "postgresql", "AtTimeZone is only supported in PostgreSQL")
+class AtTimeZoneTests(TestCase):
+    @classmethod
+    def setUpTestData(cls):
+        cls.now = datetime.datetime(2019, 1, 15, 10)
+        cls.perth_tz = pytz.timezone('Australia/Perth')
+        cls.sydney_tz = pytz.timezone('Australia/Sydney')
+        cls.adelaide_tz = pytz.timezone('Australia/Adelaide')
+
+        cls.perth = NaiveDateTimeTestModel.objects.create(
+            naive=cls.now,
+            aware=timezone.make_aware(cls.now, cls.perth_tz),
+            timezone='Australia/Perth',
+        )
+
+        cls.sydney = NaiveDateTimeTestModel.objects.create(
+            naive=cls.now,
+            aware=timezone.make_aware(cls.now, cls.sydney_tz),
+            timezone='Australia/Sydney',
+        )
+
+    def test_annotate(self):
+        self.assertQuerysetEqual(
+            NaiveDateTimeTestModel.objects.annotate(
+                naive_converted=AtTimeZone('aware', 'timezone'),
+                aware_converted=AtTimeZone('naive', 'timezone'),
+            ).values_list('naive_converted', 'aware_converted'),
+            [
+                (self.now, timezone.make_aware(self.now, self.perth_tz)),
+                (self.now, timezone.make_aware(self.now, self.sydney_tz)),
+            ],
+            transform=identity,
+        )
+
+    def test_db_aware_db_timezone(self):
+        self.assertQuerysetEqual(
+            NaiveDateTimeTestModel.objects.filter(
+                naive=AtTimeZone(
+                    'aware',
+                    'timezone',
+                )
+            ),
+            [self.perth, self.sydney],
+            transform=identity,
+        )
+
+    def test_db_aware_value_timezone(self):
+        self.assertQuerysetEqual(
+            NaiveDateTimeTestModel.objects.filter(
+                naive__lt=AtTimeZone(
+                    'aware',
+                    Value('Australia/Adelaide'),
+                )
+            ),
+            [self.perth],
+            transform=identity,
+        )
+
+    def test_db_naive_db_timezone(self):
+        self.assertQuerysetEqual(
+            NaiveDateTimeTestModel.objects.filter(
+                aware=AtTimeZone(
+                    'naive',
+                    'timezone',
+                )
+            ),
+            [self.perth, self.sydney],
+            transform=identity,
+        )
+
+    def test_db_naive_value_timezone(self):
+        self.assertQuerysetEqual(
+            NaiveDateTimeTestModel.objects.filter(
+                aware__lt=AtTimeZone(
+                    'naive',
+                    Value('Australia/Adelaide'),
+                )
+            ),
+            [self.sydney],
+            transform=identity,
+        )
+
+    def test_value_aware_db_timezone(self):
+        self.assertQuerysetEqual(
+            NaiveDateTimeTestModel.objects.filter(
+                naive__lt=AtTimeZone(
+                    timezone.make_aware(self.now, self.adelaide_tz),
+                    'timezone',
+                )
+            ),
+            [self.sydney],
+            transform=identity,
+        )
+
+    def test_value_aware_value_timezone(self):
+        self.assertQuerysetEqual(
+            NaiveDateTimeTestModel.objects.filter(
+                naive=AtTimeZone(
+                    Value(timezone.make_aware(self.now, self.adelaide_tz)),
+                    Value('Australia/Adelaide'),
+                )
+            ),
+            [self.perth, self.sydney],
+            transform=identity,
+        )
+
+    def test_value_naive_db_timezone(self):
+        self.assertQuerysetEqual(
+            NaiveDateTimeTestModel.objects.filter(
+                aware=AtTimeZone(
+                    Value(self.now),
+                    'timezone',
+                )
+            ),
+            [self.perth, self.sydney],
+            transform=identity,
+        )
+
+    def test_value_naive_value_timezone(self):
+        self.assertQuerysetEqual(
+            NaiveDateTimeTestModel.objects.filter(
+                aware__lt=AtTimeZone(
+                    Value(self.now),
+                    Value('Australia/Adelaide'),
+                )
+            ),
+            [self.sydney],
+            transform=identity,
+        )