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

9 17
import pytest
10

11

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

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

30 17
        return nobj
31

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

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

48

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

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

66 17
    @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 17
        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 17
        return guard == cls._guard_allows_change
78

79 17
    @classmethod
80 2
    def tz_change_disallowed_message(cls):
81
        """ Generate instructions on how to allow tz changes """
82 1
        msg = ('Changing time zone not allowed. Set {envar} to {gval} '
83
               'if you would like to allow this behavior')
84

85 1
        return msg.format(envar=cls._guard_var_name,
86
                          gval=cls._guard_allows_change)
87

88 17
    def __enter__(self):
89 17
        if not self.tz_change_allowed():
90 1
            msg = self.tz_change_disallowed_message()
91 1
            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 16
        self._old_tz = self.get_current_tz()
97 16
        self.set_current_tz(self.tzval)
98

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

103 16
        self._old_tz = None
104

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

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

111

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

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

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

133 8
        time.tzset()
134

135

136 17
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 17
    def get_current_tz(self):
145 8
        p = subprocess.Popen(['tzutil', '/g'], stdout=subprocess.PIPE)
146

147 8
        ctzname, err = p.communicate()
148 8
        ctzname = ctzname.decode()     # Popen returns
149

150 8
        if p.returncode:
151 0
            raise OSError('Failed to get current time zone: ' + err)
152

153 8
        return ctzname
154

155 17
    def set_current_tz(self, tzname):
156 8
        p = subprocess.Popen('tzutil /s "' + tzname + '"')
157

158 8
        out, err = p.communicate()
159

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

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

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

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

190

191 17
NotAValue = NotAValueClass()
192

193

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

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

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

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

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

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

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

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

224

225 17
ComparesEqual = ComparesEqualClass()
226

227

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

232

233 17
UnsetTz = UnsetTzClass()

Read our documentation on viewing source code .

Loading