Skip to content

Commit 4757e49

Browse files
charettescharettes
committed
Add a database datetime field that allows a timezone to be specified.
1 parent debb709 commit 4757e49

3 files changed

Lines changed: 140 additions & 3 deletions

File tree

‎sundial/fields.py‎

Lines changed: 49 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,8 +1,12 @@
11
from __future__ import unicode_literals
22

33
import django
4+
import pytz
5+
from django.conf import settings
46
from django.db import models
7+
from django.utils import timezone
58
from django.utils.encoding import force_text
9+
from django.utils.functional import cached_property
610

711
from . import forms
812
from .utils import coerce_timezone
@@ -57,3 +61,48 @@ def formfield(self, **kwargs):
5761
kwargs.setdefault('form_class', forms.TimezoneField)
5862
kwargs.setdefault('choices_form_class', forms.TimezoneChoiceField)
5963
return super(TimezoneField, self).formfield(**kwargs)
64+
65+
66+
class DateTimeField(models.DateTimeField):
67+
def __init__(self, *args, **kwargs):
68+
self.db_timezone = kwargs.pop('db_timezone', settings.TIME_ZONE) or 'UTC'
69+
super(DateTimeField, self).__init__(*args, **kwargs)
70+
71+
def deconstruct(self):
72+
name, path, args, kwargs = super(DateTimeField, self).deconstruct()
73+
kwargs['db_timezone'] = self.db_timezone
74+
return name, path, args, kwargs
75+
76+
@cached_property
77+
def db_tzinfo(self):
78+
return pytz.timezone(self.db_timezone)
79+
80+
def get_db_prep_value(self, value, connection, prepared=False):
81+
if value is None or hasattr(value, 'resolve_expression'):
82+
return value
83+
84+
if (settings.USE_TZ and
85+
not connection.features.supports_timezones and timezone.is_aware(value)):
86+
return timezone.make_naive(value, self.db_tzinfo)
87+
88+
return super(DateTimeField, self).get_db_prep_value(value, connection, prepared)
89+
90+
def _from_db_value(self, value, expression, connection):
91+
if value is None:
92+
return value
93+
94+
if settings.USE_TZ and not connection.features.supports_timezones:
95+
# At this point the value will be in UTC even if the actual field's
96+
# timezone is self.db_timezone. Replace the bogus timezone with the
97+
# appropriate one and convert back to the expected UTC.
98+
value = timezone.make_aware(
99+
value.replace(tzinfo=None), self.db_tzinfo
100+
).astimezone(timezone.utc)
101+
102+
return value
103+
104+
if django.VERSION >= (2, 0):
105+
from_db_value = _from_db_value
106+
else:
107+
def from_db_value(self, value, expression, connection, context=None):
108+
return self._from_db_value(value, expression, connection)

‎tests/test_fields/models.py‎

Lines changed: 13 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -2,11 +2,23 @@
22

33
from django.conf import settings
44
from django.db import models
5+
from django.utils.encoding import python_2_unicode_compatible
56

6-
from sundial.fields import TimezoneField
7+
from sundial.fields import DateTimeField, TimezoneField
78

89

910
class TimezoneModel(models.Model):
1011
timezone = TimezoneField(default=settings.TIME_ZONE)
1112
blank_timezone = TimezoneField(blank=True)
1213
null_timezone = TimezoneField(blank=True, null=True)
14+
15+
16+
@python_2_unicode_compatible
17+
class DbTimezoneModel(models.Model):
18+
datetime = DateTimeField(db_timezone='Europe/Amsterdam')
19+
20+
class Meta:
21+
ordering = ('pk',)
22+
23+
def __str__(self):
24+
return models.Model.__str__(self)

‎tests/test_fields/tests.py‎

Lines changed: 78 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1,14 +1,18 @@
11
from __future__ import unicode_literals
22

3+
from datetime import datetime
4+
35
import pytz
46
from django.conf import settings
57
from django.core.exceptions import ValidationError
68
from django.test import TestCase
9+
from django.test.utils import override_settings
10+
from django.utils import timezone
711

8-
from sundial.fields import TimezoneField
12+
from sundial.fields import DateTimeField, TimezoneField
913
from sundial.zones import COMMON_GROUPED_CHOICES
1014

11-
from .models import TimezoneModel
15+
from .models import DbTimezoneModel, TimezoneModel
1216

1317

1418
default_timezone = pytz.timezone(settings.TIME_ZONE)
@@ -75,3 +79,75 @@ def test_choices_formfield(self):
7579
self.assertEqual(formfield.clean(default_timezone.zone), default_timezone)
7680
self.assertEqual(list(formfield.choices), list(field.choices))
7781
self.assertEqual(formfield.initial, 'America/Montreal')
82+
83+
84+
class DateTimeFieldTests(TestCase):
85+
naive_datetime = datetime(2018, 1, 1)
86+
aware_datetime = timezone.make_aware(datetime(2018, 1, 1), pytz.timezone('Europe/Amsterdam'))
87+
88+
@classmethod
89+
def setUpTestData(cls):
90+
cls.naive_object = DbTimezoneModel.objects.create(datetime=cls.naive_datetime)
91+
with override_settings(USE_TZ=True, TIME_ZONE='UTC'):
92+
cls.utc_object = DbTimezoneModel.objects.create(datetime=cls.aware_datetime)
93+
with override_settings(USE_TZ=True, TIME_ZONE='Europe/Paris'):
94+
cls.paris_object = DbTimezoneModel.objects.create(datetime=cls.aware_datetime)
95+
96+
def assert_deconstruct_db_timezone(self, field, expected_db_timezone):
97+
name, path, args, kwargs = field.deconstruct()
98+
self.assertEqual(name, 'field')
99+
self.assertEqual(path, 'sundial.fields.DateTimeField')
100+
self.assertEqual(args, [])
101+
self.assertEqual(kwargs, {'db_timezone': expected_db_timezone})
102+
103+
def test_deconstruct(self):
104+
field = DateTimeField(name='field')
105+
self.assert_deconstruct_db_timezone(field, settings.TIME_ZONE)
106+
with self.settings(TIME_ZONE=None):
107+
field = DateTimeField(name='field')
108+
self.assert_deconstruct_db_timezone(field, 'UTC')
109+
field = DateTimeField(name='field', db_timezone='Europe/Amsterdam')
110+
self.assert_deconstruct_db_timezone(field, 'Europe/Amsterdam')
111+
112+
def test_retreival(self):
113+
self.naive_object.refresh_from_db()
114+
self.assertEqual(self.naive_object.datetime, self.naive_datetime)
115+
self.utc_object.refresh_from_db()
116+
self.assertEqual(self.utc_object.datetime, self.naive_datetime)
117+
with override_settings(USE_TZ=True, TIME_ZONE='UTC'):
118+
utc = timezone.get_current_timezone()
119+
self.naive_object.refresh_from_db()
120+
self.assertEqual(self.naive_object.datetime, self.aware_datetime)
121+
self.assertEqual(self.naive_object.datetime.tzinfo, utc)
122+
self.utc_object.refresh_from_db()
123+
self.assertEqual(self.utc_object.datetime, self.aware_datetime)
124+
self.assertEqual(self.utc_object.datetime.tzinfo, utc)
125+
self.paris_object.refresh_from_db()
126+
self.assertEqual(self.paris_object.datetime, self.aware_datetime)
127+
self.assertEqual(self.paris_object.datetime.tzinfo, utc)
128+
with override_settings(USE_TZ=True, TIME_ZONE='Europe/Paris'):
129+
self.naive_object.refresh_from_db()
130+
self.assertEqual(self.naive_object.datetime, self.aware_datetime)
131+
self.assertEqual(self.naive_object.datetime.tzinfo, utc)
132+
self.utc_object.refresh_from_db()
133+
self.assertEqual(self.utc_object.datetime, self.aware_datetime)
134+
self.assertEqual(self.utc_object.datetime.tzinfo, utc)
135+
self.paris_object.refresh_from_db()
136+
self.assertEqual(self.paris_object.datetime, self.aware_datetime)
137+
self.assertEqual(self.paris_object.datetime.tzinfo, utc)
138+
139+
def test_lookup(self):
140+
self.assertSequenceEqual(
141+
DbTimezoneModel.objects.filter(datetime=self.naive_datetime),
142+
[self.naive_object, self.utc_object, self.paris_object],
143+
)
144+
with override_settings(USE_TZ=True, TIME_ZONE='UTC'):
145+
self.assertSequenceEqual(
146+
DbTimezoneModel.objects.filter(datetime=self.aware_datetime),
147+
[self.naive_object, self.utc_object, self.paris_object],
148+
)
149+
with override_settings(USE_TZ=True, TIME_ZONE='Europe/Paris'):
150+
self.assertSequenceEqual(
151+
DbTimezoneModel.objects.filter(datetime=self.aware_datetime),
152+
[self.naive_object, self.utc_object, self.paris_object],
153+
)

0 commit comments

Comments
 (0)