Coverage for /pythoncovmergedfiles/medio/medio/usr/local/lib/python3.11/site-packages/click/testing.py: 27%
Shortcuts on this page
r m x toggle line displays
j k next/prev highlighted chunk
0 (zero) top of page
1 (one) first highlighted chunk
1from __future__ import annotations
3import collections.abc as cabc
4import contextlib
5import io
6import os
7import pdb
8import shlex
9import sys
10import tempfile
11import typing as t
12from types import TracebackType
14from . import _compat
15from . import formatting
16from . import termui
17from . import utils
18from ._compat import _find_binary_reader
20if t.TYPE_CHECKING:
21 from _typeshed import ReadableBuffer
23 from .core import Command
25if sys.platform == "win32":
26 CaptureMode: t.TypeAlias = t.Literal["sys"] # pyright: ignore[reportRedeclaration]
27else:
28 CaptureMode: t.TypeAlias = t.Literal["sys", "fd"] # pyright: ignore[reportRedeclaration]
29ExceptionInfo: t.TypeAlias = tuple[type[BaseException], BaseException, TracebackType]
32class EchoingStdin:
33 _input: t.BinaryIO
34 _output: t.BinaryIO
35 _paused: bool
37 def __init__(self, input: t.BinaryIO, output: t.BinaryIO) -> None:
38 self._input = input
39 self._output = output
40 self._paused = False
42 def __getattr__(self, x: str) -> t.Any:
43 return getattr(self._input, x)
45 def _echo(self, rv: bytes) -> bytes:
46 if not self._paused:
47 self._output.write(rv)
49 return rv
51 def read(self, n: int = -1) -> bytes:
52 return self._echo(self._input.read(n))
54 def read1(self, n: int = -1) -> bytes:
55 return self._echo(self._input.read1(n)) # type: ignore
57 def readline(self, n: int = -1) -> bytes:
58 return self._echo(self._input.readline(n))
60 def readlines(self) -> list[bytes]:
61 return [self._echo(x) for x in self._input.readlines()]
63 def __iter__(self) -> cabc.Iterator[bytes]:
64 return iter(self._echo(x) for x in self._input)
66 def __repr__(self) -> str:
67 return repr(self._input)
70@contextlib.contextmanager
71def _pause_echo(stream: EchoingStdin | None) -> cabc.Generator[None]:
72 if stream is None:
73 yield
74 else:
75 stream._paused = True
76 yield
77 stream._paused = False
80class _FDCapture:
81 """Redirect a file descriptor to a temporary file for capture.
83 Saves the current target of *targetfd* via :func:`os.dup`, then
84 redirects it to a temporary file via :func:`os.dup2`. On
85 :meth:`stop`, restores the original ``fd`` and returns the captured
86 bytes. Inspired by Pytest's ``FDCapture``.
88 .. versionadded:: 8.4.0
89 """
91 _targetfd: int
92 saved_fd: int
93 _tmpfile: t.BinaryIO | None
95 def __init__(self, targetfd: int) -> None:
96 self._targetfd = targetfd
97 self.saved_fd = -1
98 self._tmpfile = None
100 def start(self) -> None:
101 self.saved_fd = os.dup(self._targetfd)
102 self._tmpfile = tempfile.TemporaryFile(buffering=0)
103 os.dup2(self._tmpfile.fileno(), self._targetfd)
105 def stop(self) -> bytes:
106 assert self._tmpfile is not None, "_FDCapture.start() was not called"
107 os.dup2(self.saved_fd, self._targetfd)
108 os.close(self.saved_fd)
109 self.saved_fd = -1
110 self._tmpfile.seek(0)
111 data = self._tmpfile.read()
112 self._tmpfile.close()
113 self._tmpfile = None
114 return data
117class BytesIOCopy(io.BytesIO):
118 """Patch ``io.BytesIO`` to let the written stream be copied to another.
120 .. versionadded:: 8.2
121 """
123 copy_to: io.BytesIO
125 def __init__(self, copy_to: io.BytesIO) -> None:
126 super().__init__()
127 self.copy_to = copy_to
129 def flush(self) -> None:
130 super().flush()
131 self.copy_to.flush()
133 def write(self, b: ReadableBuffer) -> int:
134 self.copy_to.write(b)
135 return super().write(b)
138class StreamMixer:
139 """Mixes `<stdout>` and `<stderr>` streams.
141 The result is available in the ``output`` attribute.
143 .. versionadded:: 8.2
144 """
146 output: io.BytesIO
147 stdout: BytesIOCopy
148 stderr: BytesIOCopy
150 def __init__(self) -> None:
151 self.output = io.BytesIO()
152 self.stdout = BytesIOCopy(copy_to=self.output)
153 self.stderr = BytesIOCopy(copy_to=self.output)
156class _NamedTextIOWrapper(io.TextIOWrapper):
157 """A :class:`~io.TextIOWrapper` with custom ``name`` and ``mode``
158 that does not close its underlying buffer.
160 When ``CliRunner`` runs in ``fd`` mode, ``_original_fd`` is patched to
161 point at the saved (pre-redirection) ``fd``, so C-level consumers that call
162 :meth:`fileno` (like ``faulthandler`` or ``subprocess``) keep working. In
163 the default ``sys`` mode ``_original_fd`` stays at ``-1`` and
164 :meth:`fileno` raises :exc:`io.UnsupportedOperation`, matching the
165 pre-``8.3.3`` behavior.
166 """
168 _name: str
169 _mode: str
170 _original_fd: int
172 def __init__(
173 self,
174 buffer: t.BinaryIO,
175 name: str,
176 mode: str,
177 **kwargs: t.Any,
178 ) -> None:
179 super().__init__(buffer, **kwargs)
180 self._name = name
181 self._mode = mode
182 self._original_fd = -1
184 def close(self) -> None:
185 """The buffer this object contains belongs to some other object,
186 so prevent the default ``__del__`` implementation from closing
187 that buffer.
189 .. versionadded:: 8.3.2
190 """
192 def fileno(self) -> int:
193 """Return the file descriptor of the saved original stream when
194 ``CliRunner`` runs in ``fd`` mode. Otherwise delegate to
195 :class:`~io.TextIOWrapper`, which raises
196 :exc:`io.UnsupportedOperation` for a ``BytesIO``-backed buffer.
197 """
198 if self._original_fd >= 0:
199 return self._original_fd
200 return super().fileno()
202 @property
203 def name(self) -> str:
204 return self._name
206 @property
207 def mode(self) -> str:
208 return self._mode
211def make_input_stream(
212 input: str | bytes | t.IO[t.Any] | None, charset: str
213) -> t.BinaryIO:
214 # Is already an input stream.
215 if hasattr(input, "read"):
216 rv = _find_binary_reader(t.cast("t.IO[t.Any]", input))
218 if rv is not None:
219 return rv
221 raise TypeError("Could not find binary reader for input stream.")
223 if input is None:
224 input = b""
225 elif isinstance(input, str):
226 input = input.encode(charset)
228 return io.BytesIO(input)
231class Result:
232 """Holds the captured result of an invoked CLI script.
234 :param runner: The runner that created the result
235 :param stdout_bytes: The standard output as bytes.
236 :param stderr_bytes: The standard error as bytes.
237 :param output_bytes: A mix of ``stdout_bytes`` and ``stderr_bytes``, as the
238 user would see it in its terminal.
239 :param return_value: The value returned from the invoked command.
240 :param exit_code: The exit code as integer.
241 :param exception: The exception that happened if one did.
242 :param exc_info: Exception information (exception type, exception instance,
243 traceback type).
245 .. versionchanged:: 8.2
246 ``stderr_bytes`` no longer optional, ``output_bytes`` introduced and
247 ``mix_stderr`` has been removed.
249 .. versionadded:: 8.0
250 Added ``return_value``.
251 """
253 runner: CliRunner
254 stdout_bytes: bytes
255 stderr_bytes: bytes
256 output_bytes: bytes
257 return_value: t.Any
258 exit_code: int
259 exception: BaseException | None
260 exc_info: ExceptionInfo | None
262 def __init__(
263 self,
264 runner: CliRunner,
265 stdout_bytes: bytes,
266 stderr_bytes: bytes,
267 output_bytes: bytes,
268 return_value: t.Any,
269 exit_code: int,
270 exception: BaseException | None,
271 exc_info: ExceptionInfo | None = None,
272 ) -> None:
273 self.runner = runner
274 self.stdout_bytes = stdout_bytes
275 self.stderr_bytes = stderr_bytes
276 self.output_bytes = output_bytes
277 self.return_value = return_value
278 self.exit_code = exit_code
279 self.exception = exception
280 self.exc_info = exc_info
282 @property
283 def output(self) -> str:
284 """The terminal output as unicode string, as the user would see it.
286 .. versionchanged:: 8.2
287 No longer a proxy for ``self.stdout``. Now has its own independent stream
288 that is mixing `<stdout>` and `<stderr>`, in the order they were written.
289 """
290 return self.output_bytes.decode(self.runner.charset, "replace").replace(
291 "\r\n", "\n"
292 )
294 @property
295 def stdout(self) -> str:
296 """The standard output as unicode string."""
297 return self.stdout_bytes.decode(self.runner.charset, "replace").replace(
298 "\r\n", "\n"
299 )
301 @property
302 def stderr(self) -> str:
303 """The standard error as unicode string.
305 .. versionchanged:: 8.2
306 No longer raise an exception, always returns the `<stderr>` string.
307 """
308 return self.stderr_bytes.decode(self.runner.charset, "replace").replace(
309 "\r\n", "\n"
310 )
312 def __repr__(self) -> str:
313 exc_str = repr(self.exception) if self.exception else "okay"
314 return f"<{type(self).__name__} {exc_str}>"
317class CliRunner:
318 """The CLI runner provides functionality to invoke a Click command line
319 script for unittesting purposes in a isolated environment. This only
320 works in single-threaded systems without any concurrency as it changes the
321 global interpreter state.
323 :param charset: the character set for the input and output data.
324 :param env: a dictionary with environment variables for overriding.
325 :param echo_stdin: if this is set to `True`, then reading from `<stdin>` writes
326 to `<stdout>`. This is useful for showing examples in
327 some circumstances. Note that regular prompts
328 will automatically echo the input.
329 :param catch_exceptions: Whether to catch any exceptions other than
330 ``SystemExit`` when running :meth:`~CliRunner.invoke`.
331 :param capture: Selects the output capture strategy. ``sys`` (default)
332 captures Python-level writes only and leaves
333 :meth:`sys.stdout.fileno` raising :exc:`io.UnsupportedOperation`, so
334 user code that calls :func:`os.dup2` on ``sys.stdout.fileno()`` cannot
335 clobber the host runner's stdout. ``fd`` redirects file descriptors
336 ``1`` and ``2`` via :func:`os.dup2` to a temporary file, also catching
337 output from stale stream references, C extensions, and subprocesses.
338 ``fd`` is not supported on Windows.
340 .. versionchanged:: 8.4.0
341 Added the ``capture`` parameter. The default ``sys`` mode no longer
342 exposes the original fd through :meth:`fileno`, reverting the change
343 introduced in ``8.3.3`` that broke Pytest's ``fd``-level capture
344 teardown. Use ``capture="fd"`` to restore that behavior with proper
345 isolation. :issue:`3384`
347 .. versionchanged:: 8.2
348 Added the ``catch_exceptions`` parameter.
350 .. versionchanged:: 8.2
351 ``mix_stderr`` parameter has been removed.
352 """
354 charset: str
355 env: cabc.Mapping[str, str | None]
356 echo_stdin: bool
357 catch_exceptions: bool
358 capture: CaptureMode
360 def __init__(
361 self,
362 charset: str = "utf-8",
363 env: cabc.Mapping[str, str | None] | None = None,
364 echo_stdin: bool = False,
365 catch_exceptions: bool = True,
366 capture: CaptureMode = "sys",
367 ) -> None:
368 if capture not in {"sys", "fd"}:
369 raise ValueError(
370 f"capture={capture!r} is not valid. Choose from 'sys' or 'fd'."
371 )
372 if capture == "fd" and sys.platform == "win32":
373 raise ValueError(
374 f"capture={capture!r} is not supported on Windows. Use 'sys'."
375 )
376 self.charset = charset
377 self.env = env or {}
378 self.echo_stdin = echo_stdin
379 self.catch_exceptions = catch_exceptions
380 self.capture = capture
382 def get_default_prog_name(self, cli: Command) -> str:
383 """Given a command object it will return the default program name
384 for it. The default is the `name` attribute or ``"root"`` if not
385 set.
386 """
387 return cli.name or "root"
389 def make_env(
390 self, overrides: cabc.Mapping[str, str | None] | None = None
391 ) -> cabc.Mapping[str, str | None]:
392 """Returns the environment overrides for invoking a script."""
393 rv = dict(self.env)
394 if overrides:
395 rv.update(overrides)
396 return rv
398 @contextlib.contextmanager
399 def isolation(
400 self,
401 input: str | bytes | t.IO[t.Any] | None = None,
402 env: cabc.Mapping[str, str | None] | None = None,
403 color: bool = False,
404 ) -> cabc.Generator[tuple[io.BytesIO, io.BytesIO, io.BytesIO]]:
405 """A context manager that sets up the isolation for invoking of a
406 command line tool. This sets up `<stdin>` with the given input data
407 and `os.environ` with the overrides from the given dictionary.
408 This also rebinds some internals in Click to be mocked (like the
409 prompt functionality).
411 This is automatically done in the :meth:`invoke` method.
413 :param input: the input stream to put into `sys.stdin`.
414 :param env: the environment overrides as dictionary.
415 :param color: whether the output should contain color codes. The
416 application can still override this explicitly.
418 .. versionadded:: 8.2
419 An additional output stream is returned, which is a mix of
420 `<stdout>` and `<stderr>` streams.
422 .. versionchanged:: 8.2
423 Always returns the `<stderr>` stream.
425 .. versionchanged:: 8.0
426 `<stderr>` is opened with ``errors="backslashreplace"``
427 instead of the default ``"strict"``.
429 .. versionchanged:: 4.0
430 Added the ``color`` parameter.
431 """
432 bytes_input = make_input_stream(input, self.charset)
433 echo_input = None
435 old_stdin = sys.stdin
436 old_stdout = sys.stdout
437 old_stderr = sys.stderr
438 old_forced_width = formatting.FORCED_WIDTH
439 formatting.FORCED_WIDTH = 80
441 env = self.make_env(env)
443 stream_mixer = StreamMixer()
445 if self.echo_stdin:
446 bytes_input = echo_input = t.cast(
447 t.BinaryIO, EchoingStdin(bytes_input, stream_mixer.stdout)
448 )
450 sys.stdin = text_input = _NamedTextIOWrapper(
451 bytes_input, encoding=self.charset, name="<stdin>", mode="r"
452 )
454 if self.echo_stdin:
455 # Force unbuffered reads, otherwise TextIOWrapper reads a
456 # large chunk which is echoed early.
457 text_input._CHUNK_SIZE = 1 # type: ignore
459 sys.stdout = _NamedTextIOWrapper(
460 stream_mixer.stdout,
461 encoding=self.charset,
462 name="<stdout>",
463 mode="w",
464 )
466 sys.stderr = _NamedTextIOWrapper(
467 stream_mixer.stderr,
468 encoding=self.charset,
469 name="<stderr>",
470 mode="w",
471 errors="backslashreplace",
472 )
474 @_pause_echo(echo_input) # type: ignore
475 def visible_input(prompt: str | None = None) -> str:
476 sys.stdout.write(prompt or "")
477 try:
478 val = next(text_input).rstrip("\r\n")
479 except StopIteration as e:
480 raise EOFError() from e
481 sys.stdout.write(f"{val}\n")
482 sys.stdout.flush()
483 return val
485 @_pause_echo(echo_input) # type: ignore
486 def hidden_input(prompt: str | None = None) -> str:
487 sys.stdout.write(f"{prompt or ''}\n")
488 sys.stdout.flush()
489 try:
490 return next(text_input).rstrip("\r\n")
491 except StopIteration as e:
492 raise EOFError() from e
494 @_pause_echo(echo_input) # type: ignore
495 def _getchar(echo: bool) -> str:
496 char = sys.stdin.read(1)
498 if echo:
499 sys.stdout.write(char)
501 sys.stdout.flush()
502 return char
504 default_color = color
506 def should_strip_ansi(
507 stream: t.IO[t.Any] | None = None, color: bool | None = None
508 ) -> bool:
509 if color is None:
510 return not default_color
511 return not color
513 old_visible_prompt_func = termui.visible_prompt_func
514 old_hidden_prompt_func = termui.hidden_prompt_func
515 old__getchar_func = termui._getchar
516 old_should_strip_ansi = utils.should_strip_ansi # type: ignore
517 old__compat_should_strip_ansi = _compat.should_strip_ansi
518 old_pdb_init = pdb.Pdb.__init__
519 termui.visible_prompt_func = visible_input
520 termui.hidden_prompt_func = hidden_input
521 termui._getchar = _getchar
522 utils.should_strip_ansi = should_strip_ansi # type: ignore
523 _compat.should_strip_ansi = should_strip_ansi
525 def _patched_pdb_init(
526 self: pdb.Pdb,
527 completekey: str = "tab",
528 stdin: t.IO[str] | None = None,
529 stdout: t.IO[str] | None = None,
530 **kwargs: t.Any,
531 ) -> None:
532 """Default ``pdb.Pdb`` to real terminal streams during
533 ``CliRunner`` isolation.
535 Without this patch, ``pdb.Pdb.__init__`` inherits from
536 ``cmd.Cmd`` which falls back to ``sys.stdin``/``sys.stdout``
537 when no explicit streams are provided. During isolation
538 those are ``BytesIO``-backed wrappers, so the debugger
539 reads from an empty buffer and writes to captured output,
540 making interactive debugging impossible.
542 By defaulting to ``sys.__stdin__``/``sys.__stdout__`` (the
543 original terminal streams Python preserves regardless of
544 redirection), debuggers can interact with the user while
545 ``click.echo`` output is still captured normally.
547 This covers ``pdb.set_trace()``, ``breakpoint()``,
548 ``pdb.post_mortem()``, and debuggers that subclass
549 ``pdb.Pdb`` (ipdb, pdbpp). Explicit ``stdin``/``stdout``
550 arguments are honored and not overridden. Debuggers that
551 do not subclass ``pdb.Pdb`` (pudb, debugpy) are not
552 covered.
553 """
554 if stdin is None:
555 stdin = sys.__stdin__
556 if stdout is None:
557 stdout = sys.__stdout__
558 old_pdb_init(
559 self, completekey=completekey, stdin=stdin, stdout=stdout, **kwargs
560 )
562 pdb.Pdb.__init__ = _patched_pdb_init # type: ignore[assignment]
564 old_env = {}
565 try:
566 for key, value in env.items():
567 old_env[key] = os.environ.get(key)
568 if value is None:
569 try:
570 del os.environ[key]
571 except Exception:
572 pass
573 else:
574 os.environ[key] = value
575 yield (stream_mixer.stdout, stream_mixer.stderr, stream_mixer.output)
576 finally:
577 for key, value in old_env.items():
578 if value is None:
579 try:
580 del os.environ[key]
581 except Exception:
582 pass
583 else:
584 os.environ[key] = value
585 sys.stdout = old_stdout
586 sys.stderr = old_stderr
587 sys.stdin = old_stdin
588 termui.visible_prompt_func = old_visible_prompt_func
589 termui.hidden_prompt_func = old_hidden_prompt_func
590 termui._getchar = old__getchar_func
591 utils.should_strip_ansi = old_should_strip_ansi # type: ignore
592 _compat.should_strip_ansi = old__compat_should_strip_ansi
593 formatting.FORCED_WIDTH = old_forced_width
594 pdb.Pdb.__init__ = old_pdb_init # type: ignore[method-assign]
596 def invoke(
597 self,
598 cli: Command,
599 args: str | cabc.Sequence[str] | None = None,
600 input: str | bytes | t.IO[t.Any] | None = None,
601 env: cabc.Mapping[str, str | None] | None = None,
602 catch_exceptions: bool | None = None,
603 color: bool = False,
604 **extra: t.Any,
605 ) -> Result:
606 """Invokes a command in an isolated environment. The arguments are
607 forwarded directly to the command line script, the `extra` keyword
608 arguments are passed to the :meth:`~clickpkg.Command.main` function of
609 the command.
611 This returns a :class:`Result` object.
613 :param cli: the command to invoke
614 :param args: the arguments to invoke. It may be given as an iterable
615 or a string. When given as string it will be interpreted
616 as a Unix shell command. More details at
617 :func:`shlex.split`.
618 :param input: the input data for `sys.stdin`.
619 :param env: the environment overrides.
620 :param catch_exceptions: Whether to catch any other exceptions than
621 ``SystemExit``. If :data:`None`, the value
622 from :class:`CliRunner` is used.
623 :param extra: the keyword arguments to pass to :meth:`main`.
624 :param color: whether the output should contain color codes. The
625 application can still override this explicitly.
627 .. versionadded:: 8.2
628 The result object has the ``output_bytes`` attribute with
629 the mix of ``stdout_bytes`` and ``stderr_bytes``, as the user would
630 see it in its terminal.
632 .. versionchanged:: 8.2
633 The result object always returns the ``stderr_bytes`` stream.
635 .. versionchanged:: 8.0
636 The result object has the ``return_value`` attribute with
637 the value returned from the invoked command.
639 .. versionchanged:: 4.0
640 Added the ``color`` parameter.
642 .. versionchanged:: 3.0
643 Added the ``catch_exceptions`` parameter.
645 .. versionchanged:: 3.0
646 The result object has the ``exc_info`` attribute with the
647 traceback if available.
648 """
649 exc_info = None
650 if catch_exceptions is None:
651 catch_exceptions = self.catch_exceptions
653 # Set up fd capture before isolation replaces sys.stdout and sys.stderr.
654 cap_out: _FDCapture | None = None
655 cap_err: _FDCapture | None = None
657 if self.capture == "fd":
658 cap_out = _FDCapture(1)
659 cap_err = _FDCapture(2)
660 try:
661 cap_out.start()
662 cap_err.start()
663 except OSError:
664 cap_out = cap_err = None
666 with self.isolation(input=input, env=env, color=color) as outstreams:
667 # Point the captured streams' fileno() at the saved (original)
668 # fd so that C-level consumers like faulthandler keep working
669 # while fd 1/2 are redirected to the capture tmpfile.
670 if cap_out is not None and cap_err is not None:
671 sys.stdout._original_fd = cap_out.saved_fd # type: ignore[union-attr]
672 sys.stderr._original_fd = cap_err.saved_fd # type: ignore[union-attr]
674 return_value = None
675 exception: BaseException | None = None
676 exit_code = 0
678 if isinstance(args, str):
679 args = shlex.split(args)
681 try:
682 prog_name = extra.pop("prog_name")
683 except KeyError:
684 prog_name = self.get_default_prog_name(cli)
686 try:
687 return_value = cli.main(args=args or (), prog_name=prog_name, **extra)
688 except SystemExit as e:
689 exc_info = sys.exc_info()
690 e_code = t.cast("int | t.Any | None", e.code)
692 if e_code is None:
693 e_code = 0
695 if e_code != 0:
696 exception = e
698 if not isinstance(e_code, int):
699 sys.stdout.write(str(e_code))
700 sys.stdout.write("\n")
701 e_code = 1
703 exit_code = e_code
705 except Exception as e:
706 if not catch_exceptions:
707 raise
708 exception = e
709 exit_code = 1
710 exc_info = sys.exc_info()
711 finally:
712 sys.stdout.flush()
713 sys.stderr.flush()
715 # Stop fd capture and merge the captured bytes into
716 # the stdout/stderr BytesIO streams. BytesIOCopy mirrors
717 # those writes into outstreams[2] automatically.
718 if cap_out is not None and cap_err is not None:
719 fd_out = cap_out.stop()
720 fd_err = cap_err.stop()
721 if fd_out:
722 outstreams[0].write(fd_out)
723 if fd_err:
724 outstreams[1].write(fd_err)
726 stdout = outstreams[0].getvalue()
727 stderr = outstreams[1].getvalue()
728 output = outstreams[2].getvalue()
730 return Result(
731 runner=self,
732 stdout_bytes=stdout,
733 stderr_bytes=stderr,
734 output_bytes=output,
735 return_value=return_value,
736 exit_code=exit_code,
737 exception=exception,
738 exc_info=exc_info, # type: ignore
739 )
741 @contextlib.contextmanager
742 def isolated_filesystem(
743 self, temp_dir: str | os.PathLike[str] | None = None
744 ) -> cabc.Generator[str]:
745 """A context manager that creates a temporary directory and
746 changes the current working directory to it. This isolates tests
747 that affect the contents of the CWD to prevent them from
748 interfering with each other.
750 :param temp_dir: Create the temporary directory under this
751 directory. If given, the created directory is not removed
752 when exiting.
754 .. versionchanged:: 8.0
755 Added the ``temp_dir`` parameter.
756 """
757 cwd = os.getcwd()
758 dt = tempfile.mkdtemp(dir=temp_dir)
759 os.chdir(dt)
761 try:
762 yield dt
763 finally:
764 os.chdir(cwd)
766 if temp_dir is None:
767 import shutil
769 try:
770 shutil.rmtree(dt)
771 except OSError:
772 pass