1
|
20
|
from __future__ import unicode_literals
|
2
|
20
|
import os
|
3
|
20
|
import time
|
4
|
20
|
import subprocess
|
5
|
20
|
import warnings
|
6
|
20
|
import tempfile
|
7
|
20
|
import pickle
|
8
|
|
|
9
|
20
|
import pytest
|
10
|
|
|
11
|
|
|
12
|
20
|
class PicklableMixin(object):
|
13
|
20
|
def _get_nobj_bytes(self, obj, dump_kwargs, load_kwargs):
|
14
|
|
"""
|
15
|
|
Pickle and unpickle an object using ``pickle.dumps`` / ``pickle.loads``
|
16
|
|
"""
|
17
|
20
|
pkl = pickle.dumps(obj, **dump_kwargs)
|
18
|
20
|
return pickle.loads(pkl, **load_kwargs)
|
19
|
|
|
20
|
20
|
def _get_nobj_file(self, obj, dump_kwargs, load_kwargs):
|
21
|
|
"""
|
22
|
|
Pickle and unpickle an object using ``pickle.dump`` / ``pickle.load`` on
|
23
|
|
a temporary file.
|
24
|
|
"""
|
25
|
20
|
with tempfile.TemporaryFile('w+b') as pkl:
|
26
|
20
|
pickle.dump(obj, pkl, **dump_kwargs)
|
27
|
20
|
pkl.seek(0) # Reset the file to the beginning to read it
|
28
|
20
|
nobj = pickle.load(pkl, **load_kwargs)
|
29
|
|
|
30
|
20
|
return nobj
|
31
|
|
|
32
|
20
|
def assertPicklable(self, obj, singleton=False, asfile=False,
|
33
|
|
dump_kwargs=None, load_kwargs=None):
|
34
|
|
"""
|
35
|
|
Assert that an object can be pickled and unpickled. This assertion
|
36
|
|
assumes that the desired behavior is that the unpickled object compares
|
37
|
|
equal to the original object, but is not the same object.
|
38
|
|
"""
|
39
|
20
|
get_nobj = self._get_nobj_file if asfile else self._get_nobj_bytes
|
40
|
20
|
dump_kwargs = dump_kwargs or {}
|
41
|
20
|
load_kwargs = load_kwargs or {}
|
42
|
|
|
43
|
20
|
nobj = get_nobj(obj, dump_kwargs, load_kwargs)
|
44
|
20
|
if not singleton:
|
45
|
20
|
self.assertIsNot(obj, nobj)
|
46
|
20
|
self.assertEqual(obj, nobj)
|
47
|
|
|
48
|
|
|
49
|
20
|
class TZContextBase(object):
|
50
|
|
"""
|
51
|
|
Base class for a context manager which allows changing of time zones.
|
52
|
|
|
53
|
|
Subclasses may define a guard variable to either block or or allow time
|
54
|
|
zone changes by redefining ``_guard_var_name`` and ``_guard_allows_change``.
|
55
|
|
The default is that the guard variable must be affirmatively set.
|
56
|
|
|
57
|
|
Subclasses must define ``get_current_tz`` and ``set_current_tz``.
|
58
|
|
"""
|
59
|
20
|
_guard_var_name = "DATEUTIL_MAY_CHANGE_TZ"
|
60
|
20
|
_guard_allows_change = True
|
61
|
|
|
62
|
20
|
def __init__(self, tzval):
|
63
|
20
|
self.tzval = tzval
|
64
|
20
|
self._old_tz = None
|
65
|
|
|
66
|
20
|
@classmethod
|
67
|
2
|
def tz_change_allowed(cls):
|
68
|
|
"""
|
69
|
|
Class method used to query whether or not this class allows time zone
|
70
|
|
changes.
|
71
|
|
"""
|
72
|
20
|
guard = bool(os.environ.get(cls._guard_var_name, False))
|
73
|
|
|
74
|
|
# _guard_allows_change gives the "default" behavior - if True, the
|
75
|
|
# guard is overcoming a block. If false, the guard is causing a block.
|
76
|
|
# Whether tz_change is allowed is therefore the XNOR of the two.
|
77
|
20
|
return guard == cls._guard_allows_change
|
78
|
|
|
79
|
20
|
@classmethod
|
80
|
2
|
def tz_change_disallowed_message(cls):
|
81
|
|
""" Generate instructions on how to allow tz changes """
|
82
|
2
|
msg = ('Changing time zone not allowed. Set {envar} to {gval} '
|
83
|
|
'if you would like to allow this behavior')
|
84
|
|
|
85
|
2
|
return msg.format(envar=cls._guard_var_name,
|
86
|
|
gval=cls._guard_allows_change)
|
87
|
|
|
88
|
20
|
def __enter__(self):
|
89
|
20
|
if not self.tz_change_allowed():
|
90
|
2
|
msg = self.tz_change_disallowed_message()
|
91
|
2
|
pytest.skip(msg)
|
92
|
|
|
93
|
|
# If this is used outside of a test suite, we still want an error.
|
94
|
|
raise ValueError(msg) # pragma: no cover
|
95
|
|
|
96
|
18
|
self._old_tz = self.get_current_tz()
|
97
|
18
|
self.set_current_tz(self.tzval)
|
98
|
|
|
99
|
20
|
def __exit__(self, type, value, traceback):
|
100
|
18
|
if self._old_tz is not None:
|
101
|
18
|
self.set_current_tz(self._old_tz)
|
102
|
|
|
103
|
18
|
self._old_tz = None
|
104
|
|
|
105
|
20
|
def get_current_tz(self):
|
106
|
0
|
raise NotImplementedError
|
107
|
|
|
108
|
20
|
def set_current_tz(self):
|
109
|
0
|
raise NotImplementedError
|
110
|
|
|
111
|
|
|
112
|
20
|
class TZEnvContext(TZContextBase):
|
113
|
|
"""
|
114
|
|
Context manager that temporarily sets the `TZ` variable (for use on
|
115
|
|
*nix-like systems). Because the effect is local to the shell anyway, this
|
116
|
|
will apply *unless* a guard is set.
|
117
|
|
|
118
|
|
If you do not want the TZ environment variable set, you may set the
|
119
|
|
``DATEUTIL_MAY_NOT_CHANGE_TZ_VAR`` variable to a truthy value.
|
120
|
|
"""
|
121
|
20
|
_guard_var_name = "DATEUTIL_MAY_NOT_CHANGE_TZ_VAR"
|
122
|
20
|
_guard_allows_change = False
|
123
|
|
|
124
|
20
|
def get_current_tz(self):
|
125
|
11
|
return os.environ.get('TZ', UnsetTz)
|
126
|
|
|
127
|
20
|
def set_current_tz(self, tzval):
|
128
|
11
|
if tzval is UnsetTz and 'TZ' in os.environ:
|
129
|
11
|
del os.environ['TZ']
|
130
|
|
else:
|
131
|
11
|
os.environ['TZ'] = tzval
|
132
|
|
|
133
|
11
|
time.tzset()
|
134
|
|
|
135
|
|
|
136
|
20
|
class TZWinContext(TZContextBase):
|
137
|
|
"""
|
138
|
|
Context manager for changing local time zone on Windows.
|
139
|
|
|
140
|
|
Because the effect of this is system-wide and global, it may have
|
141
|
|
unintended side effect. Set the ``DATEUTIL_MAY_CHANGE_TZ`` environment
|
142
|
|
variable to a truthy value before using this context manager.
|
143
|
|
"""
|
144
|
20
|
def get_current_tz(self):
|
145
|
7
|
p = subprocess.Popen(['tzutil', '/g'], stdout=subprocess.PIPE)
|
146
|
|
|
147
|
7
|
ctzname, err = p.communicate()
|
148
|
7
|
ctzname = ctzname.decode() # Popen returns
|
149
|
|
|
150
|
7
|
if p.returncode:
|
151
|
0
|
raise OSError('Failed to get current time zone: ' + err)
|
152
|
|
|
153
|
7
|
return ctzname
|
154
|
|
|
155
|
20
|
def set_current_tz(self, tzname):
|
156
|
7
|
p = subprocess.Popen('tzutil /s "' + tzname + '"')
|
157
|
|
|
158
|
7
|
out, err = p.communicate()
|
159
|
|
|
160
|
7
|
if p.returncode:
|
161
|
0
|
raise OSError('Failed to set current time zone: ' +
|
162
|
|
(err or 'Unknown error.'))
|
163
|
|
|
164
|
|
|
165
|
|
###
|
166
|
|
# Utility classes
|
167
|
20
|
class NotAValueClass(object):
|
168
|
|
"""
|
169
|
|
A class analogous to NaN that has operations defined for any type.
|
170
|
|
"""
|
171
|
20
|
def _op(self, other):
|
172
|
20
|
return self # Operation with NotAValue returns NotAValue
|
173
|
|
|
174
|
20
|
def _cmp(self, other):
|
175
|
0
|
return False
|
176
|
|
|
177
|
20
|
__add__ = __radd__ = _op
|
178
|
20
|
__sub__ = __rsub__ = _op
|
179
|
20
|
__mul__ = __rmul__ = _op
|
180
|
20
|
__div__ = __rdiv__ = _op
|
181
|
20
|
__truediv__ = __rtruediv__ = _op
|
182
|
20
|
__floordiv__ = __rfloordiv__ = _op
|
183
|
|
|
184
|
20
|
__lt__ = __rlt__ = _op
|
185
|
20
|
__gt__ = __rgt__ = _op
|
186
|
20
|
__eq__ = __req__ = _op
|
187
|
20
|
__le__ = __rle__ = _op
|
188
|
20
|
__ge__ = __rge__ = _op
|
189
|
|
|
190
|
|
|
191
|
20
|
NotAValue = NotAValueClass()
|
192
|
|
|
193
|
|
|
194
|
20
|
class ComparesEqualClass(object):
|
195
|
|
"""
|
196
|
|
A class that is always equal to whatever you compare it to.
|
197
|
|
"""
|
198
|
|
|
199
|
20
|
def __eq__(self, other):
|
200
|
20
|
return True
|
201
|
|
|
202
|
20
|
def __ne__(self, other):
|
203
|
0
|
return False
|
204
|
|
|
205
|
20
|
def __le__(self, other):
|
206
|
0
|
return True
|
207
|
|
|
208
|
20
|
def __ge__(self, other):
|
209
|
0
|
return True
|
210
|
|
|
211
|
20
|
def __lt__(self, other):
|
212
|
0
|
return False
|
213
|
|
|
214
|
20
|
def __gt__(self, other):
|
215
|
0
|
return False
|
216
|
|
|
217
|
20
|
__req__ = __eq__
|
218
|
20
|
__rne__ = __ne__
|
219
|
20
|
__rle__ = __le__
|
220
|
20
|
__rge__ = __ge__
|
221
|
20
|
__rlt__ = __lt__
|
222
|
20
|
__rgt__ = __gt__
|
223
|
|
|
224
|
|
|
225
|
20
|
ComparesEqual = ComparesEqualClass()
|
226
|
|
|
227
|
|
|
228
|
20
|
class UnsetTzClass(object):
|
229
|
|
""" Sentinel class for unset time zone variable """
|
230
|
20
|
pass
|
231
|
|
|
232
|
|
|
233
|
20
|
UnsetTz = UnsetTzClass()
|