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  if sys.version_info < (3, 5):  # covers also python 2
56    if suffix is None:
57      suffix = ''
58    if prefix is None:
59      prefix = 'tmp'
60  (fd, fname) = tempfile.mkstemp(
61      suffix=suffix, prefix=prefix, dir=dirname, text=text)
62  f = io.open(
63      fd,
64      mode=mode,
65      buffering=buffering,
66      encoding=encoding,
67      errors=errors,
68      newline=newline)
69  yield f, fname
70  f.close()
71  os.remove(fname)
72
73
74@contextlib.contextmanager
75def TempFileContents(dirname,
76                     contents,
77                     encoding='utf-8',
78                     newline='',
79                     suffix=None):
80  # Note: NamedTempFile properly handles unicode encoding when using mode='w'
81  with NamedTempFile(
82      dirname=dirname,
83      mode='w',
84      encoding=encoding,
85      newline=newline,
86      suffix=suffix) as (f, fname):
87    f.write(contents)
88    f.flush()
89    yield fname
90