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()

Read our documentation on viewing source code .

Loading