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

1from __future__ import annotations 

2 

3import json 

4import re 

5 

6from io_collection.save.save_text import save_text 

7 

8 

9def save_json(location: str, key: str, obj: dict | list, levels: int = 5) -> None: 

10 """ 

11 Save dict or list as json to key at specified location. 

12 

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

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

15 

16 The json is pretty printed with an indentation level of 2. For lists, 

17 default pretty printing places each item on a new line which can produce 

18 very long files. This method will attempt map lists up to the specified 

19 number of items onto a single line. 

20 

21 Parameters 

22 ---------- 

23 location 

24 Object location (local path or S3 bucket). 

25 key 

26 Object key ending in `.json`. 

27 obj 

28 Dict or list to save. 

29 levels 

30 Maximum number of items for single line list. 

31 """ 

32 

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

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

35 raise ValueError(message) 

36 

37 contents = json.dumps(obj, indent=2) 

38 

39 for level in range(1, levels + 1): 

40 match_pattern = ( 

41 r"\[" + ",".join([r"\n\s*([A-z\"\-\d\.]+)" for i in range(level)]) + r"\n\s*\]" 

42 ) 

43 sub_pattern = "[" + ",".join([f"\\{i + 1}" for i in range(level)]) + "]" 

44 pattern = re.compile(match_pattern) 

45 contents = pattern.sub(sub_pattern, contents) 

46 

47 save_text(location, key, contents, "application/json")