Coverage for src / hallmd / utils.py: 91%

1"""Module to provide utilities for the `hallmd` package. 

2 

3Includes: 

4""" 

5import json 

6import os 

7from pathlib import Path 

8 

9import yaml 

10 

11 

12def _path_in_dict(value, data: dict) -> list: 

13 """Recursively check if a value is in a dictionary and return the list of access keys to the value.""" 

14 if isinstance(data, dict): 

15 for key, v in data.items(): 

16 path = _path_in_dict(value, v) 

17 if path: 

18 return [key] + path 

19 elif data == value: 

20 return [value] # found the value 

21 return [] # value not found 

22 

23 

24def load_thruster(thruster_dir: str | Path, thruster_filename: str = 'thruster.yml') -> dict: 

25 """Load a device configuration from the `device_dir` directory. The `device_file` must be located at 

26 `device_dir/device_name/device_file`. All other files in the directory, if referenced in `device_file`, will 

27 be converted to an absolute path. 

28 

29 !!! Example "Loading a device configuration" 

30 Currently, the only provided device configuration is for the SPT-100 thruster. 

31 ```python 

32 from hallmd.utils import load_thruster 

33 

34 device = load_thruster('devices/SPT-100') 

35 ``` 

36 

37 The format of a device file is as follows 

38 ```yaml 

39 name: MyDevice 

40 geometry: 

41 channel_length: 1 

42 inner_radius: 2 

43 outer_radius: 3 

44 magnetic_field: 

45 file: bfield.csv 

46 shielded: false 

47 ``` 

48 

49 :param device_name: name of the device configuration to load 

50 :param device_file: name of the device configuration file (default: 'device.yml'). Only supported file types are 

51 `.yml` and `.json`. 

52 :param device_dir: directory containing the devices. If None, the `hallmd.devices` directory is used. 

53 :return: dictionary containing the device configuration 

54 """ 

55 thruster_file = Path(thruster_dir) / thruster_filename 

56 

57 with open(thruster_file, 'r', encoding='utf-8') as fd: 

58 if thruster_file.suffix == '.yml': 

59 config = yaml.safe_load(fd) 

60 elif thruster_file.suffix == '.json': 

61 config = json.load(fd) 

62 else: 

63 raise ValueError( 

64 f'Unsupported file type "{thruster_file.suffix}". Only .yml and .json files are supported.' 

65 ) 

66 

67 # Convert all relative paths to absolute paths for loading bfields, etc 

68 for root, _, files in os.walk(thruster_dir): 

69 for file in files: 

70 if file != thruster_file: 

71 # Check if the posix file path from root is in the config (i.e. "./file.csv") 

72 root_path = Path(root) / file # Path like "hallmd/devices/SPT-100/path/to/file.csv" 

73 rel_path = root_path.relative_to(thruster_dir) # Just the path/to/file.csv part (relative) 

74 dict_path = _path_in_dict(rel_path.as_posix(), config) 

75 if len(dict_path) == 0: 

76 # Check if the plain filename is in the config (i.e. file.csv); will only pick first match 

77 dict_path = _path_in_dict(file, config) 

78 

79 if dict_path: 

80 d = config # pointer to the nested location in config 

81 for key in dict_path[:-2]: 

82 d = config[key] 

83 d[dict_path[-2]] = root_path.resolve().as_posix() 

84 

85 return config