codekingpro/portable-devtools

1from typing import Any, Dict
2
3import os
4import yaml
5
6
7def set_log_file_path(
8    log_config_path: str, new_filename: str = "chroma.log"
9) -> Dict[str, Any]:
10    """This works with the standard log_config.yml file.
11    It will not work with custom log configs that may use different handlers"""
12    with open(f"{log_config_path}", "r") as file:
13        log_config = yaml.safe_load(file)
14    for handler in log_config["handlers"].values():
15        if handler.get("class") == "logging.handlers.RotatingFileHandler":
16            handler["filename"] = new_filename
17
18    return log_config
19
20
21def get_directory_size(directory: str) -> int:
22    """Get the size of a directory in bytes"""
23    total = 0
24    with os.scandir(directory) as it:
25        for entry in it:
26            if entry.is_file():
27                total += entry.stat().st_size
28            elif entry.is_dir():
29                total += get_directory_size(entry.path)
30    return total
31
32
33# https://stackoverflow.com/a/1094933
34def sizeof_fmt(num: int, suffix: str = "B") -> str:
35    n: float = float(num)
36    for unit in ("", "Ki", "Mi", "Gi", "Ti", "Pi", "Ei", "Zi"):
37        if abs(n) < 1024.0:
38            return f"{n:3.1f}{unit}{suffix}"
39        n /= 1024.0
40    return f"{n:.1f}Yi{suffix}"
41