_common.py at 158da71 ⋅ dateutil/dateutil

from __future__ import unicode_literals
import os
import time
import subprocess
import warnings
import tempfile
import pickle

import pytest


class PicklableMixin(object):
    def _get_nobj_bytes(self, obj, dump_kwargs, load_kwargs):
        """
        Pickle and unpickle an object using ``pickle.dumps`` / ``pickle.loads``
        """
        pkl = pickle.dumps(obj, **dump_kwargs)
        return pickle.loads(pkl, **load_kwargs)

    def _get_nobj_file(self, obj, dump_kwargs, load_kwargs):
        """
        Pickle and unpickle an object using ``pickle.dump`` / ``pickle.load`` on
        a temporary file.
        """
        with tempfile.TemporaryFile('w+b') as pkl:
            pickle.dump(obj, pkl, **dump_kwargs)
            pkl.seek(0)         # Reset the file to the beginning to read it
            nobj = pickle.load(pkl, **load_kwargs)

        return nobj

    def assertPicklable(self, obj, singleton=False, asfile=False,
                        dump_kwargs=None, load_kwargs=None):
        """
        Assert that an object can be pickled and unpickled. This assertion
        assumes that the desired behavior is that the unpickled object compares
        equal to the original object, but is not the same object.
        """
        get_nobj = self._get_nobj_file if asfile else self._get_nobj_bytes
        dump_kwargs = dump_kwargs or {}
        load_kwargs = load_kwargs or {}

        nobj = get_nobj(obj, dump_kwargs, load_kwargs)
        if not singleton:
            self.assertIsNot(obj, nobj)
        self.assertEqual(obj, nobj)


class TZContextBase(object):
    """
    Base class for a context manager which allows changing of time zones.

    Subclasses may define a guard variable to either block or or allow time
    zone changes by redefining ``_guard_var_name`` and ``_guard_allows_change``.
    The default is that the guard variable must be affirmatively set.

    Subclasses must define ``get_current_tz`` and ``set_current_tz``.
    """
    _guard_var_name = "DATEUTIL_MAY_CHANGE_TZ"
    _guard_allows_change = True

    def __init__(self, tzval):
        self.tzval = tzval
        self._old_tz = None

    @classmethod
    def tz_change_allowed(cls):
        """
        Class method used to query whether or not this class allows time zone
        changes.
        """
        guard = bool(os.environ.get(cls._guard_var_name, False))

        # _guard_allows_change gives the "default" behavior - if True, the
        # guard is overcoming a block. If false, the guard is causing a block.
        # Whether tz_change is allowed is therefore the XNOR of the two.
        return guard == cls._guard_allows_change

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

        return msg.format(envar=cls._guard_var_name,
                          gval=cls._guard_allows_change)

    def __enter__(self):
        if not self.tz_change_allowed():
            msg = self.tz_change_disallowed_message()
            pytest.skip(msg)

            # If this is used outside of a test suite, we still want an error.
            raise ValueError(msg)  # pragma: no cover

        self._old_tz = self.get_current_tz()
        self.set_current_tz(self.tzval)

    def __exit__(self, type, value, traceback):
        if self._old_tz is not None:
            self.set_current_tz(self._old_tz)

        self._old_tz = None

    def get_current_tz(self):
        raise NotImplementedError

    def set_current_tz(self):
        raise NotImplementedError


class TZEnvContext(TZContextBase):
    """
    Context manager that temporarily sets the `TZ` variable (for use on
    *nix-like systems). Because the effect is local to the shell anyway, this
    will apply *unless* a guard is set.

    If you do not want the TZ environment variable set, you may set the
    ``DATEUTIL_MAY_NOT_CHANGE_TZ_VAR`` variable to a truthy value.
    """
    _guard_var_name = "DATEUTIL_MAY_NOT_CHANGE_TZ_VAR"
    _guard_allows_change = False

    def get_current_tz(self):
        return os.environ.get('TZ', UnsetTz)

    def set_current_tz(self, tzval):
        if tzval is UnsetTz and 'TZ' in os.environ:
            del os.environ['TZ']
        else:
            os.environ['TZ'] = tzval

        time.tzset()


class TZWinContext(TZContextBase):
    """
    Context manager for changing local time zone on Windows.

    Because the effect of this is system-wide and global, it may have
    unintended side effect. Set the ``DATEUTIL_MAY_CHANGE_TZ`` environment
    variable to a truthy value before using this context manager.
    """
    def get_current_tz(self):
        p = subprocess.Popen(['tzutil', '/g'], stdout=subprocess.PIPE)

        ctzname, err = p.communicate()
        ctzname = ctzname.decode()     # Popen returns

        if p.returncode:
            raise OSError('Failed to get current time zone: ' + err)

        return ctzname

    def set_current_tz(self, tzname):
        p = subprocess.Popen('tzutil /s "' + tzname + '"')

        out, err = p.communicate()

        if p.returncode:
            raise OSError('Failed to set current time zone: ' +
                          (err or 'Unknown error.'))


###
# Utility classes
class NotAValueClass(object):
    """
    A class analogous to NaN that has operations defined for any type.
    """
    def _op(self, other):
        return self             # Operation with NotAValue returns NotAValue

    def _cmp(self, other):
        return False

    __add__ = __radd__ = _op
    __sub__ = __rsub__ = _op
    __mul__ = __rmul__ = _op
    __div__ = __rdiv__ = _op
    __truediv__ = __rtruediv__ = _op
    __floordiv__ = __rfloordiv__ = _op

    __lt__ = __rlt__ = _op
    __gt__ = __rgt__ = _op
    __eq__ = __req__ = _op
    __le__ = __rle__ = _op
    __ge__ = __rge__ = _op


NotAValue = NotAValueClass()


class ComparesEqualClass(object):
    """
    A class that is always equal to whatever you compare it to.
    """

    def __eq__(self, other):
        return True

    def __ne__(self, other):
        return False

    def __le__(self, other):
        return True

    def __ge__(self, other):
        return True

    def __lt__(self, other):
        return False

    def __gt__(self, other):
        return False

    __req__ = __eq__
    __rne__ = __ne__
    __rle__ = __le__
    __rge__ = __ge__
    __rlt__ = __lt__
    __rgt__ = __gt__


ComparesEqual = ComparesEqualClass()


class UnsetTzClass(object):
    """ Sentinel class for unset time zone variable """
    pass


UnsetTz = UnsetTzClass()

Read our documentation on viewing source code .

Navigation	Overlay
`t` Navigate files	`h` Toggle hits
`y` Change url to tip of branch	`m` Toggle misses
`b / v` Jump to prev/next hit line	`p` Toggle partial
`z / x` Jump to prev/next missed or partial line	`1..9` Toggle flags
`shift + o` Open current page in GitHub	`a` Toggle all on
`/ or ?` Show keyboard shortcuts dialog	`c` Toggle context lines or commits

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