pferd/PFERD/tmp_dir.py

71 lines
2.1 KiB
Python
Raw Normal View History

2020-04-20 12:08:52 +02:00
"""Helper functions and classes for temporary folders."""
import logging
2020-04-20 12:08:52 +02:00
import shutil
2020-04-20 14:28:22 +02:00
from pathlib import Path
from types import TracebackType
from typing import Optional, Type
2020-04-20 12:08:52 +02:00
logger = logging.getLogger(__name__)
2020-04-20 12:08:52 +02:00
2020-04-20 14:28:22 +02:00
class TmpDir():
2020-04-20 12:08:52 +02:00
"""A temporary folder that can create files or nested temp folders."""
2020-04-20 14:28:22 +02:00
def __init__(self, path: Path):
2020-04-20 12:08:52 +02:00
"""Create a new temporary folder for the given path."""
2020-04-20 14:28:22 +02:00
self._counter = 0
self._path = path
2020-04-20 12:08:52 +02:00
def __str__(self) -> str:
"""Format the folder as a string."""
return f"Folder at {self.path}"
2020-04-20 14:28:22 +02:00
def __enter__(self) -> 'TmpDir':
2020-04-20 12:08:52 +02:00
"""Context manager entry function."""
return self
def __exit__(self,
type: Optional[Type[BaseException]],
value: Optional[BaseException],
traceback: Optional[TracebackType]) -> Optional[bool]:
"""Context manager exit function. Calls cleanup()."""
self.cleanup()
return None
2020-04-20 14:28:22 +02:00
@property
def path(self) -> Path:
"""Return the path of this folder."""
return self._path
def new_file(self, prefix: Optional[str] = None) -> Path:
2020-04-20 12:08:52 +02:00
"""Return a unique path inside the folder, but don't create a file."""
2020-04-20 14:28:22 +02:00
name = f"{prefix if prefix else 'tmp'}-{self._inc_and_get_counter():03}"
logger.debug(f"Creating temp file '{name}'")
2020-04-20 12:08:52 +02:00
return self.path.joinpath(name)
2020-04-20 14:28:22 +02:00
def new_folder(self, prefix: Optional[str] = None) -> 'TmpDir':
2020-04-20 12:08:52 +02:00
"""Create a new nested temporary folder and return its path."""
2020-04-20 14:28:22 +02:00
name = f"{prefix if prefix else 'tmp'}-{self._inc_and_get_counter():03}"
2020-04-20 12:08:52 +02:00
sub_path = self.path.joinpath(name)
sub_path.mkdir(parents=True)
logger.debug(f"Creating temp dir '{name}' at {sub_path}")
2020-04-20 14:28:22 +02:00
return TmpDir(sub_path)
2020-04-20 12:08:52 +02:00
def cleanup(self) -> None:
"""Delete this folder and all contained files."""
logger.debug(f"Deleting temp folder {self.path}")
2020-04-20 12:08:52 +02:00
shutil.rmtree(self.path.absolute())
2020-04-20 14:28:22 +02:00
def _inc_and_get_counter(self) -> int:
2020-04-20 12:08:52 +02:00
"""Get and increment the counter by one."""
2020-04-20 14:28:22 +02:00
counter = self._counter
self._counter += 1
2020-04-20 12:08:52 +02:00
return counter