Coverage for src/io_collection/save/save_tar.py: 100%

1import io 

2import tarfile 

3 

4from io_collection.keys.check_key import check_key 

5from io_collection.load.load_buffer import load_buffer 

6from io_collection.load.load_tar import load_tar 

7from io_collection.save.save_buffer import save_buffer 

8 

9 

10def save_tar(location: str, key: str, contents: list[str]) -> None: 

11 """ 

12 Save tar archive to key at specified location. 

13 

14 Method will save to the S3 bucket if the location begins with the **s3://** 

15 protocol, otherwise it assumes the location is a local path. 

16 

17 If the archive already exists, all objects in the existing archive will be 

18 copied into the new archive, along with the new contents. 

19 

20 Method currently only supports `xz` compression. 

21 

22 Parameters 

23 ---------- 

24 location 

25 Object location (local path or S3 bucket). 

26 key 

27 Object key ending in `.tar.xz`. 

28 contents 

29 List of content keys to include in the archive. 

30 """ 

31 

32 if not key.endswith(".tar.xz"): 

33 message = f"key [ {key} ] must have [ tar.xz ] extension" 

34 raise ValueError(message) 

35 

36 existing_tar = load_tar(location, key) if check_key(location, key) else None 

37 

38 with io.BytesIO() as buffer: 

39 with tarfile.open(fileobj=buffer, mode="w:xz") as tar: 

40 if existing_tar is not None: 

41 for member in existing_tar.getmembers(): 

42 tar.addfile(member, existing_tar.extractfile(member.name)) 

43 

44 for content_key in contents: 

45 content = load_buffer(location, content_key) 

46 info = tarfile.TarInfo(content_key.split("/")[-1]) 

47 info.size = content.getbuffer().nbytes 

48 tar.addfile(info, fileobj=content) 

49 

50 save_buffer(location, key, buffer)