Coverage for src/cvx/json/file.py: 100%

1# Copyright 2023 Stanford University Convex Optimization Group 

2# 

3# Licensed under the Apache License, Version 2.0 (the "License"); 

4# you may not use this file except in compliance with the License. 

5# You may obtain a copy of the License at 

6# 

7# http://www.apache.org/licenses/LICENSE-2.0 

8# 

9# Unless required by applicable law or agreed to in writing, software 

10# distributed under the License is distributed on an "AS IS" BASIS, 

11# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 

12# See the License for the specific language governing permissions and 

13# limitations under the License. 

14"""Tools to support working with json files.""" 

15 

16import json 

17from collections.abc import Iterable 

18from os import PathLike 

19from typing import Any 

20 

21import numpy as np 

22import numpy.typing as npt 

23 

24from .numpyencoder import NumpyEncoder 

25 

26FILE = str | bytes | PathLike 

27MATRIX = npt.NDArray[Any] 

28DATA = dict[str, Any] 

29 

30 

31def read_json(json_file: FILE) -> DATA: 

32 """Read a JSON file and convert its contents to a dictionary. 

33 

34 Iterables in the JSON data are converted to numpy arrays. 

35 

36 Args: 

37 json_file: Path to the JSON file to read 

38 

39 Returns: 

40 A dictionary containing the data from the JSON file 

41 """ 

42 with open(json_file) as f: 

43 json_data = json.load(f) 

44 d = {} 

45 for name, data in json_data.items(): 

46 if isinstance(data, Iterable): 

47 d[name] = np.asarray(data) 

48 else: 

49 d[name] = data 

50 

51 return d 

52 

53 

54def write_json(json_file: FILE, data: DATA) -> None: 

55 """Write data to a JSON file. 

56 

57 Uses NumpyEncoder to handle numpy data types. 

58 

59 Args: 

60 json_file: Path to the JSON file to write 

61 data: Dictionary of data to write to the file 

62 """ 

63 with open(json_file, "w") as f: 

64 json.dump(data, f, cls=NumpyEncoder)