Coverage for src/amisc/utils.py: 46%

1"""Provides some basic utilities for the package. 

2 

3Includes: 

4 

5- `load_variables`: convenience function for loading RVs from a .json config file 

6- `get_logger`: logging utility with nice formatting 

7""" 

8import json 

9import logging 

10import sys 

11from pathlib import Path 

12 

13from amisc.rv import BaseRV, NormalRV, ScalarRV, UniformRV 

14 

15LOG_FORMATTER = logging.Formatter("%(asctime)s \u2014 [%(levelname)s] \u2014 %(name)-25s \u2014 %(message)s") 

16 

17 

18def load_variables(variables: list[str], file: Path | str) -> list[BaseRV]: 

19 """Load a list of BaseRV objects from a variables json `file`. 

20 

21 :param variables: a list of str ids for variables to find in `file` 

22 :param file: json file to search for variable definitions 

23 :returns rvs: a list of corresponding `BaseRV` objects 

24 """ 

25 with open(Path(file), 'r') as fd: 

26 data = json.load(fd) 

27 

28 rvs = [] 

29 keys = ['id', 'tex', 'description', 'units', 'param_type', 'nominal', 'domain'] 

30 for str_id in variables: 

31 if str_id in data: 

32 var_info = data.get(str_id) 

33 kwargs = {key: var_info.get(key) for key in keys if var_info.get(key)} 

34 match var_info.get('rv_type', 'none'): 

35 case 'uniform_bds': 

36 bds = var_info.get('rv_params') 

37 rvs.append(UniformRV(bds[0], bds[1], **kwargs)) 

38 case 'uniform_pct': 

39 rvs.append(UniformRV(var_info.get('rv_params'), 'pct', **kwargs)) 

40 case 'uniform_tol': 

41 rvs.append(UniformRV(var_info.get('rv_params'), 'tol', **kwargs)) 

42 case 'normal': 

43 mu, std = var_info.get('rv_params') 

44 rvs.append(NormalRV(mu, std, **kwargs)) 

45 case 'none': 

46 # Make a plain stand-in scalar RV object (no uncertainty) 

47 rvs.append(ScalarRV(**kwargs)) 

48 case other: 

49 raise NotImplementedError(f'RV type "{other}" is not known.') 

50 else: 

51 raise ValueError(f'You have requested the variable {str_id}, but it was not found in {file}. ' 

52 f'Please add a definition of {str_id} to {file} or construct it on your own.') 

53 

54 return rvs 

55 

56 

57def get_logger(name: str, stdout=True, log_file: str | Path = None) -> logging.Logger: 

58 """Return a file/stdout logger with the given name. 

59 

60 :param name: the name of the logger to return 

61 :param stdout: whether to add a stdout handler to the logger 

62 :param log_file: add file logging to this file (optional) 

63 :returns: the logger 

64 """ 

65 logger = logging.getLogger(name) 

66 logger.setLevel(logging.DEBUG) 

67 logger.handlers.clear() 

68 if stdout: 

69 std_handler = logging.StreamHandler(sys.stdout) 

70 std_handler.setFormatter(LOG_FORMATTER) 

71 logger.addHandler(std_handler) 

72 if log_file is not None: 

73 f_handler = logging.FileHandler(log_file, mode='a', encoding='utf-8') 

74 f_handler.setLevel(logging.DEBUG) 

75 f_handler.setFormatter(LOG_FORMATTER) 

76 logger.addHandler(f_handler) 

77 

78 return logger