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

1from __future__ import annotations 

2 

3import json 

4from pathlib import Path 

5 

6from io_collection.load.load_buffer import _load_buffer_from_s3 

7 

8 

9def load_json(location: str, key: str) -> list | dict: 

10 """ 

11 Load key as dict or list from specified location. 

12 

13 Method will load from 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 ending in `.json`. 

22 

23 Returns 

24 ------- 

25 : 

26 Loaded json. 

27 """ 

28 

29 if not key.endswith(".json"): 

30 message = f"key [ {key} ] must have [ json ] extension" 

31 raise ValueError(message) 

32 

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

34 return _load_json_from_s3(location[5:], key) 

35 return _load_json_from_fs(location, key) 

36 

37 

38def _load_json_from_fs(path: str, key: str) -> list | dict: 

39 """ 

40 Load key as dict or list from local file system. 

41 

42 Parameters 

43 ---------- 

44 path 

45 Local object path. 

46 key 

47 Object key ending in `.json`. 

48 

49 Returns 

50 ------- 

51 : 

52 Loaded json. 

53 """ 

54 

55 full_path = Path(path) / key 

56 with full_path.open("r") as fileobj: 

57 return json.loads(fileobj.read()) 

58 

59 

60def _load_json_from_s3(bucket: str, key: str) -> list | dict: 

61 """ 

62 Load key as dict or list from AWS S3 bucket. 

63 

64 Parameters 

65 ---------- 

66 bucket 

67 AWS S3 bucket name. 

68 key 

69 Object key ending in `.json`. 

70 

71 Returns 

72 ------- 

73 : 

74 Loaded json. 

75 """ 

76 

77 buffer = _load_buffer_from_s3(bucket, key) 

78 return json.loads(buffer.getvalue().decode("utf-8"))