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

1import io 

2from pathlib import Path 

3 

4import boto3 

5 

6 

7def save_buffer( 

8 location: str, key: str, buffer: io.BytesIO, content_type: str = "binary/octet-stream" 

9) -> None: 

10 """ 

11 Save buffer to key at specified location. 

12 

13 Method will save to the S3 bucket if the location begins with the 

14 **s3://** protocol, otherwise it assumes the location is a local path. 

15 

16 Parameters 

17 ---------- 

18 location 

19 Object location (local path or S3 bucket). 

20 key 

21 Object key. 

22 buffer 

23 Content buffer. 

24 content_type 

25 Content type (S3 only). 

26 """ 

27 

28 if location[:5] == "s3://": 

29 _save_buffer_to_s3(location[5:], key, buffer, content_type) 

30 else: 

31 _save_buffer_to_fs(location, key, buffer) 

32 

33 

34def _save_buffer_to_fs(path: str, key: str, buffer: io.BytesIO) -> None: 

35 """ 

36 Save buffer to key on local file system. 

37 

38 Directories that do not exist will be created. 

39 

40 Parameters 

41 ---------- 

42 path 

43 Local object path. 

44 key 

45 Object key. 

46 buffer 

47 Content buffer. 

48 """ 

49 

50 full_path = Path(path) / key 

51 full_path.parent.mkdir(parents=True, exist_ok=True) 

52 

53 with full_path.open("wb") as fileobj: 

54 fileobj.write(buffer.getvalue()) 

55 

56 

57def _save_buffer_to_s3(bucket: str, key: str, buffer: io.BytesIO, content_type: str) -> None: 

58 """ 

59 Save buffer to key in AWS S3 bucket. 

60 

61 Parameters 

62 ---------- 

63 bucket 

64 AWS S3 bucket name. 

65 key 

66 Object key. 

67 buffer 

68 Content buffer. 

69 content_type 

70 Content type. 

71 """ 

72 

73 s3_client = boto3.client("s3") 

74 s3_client.put_object(Bucket=bucket, Key=key, Body=buffer.getvalue(), ContentType=content_type)