1import asyncio
2import builtins
3import locale
4import logging
5import os
6import shutil
7import sys
8import sysconfig
9import threading
10import urllib.request
11import warnings
12from test import support
13from test.libregrtest.utils import print_warning
14try:
15    import _multiprocessing, multiprocessing.process
16except ImportError:
17    multiprocessing = None
18
19
20# Unit tests are supposed to leave the execution environment unchanged
21# once they complete.  But sometimes tests have bugs, especially when
22# tests fail, and the changes to environment go on to mess up other
23# tests.  This can cause issues with buildbot stability, since tests
24# are run in random order and so problems may appear to come and go.
25# There are a few things we can save and restore to mitigate this, and
26# the following context manager handles this task.
27
28class saved_test_environment:
29    """Save bits of the test environment and restore them at block exit.
30
31        with saved_test_environment(testname, verbose, quiet):
32            #stuff
33
34    Unless quiet is True, a warning is printed to stderr if any of
35    the saved items was changed by the test.  The attribute 'changed'
36    is initially False, but is set to True if a change is detected.
37
38    If verbose is more than 1, the before and after state of changed
39    items is also printed.
40    """
41
42    changed = False
43
44    def __init__(self, testname, verbose=0, quiet=False, *, pgo=False):
45        self.testname = testname
46        self.verbose = verbose
47        self.quiet = quiet
48        self.pgo = pgo
49
50    # To add things to save and restore, add a name XXX to the resources list
51    # and add corresponding get_XXX/restore_XXX functions.  get_XXX should
52    # return the value to be saved and compared against a second call to the
53    # get function when test execution completes.  restore_XXX should accept
54    # the saved value and restore the resource using it.  It will be called if
55    # and only if a change in the value is detected.
56    #
57    # Note: XXX will have any '.' replaced with '_' characters when determining
58    # the corresponding method names.
59
60    resources = ('sys.argv', 'cwd', 'sys.stdin', 'sys.stdout', 'sys.stderr',
61                 'os.environ', 'sys.path', 'sys.path_hooks', '__import__',
62                 'warnings.filters', 'asyncore.socket_map',
63                 'logging._handlers', 'logging._handlerList', 'sys.gettrace',
64                 'sys.warnoptions',
65                 # multiprocessing.process._cleanup() may release ref
66                 # to a thread, so check processes first.
67                 'multiprocessing.process._dangling', 'threading._dangling',
68                 'sysconfig._CONFIG_VARS', 'sysconfig._INSTALL_SCHEMES',
69                 'files', 'locale', 'warnings.showwarning',
70                 'shutil_archive_formats', 'shutil_unpack_formats',
71                 'asyncio.events._event_loop_policy',
72                 'urllib.requests._url_tempfiles', 'urllib.requests._opener',
73                )
74
75    def get_urllib_requests__url_tempfiles(self):
76        return list(urllib.request._url_tempfiles)
77    def restore_urllib_requests__url_tempfiles(self, tempfiles):
78        for filename in tempfiles:
79            support.unlink(filename)
80
81    def get_urllib_requests__opener(self):
82        return urllib.request._opener
83    def restore_urllib_requests__opener(self, opener):
84        urllib.request._opener = opener
85
86    def get_asyncio_events__event_loop_policy(self):
87        return support.maybe_get_event_loop_policy()
88    def restore_asyncio_events__event_loop_policy(self, policy):
89        asyncio.set_event_loop_policy(policy)
90
91    def get_sys_argv(self):
92        return id(sys.argv), sys.argv, sys.argv[:]
93    def restore_sys_argv(self, saved_argv):
94        sys.argv = saved_argv[1]
95        sys.argv[:] = saved_argv[2]
96
97    def get_cwd(self):
98        return os.getcwd()
99    def restore_cwd(self, saved_cwd):
100        os.chdir(saved_cwd)
101
102    def get_sys_stdout(self):
103        return sys.stdout
104    def restore_sys_stdout(self, saved_stdout):
105        sys.stdout = saved_stdout
106
107    def get_sys_stderr(self):
108        return sys.stderr
109    def restore_sys_stderr(self, saved_stderr):
110        sys.stderr = saved_stderr
111
112    def get_sys_stdin(self):
113        return sys.stdin
114    def restore_sys_stdin(self, saved_stdin):
115        sys.stdin = saved_stdin
116
117    def get_os_environ(self):
118        return id(os.environ), os.environ, dict(os.environ)
119    def restore_os_environ(self, saved_environ):
120        os.environ = saved_environ[1]
121        os.environ.clear()
122        os.environ.update(saved_environ[2])
123
124    def get_sys_path(self):
125        return id(sys.path), sys.path, sys.path[:]
126    def restore_sys_path(self, saved_path):
127        sys.path = saved_path[1]
128        sys.path[:] = saved_path[2]
129
130    def get_sys_path_hooks(self):
131        return id(sys.path_hooks), sys.path_hooks, sys.path_hooks[:]
132    def restore_sys_path_hooks(self, saved_hooks):
133        sys.path_hooks = saved_hooks[1]
134        sys.path_hooks[:] = saved_hooks[2]
135
136    def get_sys_gettrace(self):
137        return sys.gettrace()
138    def restore_sys_gettrace(self, trace_fxn):
139        sys.settrace(trace_fxn)
140
141    def get___import__(self):
142        return builtins.__import__
143    def restore___import__(self, import_):
144        builtins.__import__ = import_
145
146    def get_warnings_filters(self):
147        return id(warnings.filters), warnings.filters, warnings.filters[:]
148    def restore_warnings_filters(self, saved_filters):
149        warnings.filters = saved_filters[1]
150        warnings.filters[:] = saved_filters[2]
151
152    def get_asyncore_socket_map(self):
153        asyncore = sys.modules.get('asyncore')
154        # XXX Making a copy keeps objects alive until __exit__ gets called.
155        return asyncore and asyncore.socket_map.copy() or {}
156    def restore_asyncore_socket_map(self, saved_map):
157        asyncore = sys.modules.get('asyncore')
158        if asyncore is not None:
159            asyncore.close_all(ignore_all=True)
160            asyncore.socket_map.update(saved_map)
161
162    def get_shutil_archive_formats(self):
163        # we could call get_archives_formats() but that only returns the
164        # registry keys; we want to check the values too (the functions that
165        # are registered)
166        return shutil._ARCHIVE_FORMATS, shutil._ARCHIVE_FORMATS.copy()
167    def restore_shutil_archive_formats(self, saved):
168        shutil._ARCHIVE_FORMATS = saved[0]
169        shutil._ARCHIVE_FORMATS.clear()
170        shutil._ARCHIVE_FORMATS.update(saved[1])
171
172    def get_shutil_unpack_formats(self):
173        return shutil._UNPACK_FORMATS, shutil._UNPACK_FORMATS.copy()
174    def restore_shutil_unpack_formats(self, saved):
175        shutil._UNPACK_FORMATS = saved[0]
176        shutil._UNPACK_FORMATS.clear()
177        shutil._UNPACK_FORMATS.update(saved[1])
178
179    def get_logging__handlers(self):
180        # _handlers is a WeakValueDictionary
181        return id(logging._handlers), logging._handlers, logging._handlers.copy()
182    def restore_logging__handlers(self, saved_handlers):
183        # Can't easily revert the logging state
184        pass
185
186    def get_logging__handlerList(self):
187        # _handlerList is a list of weakrefs to handlers
188        return id(logging._handlerList), logging._handlerList, logging._handlerList[:]
189    def restore_logging__handlerList(self, saved_handlerList):
190        # Can't easily revert the logging state
191        pass
192
193    def get_sys_warnoptions(self):
194        return id(sys.warnoptions), sys.warnoptions, sys.warnoptions[:]
195    def restore_sys_warnoptions(self, saved_options):
196        sys.warnoptions = saved_options[1]
197        sys.warnoptions[:] = saved_options[2]
198
199    # Controlling dangling references to Thread objects can make it easier
200    # to track reference leaks.
201    def get_threading__dangling(self):
202        # This copies the weakrefs without making any strong reference
203        return threading._dangling.copy()
204    def restore_threading__dangling(self, saved):
205        threading._dangling.clear()
206        threading._dangling.update(saved)
207
208    # Same for Process objects
209    def get_multiprocessing_process__dangling(self):
210        if not multiprocessing:
211            return None
212        # Unjoined process objects can survive after process exits
213        multiprocessing.process._cleanup()
214        # This copies the weakrefs without making any strong reference
215        return multiprocessing.process._dangling.copy()
216    def restore_multiprocessing_process__dangling(self, saved):
217        if not multiprocessing:
218            return
219        multiprocessing.process._dangling.clear()
220        multiprocessing.process._dangling.update(saved)
221
222    def get_sysconfig__CONFIG_VARS(self):
223        # make sure the dict is initialized
224        sysconfig.get_config_var('prefix')
225        return (id(sysconfig._CONFIG_VARS), sysconfig._CONFIG_VARS,
226                dict(sysconfig._CONFIG_VARS))
227    def restore_sysconfig__CONFIG_VARS(self, saved):
228        sysconfig._CONFIG_VARS = saved[1]
229        sysconfig._CONFIG_VARS.clear()
230        sysconfig._CONFIG_VARS.update(saved[2])
231
232    def get_sysconfig__INSTALL_SCHEMES(self):
233        return (id(sysconfig._INSTALL_SCHEMES), sysconfig._INSTALL_SCHEMES,
234                sysconfig._INSTALL_SCHEMES.copy())
235    def restore_sysconfig__INSTALL_SCHEMES(self, saved):
236        sysconfig._INSTALL_SCHEMES = saved[1]
237        sysconfig._INSTALL_SCHEMES.clear()
238        sysconfig._INSTALL_SCHEMES.update(saved[2])
239
240    def get_files(self):
241        return sorted(fn + ('/' if os.path.isdir(fn) else '')
242                      for fn in os.listdir())
243    def restore_files(self, saved_value):
244        fn = support.TESTFN
245        if fn not in saved_value and (fn + '/') not in saved_value:
246            if os.path.isfile(fn):
247                support.unlink(fn)
248            elif os.path.isdir(fn):
249                support.rmtree(fn)
250
251    _lc = [getattr(locale, lc) for lc in dir(locale)
252           if lc.startswith('LC_')]
253    def get_locale(self):
254        pairings = []
255        for lc in self._lc:
256            try:
257                pairings.append((lc, locale.setlocale(lc, None)))
258            except (TypeError, ValueError):
259                continue
260        return pairings
261    def restore_locale(self, saved):
262        for lc, setting in saved:
263            locale.setlocale(lc, setting)
264
265    def get_warnings_showwarning(self):
266        return warnings.showwarning
267    def restore_warnings_showwarning(self, fxn):
268        warnings.showwarning = fxn
269
270    def resource_info(self):
271        for name in self.resources:
272            method_suffix = name.replace('.', '_')
273            get_name = 'get_' + method_suffix
274            restore_name = 'restore_' + method_suffix
275            yield name, getattr(self, get_name), getattr(self, restore_name)
276
277    def __enter__(self):
278        self.saved_values = dict((name, get()) for name, get, restore
279                                                   in self.resource_info())
280        return self
281
282    def __exit__(self, exc_type, exc_val, exc_tb):
283        saved_values = self.saved_values
284        del self.saved_values
285
286        # Some resources use weak references
287        support.gc_collect()
288
289        # Read support.environment_altered, set by support helper functions
290        self.changed |= support.environment_altered
291
292        for name, get, restore in self.resource_info():
293            current = get()
294            original = saved_values.pop(name)
295            # Check for changes to the resource's value
296            if current != original:
297                self.changed = True
298                restore(original)
299                if not self.quiet and not self.pgo:
300                    print_warning(f"{name} was modified by {self.testname}")
301                    print(f"  Before: {original}\n  After:  {current} ",
302                          file=sys.stderr, flush=True)
303        return False
304