232 lines
7.9 KiB
Python
232 lines
7.9 KiB
Python
import os
|
|
import warnings
|
|
from enum import Enum
|
|
from pathlib import Path
|
|
from typing import List, Union
|
|
|
|
try:
|
|
from typing import Literal
|
|
except ImportError:
|
|
from typing_extensions import Literal
|
|
|
|
|
|
def check_path_suffix(path_str: str,
|
|
allowed_suffix: Union[str, List[str]] = '') -> bool:
|
|
"""Check whether the suffix of the path is allowed.
|
|
|
|
Args:
|
|
path_str (str):
|
|
Path to check.
|
|
allowed_suffix (List[str], optional):
|
|
What extension names are allowed.
|
|
Offer a list like ['.jpg', ',jpeg'].
|
|
When it's [], all will be received.
|
|
Use [''] then directory is allowed.
|
|
Defaults to [].
|
|
|
|
Returns:
|
|
bool:
|
|
True: suffix test passed
|
|
False: suffix test failed
|
|
"""
|
|
if isinstance(allowed_suffix, str):
|
|
allowed_suffix = [allowed_suffix]
|
|
pathinfo = Path(path_str)
|
|
suffix = pathinfo.suffix.lower()
|
|
if len(allowed_suffix) == 0:
|
|
return True
|
|
if pathinfo.is_dir():
|
|
if '' in allowed_suffix:
|
|
return True
|
|
else:
|
|
return False
|
|
else:
|
|
for index, tmp_suffix in enumerate(allowed_suffix):
|
|
if not tmp_suffix.startswith('.'):
|
|
tmp_suffix = '.' + tmp_suffix
|
|
allowed_suffix[index] = tmp_suffix.lower()
|
|
if suffix in allowed_suffix:
|
|
return True
|
|
else:
|
|
return False
|
|
|
|
|
|
class Existence(Enum):
|
|
"""State of file existence."""
|
|
FileExist = 0
|
|
DirectoryExistEmpty = 1
|
|
DirectoryExistNotEmpty = 2
|
|
MissingParent = 3
|
|
DirectoryNotExist = 4
|
|
FileNotExist = 5
|
|
|
|
|
|
def check_path_existence(
|
|
path_str: str,
|
|
path_type: Literal['file', 'dir', 'auto'] = 'auto',
|
|
) -> Existence:
|
|
"""Check whether a file or a directory exists at the expected path.
|
|
|
|
Args:
|
|
path_str (str):
|
|
Path to check.
|
|
path_type (Literal[, optional):
|
|
What kind of file do we expect at the path.
|
|
Choose among `file`, `dir`, `auto`.
|
|
Defaults to 'auto'. path_type = path_type.lower()
|
|
|
|
Raises:
|
|
KeyError: if `path_type` conflicts with `path_str`
|
|
|
|
Returns:
|
|
Existence:
|
|
0. FileExist: file at path_str exists.
|
|
1. DirectoryExistEmpty: folder at path exists and.
|
|
2. DirectoryExistNotEmpty: folder at path_str exists and not empty.
|
|
3. MissingParent: its parent doesn't exist.
|
|
4. DirectoryNotExist: expect a folder at path_str, but not found.
|
|
5. FileNotExist: expect a file at path_str, but not found.
|
|
"""
|
|
path_type = path_type.lower()
|
|
assert path_type in {'file', 'dir', 'auto'}
|
|
pathinfo = Path(path_str)
|
|
if not pathinfo.parent.is_dir():
|
|
return Existence.MissingParent
|
|
suffix = pathinfo.suffix.lower()
|
|
if path_type == 'dir' or\
|
|
path_type == 'auto' and suffix == '':
|
|
if pathinfo.is_dir():
|
|
if len(os.listdir(path_str)) == 0:
|
|
return Existence.DirectoryExistEmpty
|
|
else:
|
|
return Existence.DirectoryExistNotEmpty
|
|
else:
|
|
return Existence.DirectoryNotExist
|
|
elif path_type == 'file' or\
|
|
path_type == 'auto' and suffix != '':
|
|
if pathinfo.is_file():
|
|
return Existence.FileExist
|
|
elif pathinfo.is_dir():
|
|
if len(os.listdir(path_str)) == 0:
|
|
return Existence.DirectoryExistEmpty
|
|
else:
|
|
return Existence.DirectoryExistNotEmpty
|
|
if path_str.endswith('/'):
|
|
return Existence.DirectoryNotExist
|
|
else:
|
|
return Existence.FileNotExist
|
|
|
|
|
|
def prepare_output_path(output_path: str,
|
|
allowed_suffix: List[str] = [],
|
|
tag: str = 'output file',
|
|
path_type: Literal['file', 'dir', 'auto'] = 'auto',
|
|
overwrite: bool = True) -> None:
|
|
"""Check output folder or file.
|
|
|
|
Args:
|
|
output_path (str): could be folder or file.
|
|
allowed_suffix (List[str], optional):
|
|
Check the suffix of `output_path`. If folder, should be [] or [''].
|
|
If could both be folder or file, should be [suffixs..., ''].
|
|
Defaults to [].
|
|
tag (str, optional): The `string` tag to specify the output type.
|
|
Defaults to 'output file'.
|
|
path_type (Literal[, optional):
|
|
Choose `file` for file and `dir` for folder.
|
|
Choose `auto` if allowed to be both.
|
|
Defaults to 'auto'.
|
|
overwrite (bool, optional):
|
|
Whether overwrite the existing file or folder.
|
|
Defaults to True.
|
|
|
|
Raises:
|
|
FileNotFoundError: suffix does not match.
|
|
FileExistsError: file or folder already exists and `overwrite` is
|
|
False.
|
|
|
|
Returns:
|
|
None
|
|
"""
|
|
if path_type.lower() == 'dir':
|
|
allowed_suffix = []
|
|
exist_result = check_path_existence(output_path, path_type=path_type)
|
|
if exist_result == Existence.MissingParent:
|
|
warnings.warn(
|
|
f'The parent folder of {tag} does not exist: {output_path},' +
|
|
f' will make dir {Path(output_path).parent.absolute().__str__()}')
|
|
os.makedirs(
|
|
Path(output_path).parent.absolute().__str__(), exist_ok=True)
|
|
|
|
elif exist_result == Existence.DirectoryNotExist:
|
|
os.mkdir(output_path)
|
|
print(f'Making directory {output_path} for saving results.')
|
|
elif exist_result == Existence.FileNotExist:
|
|
suffix_matched = \
|
|
check_path_suffix(output_path, allowed_suffix=allowed_suffix)
|
|
if not suffix_matched:
|
|
raise FileNotFoundError(
|
|
f'The {tag} should be {", ".join(allowed_suffix)}: '
|
|
f'{output_path}.')
|
|
elif exist_result == Existence.FileExist:
|
|
if not overwrite:
|
|
raise FileExistsError(
|
|
f'{output_path} exists (set overwrite = True to overwrite).')
|
|
else:
|
|
print(f'Overwriting {output_path}.')
|
|
elif exist_result == Existence.DirectoryExistEmpty:
|
|
pass
|
|
elif exist_result == Existence.DirectoryExistNotEmpty:
|
|
if not overwrite:
|
|
raise FileExistsError(
|
|
f'{output_path} is not empty (set overwrite = '
|
|
'True to overwrite the files).')
|
|
else:
|
|
print(f'Overwriting {output_path} and its files.')
|
|
else:
|
|
raise FileNotFoundError(f'No Existence type for {output_path}.')
|
|
|
|
|
|
def check_input_path(
|
|
input_path: str,
|
|
allowed_suffix: List[str] = [],
|
|
tag: str = 'input file',
|
|
path_type: Literal['file', 'dir', 'auto'] = 'auto',
|
|
):
|
|
"""Check input folder or file.
|
|
|
|
Args:
|
|
input_path (str): input folder or file path.
|
|
allowed_suffix (List[str], optional):
|
|
Check the suffix of `input_path`. If folder, should be [] or [''].
|
|
If could both be folder or file, should be [suffixs..., ''].
|
|
Defaults to [].
|
|
tag (str, optional): The `string` tag to specify the output type.
|
|
Defaults to 'output file'.
|
|
path_type (Literal[, optional):
|
|
Choose `file` for file and `directory` for folder.
|
|
Choose `auto` if allowed to be both.
|
|
Defaults to 'auto'.
|
|
|
|
Raises:
|
|
FileNotFoundError: file does not exists or suffix does not match.
|
|
|
|
Returns:
|
|
None
|
|
"""
|
|
if path_type.lower() == 'dir':
|
|
allowed_suffix = []
|
|
exist_result = check_path_existence(input_path, path_type=path_type)
|
|
|
|
if exist_result in [
|
|
Existence.FileExist, Existence.DirectoryExistEmpty,
|
|
Existence.DirectoryExistNotEmpty
|
|
]:
|
|
suffix_matched = \
|
|
check_path_suffix(input_path, allowed_suffix=allowed_suffix)
|
|
if not suffix_matched:
|
|
raise FileNotFoundError(
|
|
f'The {tag} should be {", ".join(allowed_suffix)}:' +
|
|
f'{input_path}.')
|
|
else:
|
|
raise FileNotFoundError(f'The {tag} does not exist: {input_path}.') |