| 244 | |
| 245 | @contextmanager |
| 246 | def temp_dir(dir, erase_after=False, with_sentinel=True): |
| 247 | created_by_me = False |
| 248 | if not os.path.exists(dir): |
| 249 | if os.pardir in dir: |
| 250 | raise RuntimeError("workdir contains os.pardir ('..')") |
| 251 | if erase_after and with_sentinel: |
| 252 | closest_dir, fn = get_closest_dir(dir) |
| 253 | sentinel = os.path.join(closest_dir, fn + ".inuse") |
| 254 | open(sentinel, "w").close() |
| 255 | os.makedirs(dir) |
| 256 | created_by_me = True |
| 257 | else: |
| 258 | assert os.path.isdir(dir) |
| 259 | yield |
| 260 | if erase_after and created_by_me: |
| 261 | # erase all files in workdir |
| 262 | shutil.rmtree(dir) |
| 263 | if with_sentinel: |
| 264 | # put dir back as starting point for recursive remove |
| 265 | os.mkdir(dir) |
| 266 | |
| 267 | # also try to erase any other empty directories up to |
| 268 | # sentinel file |
| 269 | os.removedirs(dir) |
| 270 | |
| 271 | # remove sentinel file |
| 272 | os.remove(sentinel) |