1 68
from __future__ import unicode_literals
2 68
import os
3 68
import time
4 68
import subprocess
5 68
import warnings
6 68
import tempfile
7 68
import pickle
8

9 68
import pytest
10

11

12 68
class PicklableMixin(object):
13 68
    def _get_nobj_bytes(self, obj, dump_kwargs, load_kwargs):
14
        """
15
        Pickle and unpickle an object using ``pickle.dumps`` / ``pickle.loads``
16
        """
17 68
        pkl = pickle.dumps(obj, **dump_kwargs)
18 68
        return pickle.loads(pkl, **load_kwargs)
19

20 68
    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 68
        with tempfile.TemporaryFile('w+b') as pkl:
26 68
            pickle.dump(obj, pkl, **dump_kwargs)
27 68
            pkl.seek(0)         # Reset the file to the beginning to read it
28 68
            nobj = pickle.load(pkl, **load_kwargs)
29

30 68
        return nobj
31

32 68
    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 68
        get_nobj = self._get_nobj_file if asfile else self._get_nobj_bytes
40 68
        dump_kwargs = dump_kwargs or {}
41 68
        load_kwargs = load_kwargs or {}
42

43 68
        nobj = get_nobj(obj, dump_kwargs, load_kwargs)
44 68
        if not singleton:
45 68
            self.assertIsNot(obj, nobj)
46 68
        self.assertEqual(obj, nobj)
47

48

49 68
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 68
    _guard_var_name = "DATEUTIL_MAY_CHANGE_TZ"
60 68
    _guard_allows_change = True
61

62 68
    def __init__(self, tzval):
63 68
        self.tzval = tzval
64 68
        self._old_tz = None
65

66 68
    @classmethod
67 21
    def tz_change_allowed(cls):
68
        """
69
        Class method used to query whether or not this class allows time zone
70
        changes.
71
        """
72 68
        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 68
        return guard == cls._guard_allows_change
78

79 68
    @classmethod
80 21
    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 68
    def __enter__(self):
89 68
        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 66
        self._old_tz = self.get_current_tz()
97 66
        self.set_current_tz(self.tzval)
98

99 68
    def __exit__(self, type, value, traceback):
100 66
        if self._old_tz is not None:
101 66
            self.set_current_tz(self._old_tz)
102

103 66
        self._old_tz = None
104

105 68
    def get_current_tz(self):
106 0
        raise NotImplementedError
107

108 68
    def set_current_tz(self):
109 0
        raise NotImplementedError
110

111

112 68
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 68
    _guard_var_name = "DATEUTIL_MAY_NOT_CHANGE_TZ_VAR"
122 68
    _guard_allows_change = False
123

124 68
    def get_current_tz(self):
125 59
        return os.environ.get('TZ', UnsetTz)
126

127 68
    def set_current_tz(self, tzval):
128 59
        if tzval is UnsetTz and 'TZ' in os.environ:
129 59
            del os.environ['TZ']
130
        else:
131 59
            os.environ['TZ'] = tzval
132

133 59
        time.tzset()
134

135

136 68
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 68
    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 68
    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 68
class NotAValueClass(object):
168
    """
169
    A class analogous to NaN that has operations defined for any type.
170
    """
171 68
    def _op(self, other):
172 68
        return self             # Operation with NotAValue returns NotAValue
173

174 68
    def _cmp(self, other):
175 0
        return False
176

177 68
    __add__ = __radd__ = _op
178 68
    __sub__ = __rsub__ = _op
179 68
    __mul__ = __rmul__ = _op
180 68
    __div__ = __rdiv__ = _op
181 68
    __truediv__ = __rtruediv__ = _op
182 68
    __floordiv__ = __rfloordiv__ = _op
183

184 68
    __lt__ = __rlt__ = _op
185 68
    __gt__ = __rgt__ = _op
186 68
    __eq__ = __req__ = _op
187 68
    __le__ = __rle__ = _op
188 68
    __ge__ = __rge__ = _op
189

190

191 68
NotAValue = NotAValueClass()
192

193

194 68
class ComparesEqualClass(object):
195
    """
196
    A class that is always equal to whatever you compare it to.
197
    """
198

199 68
    def __eq__(self, other):
200 68
        return True
201

202 68
    def __ne__(self, other):
203 0
        return False
204

205 68
    def __le__(self, other):
206 0
        return True
207

208 68
    def __ge__(self, other):
209 0
        return True
210

211 68
    def __lt__(self, other):
212 0
        return False
213

214 68
    def __gt__(self, other):
215 0
        return False
216

217 68
    __req__ = __eq__
218 68
    __rne__ = __ne__
219 68
    __rle__ = __le__
220 68
    __rge__ = __ge__
221 68
    __rlt__ = __lt__
222 68
    __rgt__ = __gt__
223

224

225 68
ComparesEqual = ComparesEqualClass()
226

227

228 68
class UnsetTzClass(object):
229
    """ Sentinel class for unset time zone variable """
230 68
    pass
231

232

233 68
UnsetTz = UnsetTzClass()

Read our documentation on viewing source code .

Loading