Skip to content

Commit bb376b1

Browse files
committed
RF: Reimplement tmpdirs tools as contextmanagers
Simplifying these tools as wrapped functions makes the logic easier to read.
1 parent 63b5586 commit bb376b1

File tree

1 file changed

+30
-34
lines changed

1 file changed

+30
-34
lines changed

nibabel/tmpdirs.py

Lines changed: 30 additions & 34 deletions
Original file line numberDiff line numberDiff line change
@@ -10,6 +10,19 @@
1010
"""
1111
import os
1212
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+
1326

1427
from .deprecated import deprecate_with_version
1528

@@ -42,7 +55,8 @@ def __init__(self, suffix='', prefix=tempfile.template, dir=None):
4255
return super().__init__(suffix, prefix, dir)
4356

4457

45-
class InTemporaryDirectory(tempfile.TemporaryDirectory):
58+
@contextmanager
59+
def InTemporaryDirectory():
4660
"""Create, return, and change directory to a temporary directory
4761
4862
Notes
@@ -65,18 +79,12 @@ class InTemporaryDirectory(tempfile.TemporaryDirectory):
6579
>>> os.getcwd() == my_cwd
6680
True
6781
"""
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
7784

7885

79-
class InGivenDirectory:
86+
@contextmanager
87+
def InGivenDirectory(path=None):
8088
"""Change directory to given directory for duration of ``with`` block
8189
8290
Useful when you want to use `InTemporaryDirectory` for the final test, but
@@ -98,27 +106,15 @@ class InGivenDirectory:
98106
You can then look at the temporary file outputs to debug what is happening,
99107
fix, and finally replace ``InGivenDirectory`` with ``InTemporaryDirectory``
100108
again.
101-
"""
102-
103-
def __init__(self, path=None):
104-
"""Initialize directory context manager
105109
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

Comments
 (0)