]>
Commit | Line | Data |
---|---|---|
53e6db90 DC |
1 | # -*- coding: utf-8 -*- |
2 | # Copyright 2017 Google Inc. All Rights Reserved. | |
3 | # | |
4 | # Licensed under the Apache License, Version 2.0 (the "License"); | |
5 | # you may not use this file except in compliance with the License. | |
6 | # You may obtain a copy of the License at | |
7 | # | |
8 | # http://www.apache.org/licenses/LICENSE-2.0 | |
9 | # | |
10 | # Unless required by applicable law or agreed to in writing, software | |
11 | # distributed under the License is distributed on an "AS IS" BASIS, | |
12 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. | |
13 | # See the License for the specific language governing permissions and | |
14 | # limitations under the License. | |
15 | """Utilities for tests.""" | |
16 | ||
17 | import contextlib | |
18 | import io | |
19 | import os | |
20 | import sys | |
21 | import tempfile | |
22 | ||
23 | ||
24 | @contextlib.contextmanager | |
25 | def stdout_redirector(stream): # pylint: disable=invalid-name | |
26 | old_stdout = sys.stdout | |
27 | sys.stdout = stream | |
28 | try: | |
29 | yield | |
30 | finally: | |
31 | sys.stdout = old_stdout | |
32 | ||
33 | ||
34 | # NamedTemporaryFile is useless because on Windows the temporary file would be | |
35 | # created with O_TEMPORARY, which would not allow the file to be opened a | |
36 | # second time, even by the same process, unless the same flag is used. | |
37 | # Thus we provide a simplified version ourselves. | |
38 | # | |
39 | # Note: returns a tuple of (io.file_obj, file_path), instead of a file_obj with | |
40 | # a .name attribute | |
41 | # | |
42 | # Note: `buffering` is set to -1 despite documentation of NamedTemporaryFile | |
43 | # says None. This is probably a problem with the python documentation. | |
44 | @contextlib.contextmanager | |
45 | def NamedTempFile(mode='w+b', | |
46 | buffering=-1, | |
47 | encoding=None, | |
48 | errors=None, | |
49 | newline=None, | |
50 | suffix=None, | |
51 | prefix=None, | |
52 | dirname=None, | |
53 | text=False): | |
54 | """Context manager creating a new temporary file in text mode.""" | |
55 | (fd, fname) = tempfile.mkstemp( | |
56 | suffix=suffix, prefix=prefix, dir=dirname, text=text) | |
57 | f = io.open( | |
58 | fd, | |
59 | mode=mode, | |
60 | buffering=buffering, | |
61 | encoding=encoding, | |
62 | errors=errors, | |
63 | newline=newline) | |
64 | yield f, fname | |
65 | f.close() | |
66 | os.remove(fname) | |
67 | ||
68 | ||
69 | @contextlib.contextmanager | |
70 | def TempFileContents(dirname, | |
71 | contents, | |
72 | encoding='utf-8', | |
73 | newline='', | |
74 | suffix=None): | |
75 | # Note: NamedTempFile properly handles unicode encoding when using mode='w' | |
76 | with NamedTempFile( | |
77 | dirname=dirname, | |
78 | mode='w', | |
79 | encoding=encoding, | |
80 | newline=newline, | |
81 | suffix=suffix) as (f, fname): | |
82 | f.write(contents) | |
83 | f.flush() | |
84 | yield fname |