Clean up temp_folder

This commit is contained in:
I-Al-Istannen 2020-04-20 14:28:22 +02:00
parent 53ad1c924b
commit a3d6dc7873

View File

@ -1,25 +1,24 @@
"""Helper functions and classes for temporary folders."""
from typing import Optional, Type
from types import TracebackType
import pathlib
import shutil
from pathlib import Path
from types import TracebackType
from typing import Optional, Type
class TempFolder():
class TmpDir():
"""A temporary folder that can create files or nested temp folders."""
def __init__(self, path: pathlib.Path):
def __init__(self, path: Path):
"""Create a new temporary folder for the given path."""
self.counter = 0
self.path = path
self._counter = 0
self._path = path
def __str__(self) -> str:
"""Format the folder as a string."""
return f"Folder at {self.path}"
def __enter__(self) -> 'TempFolder':
def __enter__(self) -> 'TmpDir':
"""Context manager entry function."""
return self
@ -31,26 +30,31 @@ class TempFolder():
self.cleanup()
return None
def new_file(self) -> pathlib.Path:
@property
def path(self) -> Path:
"""Return the path of this folder."""
return self._path
def new_file(self, prefix: Optional[str] = None) -> Path:
"""Return a unique path inside the folder, but don't create a file."""
name = f"tmp-{self.__inc_and_get_counter():03}"
name = f"{prefix if prefix else 'tmp'}-{self._inc_and_get_counter():03}"
return self.path.joinpath(name)
def new_folder(self, prefix: str = "") -> 'TempFolder':
def new_folder(self, prefix: Optional[str] = None) -> 'TmpDir':
"""Create a new nested temporary folder and return its path."""
name = f"{prefix if prefix else 'tmp'}-{self.__inc_and_get_counter():03}"
name = f"{prefix if prefix else 'tmp'}-{self._inc_and_get_counter():03}"
sub_path = self.path.joinpath(name)
sub_path.mkdir(parents=True)
return TempFolder(sub_path)
return TmpDir(sub_path)
def cleanup(self) -> None:
"""Delete this folder and all contained files."""
shutil.rmtree(self.path.absolute())
def __inc_and_get_counter(self) -> int:
def _inc_and_get_counter(self) -> int:
"""Get and increment the counter by one."""
counter = self.counter
self.counter += 1
counter = self._counter
self._counter += 1
return counter