Coverage for src/arcade_collection/output/extract_tick_json.py: 100%

1from __future__ import annotations 

2 

3import json 

4from typing import TYPE_CHECKING 

5 

6import numpy as np 

7 

8if TYPE_CHECKING: 

9 import tarfile 

10 

11 

12def extract_tick_json( 

13 tar: tarfile.TarFile, 

14 key: str, 

15 tick: float, 

16 extension: str | None = None, 

17 field: str | None = None, 

18) -> list: 

19 """ 

20 Extract json for specified tick from tar archive. 

21 

22 For v3 simulations, each tick is saved as a separate json file in the 

23 archive. The file names are formatted as ``<key>_<tick>.json`` or 

24 ``<key>_<tick>.<extension>.json`` where tick is padded to six digits. Use an 

25 integer tick to extract the desired tick from the archive. 

26 

27 For v2 simulations, all ticks are saved to the same file in the archive. The 

28 file name is given as ``<key>.json`` or ``<key>.<extension>.json``. Use a 

29 float tick and field to extract the desired tick from the archive. 

30 

31 .. code-block:: python 

32 

33 { 

34 "timepoints": [ 

35 { 

36 "time": <time>, "field": <field>, ... 

37 } 

38 ] 

39 } 

40 

41 Parameters 

42 ---------- 

43 tar 

44 Tar archive. 

45 key 

46 Simulation key. 

47 tick 

48 Tick to extract. 

49 extension 

50 Additional extension in file name. 

51 field 

52 Field in json to extract (only used with float ticks). 

53 

54 Returns 

55 ------- 

56 : 

57 Archive contents for specified tick. 

58 """ 

59 

60 formatted_tick = f"_{tick:06d}" if isinstance(tick, (int, np.integer)) else "" 

61 

62 if extension is None: 

63 member = tar.extractfile(f"{key}{formatted_tick}.json") 

64 else: 

65 member = tar.extractfile(f"{key}{formatted_tick}.{extension}.json") 

66 

67 if member is None: 

68 message = "File does not exist in archive." 

69 raise RuntimeError(message) 

70 

71 tick_json = json.loads(member.read().decode("utf-8")) 

72 

73 if isinstance(tick, float): 

74 tick_json = next(item for item in tick_json["timepoints"] if item["time"] == tick)[field] 

75 

76 return tick_json