• Home
  • Line#
  • Scopes#
  • Navigate#
  • Raw
  • Download
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
17import contextlib
18import io
19import os
20import sys
21import tempfile
22
23
24@contextlib.contextmanager
25def 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
45def 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
70def 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
85