Coverage for src/io_collection/load/load_pickle.py: 100%

1import pickle 

2 

3from io_collection.load.load_buffer import load_buffer 

4 

5 

6def load_pickle(location: str, key: str) -> object: 

7 """ 

8 Load key as pickled object from specified location. 

9 

10 Method will load from the S3 bucket if the location begins with the 

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

12 

13 Parameters 

14 ---------- 

15 location 

16 Object location (local path or S3 bucket). 

17 key 

18 Object key ending in `.pkl`. 

19 

20 Returns 

21 ------- 

22 : 

23 Loaded object. 

24 """ 

25 

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

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

28 raise ValueError(message) 

29 

30 buffer = load_buffer(location, key) 

31 return pickle.loads(buffer.read())