Coverage for foxplot / decode.py: 94%

1#!/usr/bin/env python3 

2# -*- coding: utf-8 -*- 

3# 

4# SPDX-License-Identifier: Apache-2.0 

5 

6"""Decode a series of dictionaries from file.""" 

7 

8import json 

9import sys 

10from pathlib import PosixPath 

11from typing import Generator, Union 

12 

13import mpacklog 

14 

15from .exceptions import FoxplotError 

16 

17 

18def decode(file_path: Union[str, PosixPath]) -> Generator[dict, None, None]: 

19 """Unpack a series of dictionaries from a given file. 

20 

21 Args: 

22 file_path: Path to the file to read from (can be "stdin"). 

23 

24 Yields: 

25 Unpacked dictionaries. 

26 """ 

27 file_path = str(file_path) 

28 if file_path == "stdin": 

29 yield from decode_json(file=sys.stdin) 

30 elif file_path.endswith(".json") or file_path.endswith(".jsonl"): 

31 with open(file_path, "r", encoding="utf-8") as file: 

32 yield from decode_json(file=file) 

33 elif file_path.endswith(".mpack"): 

34 yield from mpacklog.read_log(file_path) 

35 else: # unknown file extension 

36 raise FoxplotError(f"Unknown file type in '{file_path}'") 

37 

38 

39def decode_json(file, chunk_size=100_000) -> Generator[dict, None, None]: 

40 """Decode dictionaries from a line-delimited JSON file. 

41 

42 Args: 

43 file: File stream (for instance ``sys.stdin``). 

44 chunk_size: Number of bytes read at a time from the standard input. 

45 

46 Yields: 

47 dict: Dictionary read from file. 

48 """ 

49 decoder = json.JSONDecoder() 

50 buffer = "" 

51 while True: 

52 data = file.read(chunk_size) 

53 if not data: # end of file 

54 break 

55 buffer += data 

56 while buffer: 

57 try: 

58 result, index = decoder.raw_decode(buffer) 

59 buffer = buffer[index:].lstrip() 

60 yield result 

61 except json.JSONDecodeError: 

62 break