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