10
10
"""
11
11
import os
12
12
import tempfile
13
+ from contextlib import contextmanager
14
+
15
+ try :
16
+ from contextlib import chdir as _chdir
17
+ except ImportError : # PY310
18
+
19
+ @contextmanager
20
+ def _chdir (path ):
21
+ cwd = os .getcwd ()
22
+ os .chdir (path )
23
+ yield
24
+ os .chdir (cwd )
25
+
13
26
14
27
from .deprecated import deprecate_with_version
15
28
@@ -42,7 +55,8 @@ def __init__(self, suffix='', prefix=tempfile.template, dir=None):
42
55
return super ().__init__ (suffix , prefix , dir )
43
56
44
57
45
- class InTemporaryDirectory (tempfile .TemporaryDirectory ):
58
+ @contextmanager
59
+ def InTemporaryDirectory ():
46
60
"""Create, return, and change directory to a temporary directory
47
61
48
62
Notes
@@ -65,18 +79,12 @@ class InTemporaryDirectory(tempfile.TemporaryDirectory):
65
79
>>> os.getcwd() == my_cwd
66
80
True
67
81
"""
68
-
69
- def __enter__ (self ):
70
- self ._pwd = os .getcwd ()
71
- os .chdir (self .name )
72
- return super ().__enter__ ()
73
-
74
- def __exit__ (self , exc , value , tb ):
75
- os .chdir (self ._pwd )
76
- return super ().__exit__ (exc , value , tb )
82
+ with tempfile .TemporaryDirectory () as tmpdir , _chdir (tmpdir ):
83
+ yield tmpdir
77
84
78
85
79
- class InGivenDirectory :
86
+ @contextmanager
87
+ def InGivenDirectory (path = None ):
80
88
"""Change directory to given directory for duration of ``with`` block
81
89
82
90
Useful when you want to use `InTemporaryDirectory` for the final test, but
@@ -98,27 +106,15 @@ class InGivenDirectory:
98
106
You can then look at the temporary file outputs to debug what is happening,
99
107
fix, and finally replace ``InGivenDirectory`` with ``InTemporaryDirectory``
100
108
again.
101
- """
102
-
103
- def __init__ (self , path = None ):
104
- """Initialize directory context manager
105
109
106
- Parameters
107
- ----------
108
- path : None or str, optional
109
- path to change directory to, for duration of ``with`` block.
110
- Defaults to ``os.getcwd()`` if None
111
- """
112
- if path is None :
113
- path = os .getcwd ()
114
- self .path = os .path .abspath (path )
115
-
116
- def __enter__ (self ):
117
- self ._pwd = os .path .abspath (os .getcwd ())
118
- if not os .path .isdir (self .path ):
119
- os .mkdir (self .path )
120
- os .chdir (self .path )
121
- return self .path
122
-
123
- def __exit__ (self , exc , value , tb ):
124
- os .chdir (self ._pwd )
110
+ Parameters
111
+ ----------
112
+ path : None or str, optional
113
+ path to change directory to, for duration of ``with`` block.
114
+ Defaults to ``os.getcwd()`` if None
115
+ """
116
+ if path is None :
117
+ path = os .getcwd ()
118
+ os .makedirs (path , exist_ok = True )
119
+ with _chdir (path ):
120
+ yield os .path .abspath (path )
0 commit comments