File size: 2,428 Bytes
8e04e6f | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 | import os
import pathlib
import tarfile
def flatten_directory(
directory: pathlib.Path,
replace_ending: bool = False,
old_ending: str = ".ent",
new_ending: str = ".pdb",
) -> None:
"""
Flattens the directory structure by moving all files to the top level and optionally replacing file endings.
Args:
directory: The directory to flatten.
replace_ending: A boolean flag indicating whether to replace the file ending. Default is False.
old_ending: The old file ending to be replaced. Default is '.ent'.
new_ending: The new file ending to replace the old one with. Default is '.pdb'.
Returns:
None
"""
for file_path in directory.rglob("*"):
if file_path.is_file():
new_file_path = directory / (
file_path.stem
+ (
file_path.suffix.replace(old_ending, new_ending)
if replace_ending
else file_path.suffix
)
)
file_path.rename(new_file_path)
for dir_path in directory.rglob("*"):
if dir_path.is_dir() and not any(dir_path.iterdir()):
dir_path.rmdir()
def extract_archive(
tar_file: pathlib.Path, destination_dir: pathlib.Path, extracted_dir_name: str
) -> pathlib.Path:
"""
Extracts the contents of the tar archive to a specified directory with a given name.
Args:
tar_file: The path to the tar archive file.
destination_dir: The directory where the archive contents will be extracted.
extracted_dir_name: The desired name for the directory containing the extracted contents.
Returns:
The path to the directory containing the extracted contents.
"""
# Create the destination directory if it doesn't exist
destination_dir.mkdir(parents=True, exist_ok=True)
# Create the directory for the extracted contents
renamed_extracted_dir = destination_dir / extracted_dir_name
# Extract the contents of the archive to the extracted directory
with tarfile.open(tar_file, "r:gz") as tar:
top_dir = os.path.commonpath(tar.getnames())
tar.extractall(path=destination_dir, filter="data")
# Remove the tar archive file
tar_file.unlink()
extracted_dir = pathlib.Path(destination_dir) / top_dir
extracted_dir.rename(renamed_extracted_dir)
return renamed_extracted_dir
|