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

1import io 

2import pickle 

3 

4from io_collection.save.save_buffer import save_buffer 

5 

6 

7def save_pickle(location: str, key: str, obj: object) -> None: 

8 """ 

9 Save pickled object to key at specified location. 

10 

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

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

13 

14 Parameters 

15 ---------- 

16 location 

17 Object location (local path or S3 bucket). 

18 key 

19 Object key ending in `.pkl`. 

20 obj 

21 Object to save. 

22 """ 

23 

24 if not key.endswith(".pkl"): 

25 message = f"key [ {key} ] must have [ pkl ] extension" 

26 raise ValueError(message) 

27 

28 with io.BytesIO() as buffer: 

29 pickle.dump(obj, buffer) 

30 save_buffer(location, key, buffer)