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}.')