1# -*- coding: utf-8 -*-
2#
3# Copyright (C) 2013 Jason R. Coombs <jaraco@jaraco.com>
4# All rights reserved.
5#
6# This software is licensed as described in the file COPYING, which
7# you should have received as part of this distribution.
8
9import unittest
10import datetime
11import time
12
13import jsonpickle
14from jsonpickle import tags
15
16
17class ObjWithDate(object):
18    def __init__(self):
19        ts = datetime.datetime.now()
20        self.data = dict(a='a', ts=ts)
21        self.data_ref = dict(b='b', ts=ts)
22
23
24# UTC implementation from Python 2.7 docs
25class UTC(datetime.tzinfo):
26    """UTC"""
27
28    def utcoffset(self, dt):
29        return datetime.timedelta()
30
31    def tzname(self, dt):
32        return 'UTC'
33
34    def dst(self, dt):
35        return datetime.timedelta()
36
37
38utc = UTC()
39
40
41class TimestampedVariable(object):
42    def __init__(self, value=None):
43        self._value = value
44        self._dt_read = datetime.datetime.utcnow()
45        self._dt_write = self._dt_read
46
47    def get(self, default_value=None):
48        if self._dt_read is None and self._dt_write is None:
49            value = default_value
50            self._value = value
51            self._dt_write = datetime.datetime.utcnow()
52        else:
53            value = self._value
54        self._dt_read = datetime.datetime.utcnow()
55        return value
56
57    def set(self, new_value):
58        self._dt_write = datetime.datetime.utcnow()
59        self._value = new_value
60
61    def __repr__(self):
62        dt_now = datetime.datetime.utcnow()
63        td_read = dt_now - self._dt_read
64        td_write = dt_now - self._dt_write
65        s = '<TimestampedVariable>\n'
66        s += ' value: ' + str(self._value) + '\n'
67        s += ' dt_read: ' + str(self._dt_read) + ' (%s ago)' % td_read + '\n'
68        s += ' dt_write: ' + str(self._dt_write) + ' (%s ago)' % td_write + '\n'
69        return s
70
71    def erasable(self, td=datetime.timedelta(seconds=1)):
72        dt_now = datetime.datetime.utcnow()
73        td_read = dt_now - self._dt_read
74        td_write = dt_now - self._dt_write
75        return td_read > td and td_write > td
76
77
78class PersistantVariables(object):
79    def __init__(self):
80        self._data = {}
81
82    def __getitem__(self, key):
83        return self._data.setdefault(key, TimestampedVariable(None))
84
85    def __setitem__(self, key, value):
86        return self._data.setdefault(key, TimestampedVariable(value))
87
88    def __repr__(self):
89        return str(self._data)
90
91
92class DateTimeInnerReferenceTestCase(unittest.TestCase):
93    def test_object_with_inner_datetime_refs(self):
94        pvars = PersistantVariables()
95        pvars['z'] = 1
96        pvars['z2'] = 2
97        pickled = jsonpickle.encode(pvars)
98        obj = jsonpickle.decode(pickled)
99
100        # ensure the references are valid
101        self.assertTrue(obj['z']._dt_read is obj['z']._dt_write)
102        self.assertTrue(obj['z2']._dt_read is obj['z2']._dt_write)
103
104        # ensure the values are valid
105        self.assertEqual(obj['z'].get(), 1)
106        self.assertEqual(obj['z2'].get(), 2)
107
108        # ensure get() updates _dt_read
109        self.assertTrue(obj['z']._dt_read is not obj['z']._dt_write)
110        self.assertTrue(obj['z2']._dt_read is not obj['z2']._dt_write)
111
112
113class DateTimeSimpleTestCase(unittest.TestCase):
114    def _roundtrip(self, obj):
115        """
116        pickle and then unpickle object, then assert the new object is the
117        same as the original.
118        """
119        pickled = jsonpickle.encode(obj)
120        unpickled = jsonpickle.decode(pickled)
121        self.assertEqual(obj, unpickled)
122
123    def test_datetime(self):
124        """
125        jsonpickle should pickle a datetime object
126        """
127        self._roundtrip(datetime.datetime.now())
128
129    def test_date(self):
130        """
131        jsonpickle should pickle a date object
132        """
133        self._roundtrip(datetime.datetime.today())
134
135    def test_time(self):
136        """
137        jsonpickle should pickle a time object
138        """
139        self._roundtrip(datetime.datetime.now().time())
140
141    def test_timedelta(self):
142        """
143        jsonpickle should pickle a timedelta object
144        """
145        self._roundtrip(datetime.timedelta(days=3))
146
147    def test_utc(self):
148        """
149        jsonpickle should be able to encode and decode a datetime with a
150        simple, pickleable UTC tzinfo.
151        """
152        self._roundtrip(datetime.datetime.utcnow().replace(tzinfo=utc))
153
154    def test_unpickleable(self):
155        """
156        If 'unpickleable' is set on the Pickler, the date objects should be
157        simple, human-readable strings.
158        """
159        obj = datetime.datetime.now()
160        pickler = jsonpickle.pickler.Pickler(unpicklable=False)
161        flattened = pickler.flatten(obj)
162        self.assertEqual(obj.isoformat(), flattened)
163
164    def test_object_with_datetime(self):
165        test_obj = ObjWithDate()
166        json = jsonpickle.encode(test_obj)
167        test_obj_decoded = jsonpickle.decode(json)
168        self.assertEqual(test_obj_decoded.data['ts'], test_obj_decoded.data_ref['ts'])
169
170
171class DateTimeAdvancedTestCase(unittest.TestCase):
172    def setUp(self):
173        self.pickler = jsonpickle.pickler.Pickler()
174        self.unpickler = jsonpickle.unpickler.Unpickler()
175
176    def tearDown(self):
177        self.pickler.reset()
178        self.unpickler.reset()
179
180    def test_struct_time(self):
181        expect = time.struct_time([1, 2, 3, 4, 5, 6, 7, 8, 9])
182        json = jsonpickle.encode(expect)
183        actual = jsonpickle.decode(json)
184        self.assertEqual(type(actual), time.struct_time)
185        self.assertEqual(expect, actual)
186
187    def test_struct_time_chars(self):
188        expect = time.struct_time('123456789')
189
190        flattened = self.pickler.flatten(expect)
191        actual = self.unpickler.restore(flattened)
192        self.assertEqual(expect, actual)
193
194    def test_datetime_structure(self):
195        obj = datetime.datetime.now()
196
197        flattened = self.pickler.flatten(obj)
198        self.assertTrue(tags.OBJECT in flattened)
199        self.assertTrue('__reduce__' in flattened)
200
201        inflated = self.unpickler.restore(flattened)
202        self.assertEqual(obj, inflated)
203
204    def test_datetime_inside_int_keys_defaults(self):
205        t = datetime.time(hour=10)
206        s = jsonpickle.encode({1: t, 2: t})
207        d = jsonpickle.decode(s)
208        self.assertEqual(d["1"], d["2"])
209        self.assertTrue(d["1"] is d["2"])
210        self.assertTrue(isinstance(d["1"], datetime.time))
211
212    def test_datetime_inside_int_keys_with_keys_enabled(self):
213        t = datetime.time(hour=10)
214        s = jsonpickle.encode({1: t, 2: t}, keys=True)
215        d = jsonpickle.decode(s, keys=True)
216        self.assertEqual(d[1], d[2])
217        self.assertTrue(d[1] is d[2])
218        self.assertTrue(isinstance(d[1], datetime.time))
219
220    def test_datetime_repr_not_unpicklable(self):
221        obj = datetime.datetime.now()
222        pickler = jsonpickle.pickler.Pickler(unpicklable=False)
223        flattened = pickler.flatten(obj)
224        self.assertFalse(tags.REPR in flattened)
225        self.assertFalse(tags.OBJECT in flattened)
226        self.assertEqual(obj.isoformat(), flattened)
227
228    def test_datetime_dict_keys_defaults(self):
229        """Test that we handle datetime objects as keys."""
230        datetime_dict = {datetime.datetime(2008, 12, 31): True}
231        pickled = jsonpickle.encode(datetime_dict)
232        expect = {'datetime.datetime(2008, 12, 31, 0, 0)': True}
233        actual = jsonpickle.decode(pickled)
234        self.assertEqual(expect, actual)
235
236    def test_datetime_dict_keys_with_keys_enabled(self):
237        """Test that we handle datetime objects as keys."""
238        datetime_dict = {datetime.datetime(2008, 12, 31): True}
239        pickled = jsonpickle.encode(datetime_dict, keys=True)
240        expect = datetime_dict
241        actual = jsonpickle.decode(pickled, keys=True)
242        self.assertEqual(expect, actual)
243
244
245def suite():
246    suite = unittest.TestSuite()
247    suite.addTest(unittest.makeSuite(DateTimeSimpleTestCase))
248    suite.addTest(unittest.makeSuite(DateTimeAdvancedTestCase))
249    suite.addTest(unittest.makeSuite(DateTimeInnerReferenceTestCase))
250    return suite
251
252
253if __name__ == '__main__':
254    unittest.main(defaultTest='suite')
255