Coverage for src/cvx/bson/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 bson files and strings.""" 

15 

16from __future__ import annotations 

17 

18from os import PathLike 

19from typing import Any 

20 

21import numpy.typing as npt 

22import pandas as pd 

23 

24import bson 

25 

26from .io import decode, encode 

27 

28FILE = str | bytes | PathLike 

29MATRIX = npt.NDArray[Any] | pd.DataFrame 

30MATRICES = dict[str, MATRIX] 

31 

32 

33def read_bson(file: FILE) -> MATRICES: 

34 """Read a bson file and prepare the bson_document needed. 

35 

36 Args: 

37 file: Path to the BSON file to read 

38 

39 Returns: 

40 A dictionary of numpy arrays 

41 """ 

42 with open(file, mode="rb") as openfile: 

43 # Reading from bson file 

44 return from_bson(bson_str=openfile.read()) 

45 

46 

47def write_bson(file: FILE, data: MATRICES): 

48 """Write dictionary into a bson file. 

49 

50 Args: 

51 file: file 

52 data: dictionary of numpy arrays 

53 """ 

54 with open(file=file, mode="wb") as bson_file: 

55 bson_file.write(to_bson(data=data)) 

56 

57 

58def to_bson(data: MATRICES) -> bytes: 

59 """Convert a dictionary of numpy arrays into a bson string. 

60 

61 Args: 

62 data: dictionary of numpy arrays 

63 """ 

64 return bytes(bson.dumps({name: encode(matrix) for name, matrix in data.items()})) 

65 

66 

67def from_bson(bson_str: bytes) -> MATRICES: 

68 """Convert a bson string into a dictionary of numpy arrays.""" 

69 return {name: decode(value) for name, value in bson.loads(bson_str).items()}