Coverage for console_utils.py: 95%

1""" 

2Utility functions for console applications. 

3""" 

4import os 

5import sys 

6import logging 

7 

8from typing import Union 

9from typing import List 

10 

11from colorlog import ColoredFormatter 

12 

13 

14def setup_logging( 

15 path: str = "/tmp", 

16 file_name: Union[str, List[str]] = "file.log", 

17 level: int = logging.INFO, 

18 use_global: bool = True, 

19) -> logging.Logger: 

20 """Setup logging configuration. 

21 Output logs to a file and stdout.""" 

22 

23 if isinstance(file_name, str): 

24 file_names = [file_name] 

25 else: 

26 file_names = file_name 

27 

28 if not os.path.exists(path): 

29 os.makedirs(path) 

30 

31 logger_name = file_names[0] 

32 if use_global: 

33 logger = logging.getLogger() 

34 else: 

35 logger = logging.getLogger(logger_name) 

36 logger.setLevel(level) 

37 logger.handlers.clear() 

38 

39 formatter = ColoredFormatter( 

40 fmt="[%(asctime)s.%(msecs)03d] %(log_color)s%(levelname)s: %(message)s", 

41 datefmt="%Y-%m-%d %H:%M:%S", 

42 log_colors={ 

43 "DEBUG": "cyan", 

44 "INFO": "white", 

45 # "INFO": None, 

46 "WARNING": "yellow", 

47 "ERROR": "red", 

48 "CRITICAL": "bold_red", 

49 } 

50 ) 

51 

52 for file_name in file_names: 

53 file_handler = logging.FileHandler(f"{path}/{file_name}") 

54 file_handler.setFormatter(formatter) 

55 logger.addHandler(file_handler) 

56 

57 stdout_handler = logging.StreamHandler(sys.stdout) 

58 stdout_handler.setFormatter(formatter) 

59 logger.addHandler(stdout_handler) 

60 return logger 

61 

62 

63def bytes_to_human(num_bytes: int | float) -> str: 

64 """Convert a byte count to a human-readable string.""" 

65 if num_bytes < 0: 

66 raise ValueError("num_bytes must be non-negative") 

67 units = ["B", "KB", "MB", "GB", "TB", "PB", "EB"] 

68 value = float(num_bytes) 

69 for unit in units: 

70 if value < 1024.0 or unit == units[-1]: 

71 if unit == "B": 

72 return f"{int(value)} {unit}" 

73 return f"{value:.2f} {unit}" 

74 value /= 1024.0 

75 return f"{value:.2f} EB"