def download_model(self, dir_path: str) -> str:
if not dir_path.endswith('.tar.gz'):
return dir_path
local_dir_path = str(self.download_manager.get_local_file(
dir_path, auto_uncompress=False
)).replace('.tar.gz', '')
copy_directory_with_source_meta(dir_path, local_dir_path)
return local_dir_path
def test_should_not_skip_if_source_url_is_different(
self, target_directory: Path, source_path: Path):
target_directory.mkdir(parents=True, exist_ok=True)
target_directory.joinpath(SOURCE_URL_META_FILENAME).write_text("other")
target_directory.joinpath(MODEL_FILE_1).write_bytes(MODEL_DATA_1)
source_path.mkdir(parents=True, exist_ok=True)
source_path.joinpath(MODEL_FILE_1).write_bytes(MODEL_DATA_2)
copy_directory_with_source_meta(source_url=str(source_path),
target_directory=str(target_directory),
force=False)
assert target_directory.joinpath(
MODEL_FILE_1).read_bytes() == MODEL_DATA_2
def test_should_copy_to_not_yet_existing_directory(self,
target_directory: Path,
source_path: Path):
source_path.mkdir(parents=True, exist_ok=True)
source_path.joinpath(MODEL_FILE_1).write_bytes(MODEL_DATA_1)
copy_directory_with_source_meta(source_url=str(source_path),
target_directory=str(target_directory),
force=False)
assert target_directory.joinpath(
MODEL_FILE_1).read_bytes() == MODEL_DATA_1
assert (target_directory.joinpath(
SOURCE_URL_META_FILENAME).read_text() == str(source_path))
def test_should_extract_from_zip(self, target_directory: Path,
source_path: Path):
source_path.mkdir(parents=True, exist_ok=True)
zip_file_path = source_path.joinpath('archive1.zip')
with zipfile.ZipFile(str(zip_file_path), mode='w') as zip_file:
zip_file.writestr(MODEL_FILE_1, MODEL_DATA_1)
copy_directory_with_source_meta(source_url=str(zip_file_path),
target_directory=str(target_directory),
force=False)
assert target_directory.joinpath(
MODEL_FILE_1).read_bytes() == MODEL_DATA_1
assert (target_directory.joinpath(
SOURCE_URL_META_FILENAME).read_text() == str(zip_file_path))
def test_should_extract_from_tar_gz(self, target_directory: Path,
source_path: Path):
source_path.mkdir(parents=True, exist_ok=True)
tar_gz_file_path = source_path.joinpath('archive1.tar.gz')
with tarfile.open(str(tar_gz_file_path), mode='w:gz') as tar_file:
buf = BytesIO(MODEL_DATA_1)
tar_info = tarfile.TarInfo(name=MODEL_FILE_1)
tar_info.size = len(buf.getvalue())
tar_file.addfile(tar_info, buf)
copy_directory_with_source_meta(source_url=str(tar_gz_file_path),
target_directory=str(target_directory),
force=False)
assert target_directory.joinpath(
MODEL_FILE_1).read_bytes() == MODEL_DATA_1
assert (target_directory.joinpath(
SOURCE_URL_META_FILENAME).read_text() == str(tar_gz_file_path))