Skip to content

toolkit

Modules:

Classes:

Functions:

Attributes:

critical_once module-attribute 

critical_once = partial(log_once, 'CRITICAL')

debug_once module-attribute 

debug_once = partial(log_once, 'DEBUG')

error_once module-attribute 

error_once = partial(log_once, 'ERROR')

info_once module-attribute 

info_once = partial(log_once, 'INFO')

success_once module-attribute 

success_once = partial(log_once, 'SUCCESS')

timer module-attribute 

timer = Timer

trace_once module-attribute 

trace_once = partial(log_once, 'TRACE')

warning_once module-attribute 

warning_once = partial(log_once, 'WARNING')

BaseConfig

Bases: BaseSettings

Methods:

Attributes:

Source code in src/toolkit/exp/_config.py 
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
class BaseConfig(ps.BaseSettings):
    model_config = ps.SettingsConfigDict(cli_parse_args=True, yaml_file="params.yaml")

    @classmethod
    def settings_customise_sources(
        cls,
        settings_cls: type[ps.BaseSettings],
        init_settings: ps.PydanticBaseSettingsSource,
        env_settings: ps.PydanticBaseSettingsSource,
        dotenv_settings: ps.PydanticBaseSettingsSource,
        file_secret_settings: ps.PydanticBaseSettingsSource,
    ) -> tuple[ps.PydanticBaseSettingsSource, ...]:
        return (
            init_settings,
            env_settings,
            dotenv_settings,
            file_secret_settings,
            ps.YamlConfigSettingsSource(settings_cls),
        )

model_config class-attribute instance-attribute 

model_config = SettingsConfigDict(cli_parse_args=True, yaml_file='params.yaml')

settings_customise_sources classmethod 

settings_customise_sources(settings_cls: type[BaseSettings], init_settings: PydanticBaseSettingsSource, env_settings: PydanticBaseSettingsSource, dotenv_settings: PydanticBaseSettingsSource, file_secret_settings: PydanticBaseSettingsSource) -> tuple[PydanticBaseSettingsSource, ...]
Source code in src/toolkit/exp/_config.py 
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
@classmethod
def settings_customise_sources(
    cls,
    settings_cls: type[ps.BaseSettings],
    init_settings: ps.PydanticBaseSettingsSource,
    env_settings: ps.PydanticBaseSettingsSource,
    dotenv_settings: ps.PydanticBaseSettingsSource,
    file_secret_settings: ps.PydanticBaseSettingsSource,
) -> tuple[ps.PydanticBaseSettingsSource, ...]:
    return (
        init_settings,
        env_settings,
        dotenv_settings,
        file_secret_settings,
        ps.YamlConfigSettingsSource(settings_cls),
    )

Experiment

Methods:

Attributes:

Source code in src/toolkit/exp/_exp.py 
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
class Experiment:
    _exp: comet.BaseExperiment

    def __init__(self, exp: comet.BaseExperiment | None = None) -> None:
        self._exp = exp or comet.get_running_experiment() or unittest.mock.Mock()

    @property
    def project_name(self) -> str:
        return self._exp.project_name or "general"

    @property
    def name(self) -> str:
        return self._exp.get_name()

    def get_parameter(self, name: str) -> Any:
        return self._exp.get_parameter(name)

    def log_parameter(self, name: str, value: Any) -> None:
        self._exp.log_parameter(name, value)

    def log_parameters(self, parameters: dict[str, Any]) -> None:
        self._exp.log_parameters(parameters)

    def log_asset(self, path: tp.StrPath, name: str | None = None) -> None:
        path: Path = Path(path)
        if tk.environ.get_bool("EXP_LOG_ASSET_TO_COMET", default=False):
            self._exp.log_asset(path, name)
        if name is None:
            try:
                name = str(path.absolute().relative_to(Path.cwd()))
            except ValueError:
                name = path.name
        target_path: Path = self.exp_dir / name
        target_path.parent.mkdir(parents=True, exist_ok=True)
        shutil.copy2(path, target_path)

    @property
    def exp_dir(self) -> Path:
        return self.project_dir / self.name

    @property
    def project_dir(self) -> Path:
        return Path(
            tk.environ.get_str("EXP_PROJECT_DIR")
            or Path.home() / "exp" / self.project_name
        )

    def log_other(self, key: Any, value: Any) -> None:
        self._exp.log_other(key, value)

    def log_others(self, dictionary: Mapping[str, Any]) -> None:
        self._exp.log_others(dictionary)

exp_dir property 

exp_dir: Path

name property 

name: str

project_dir property 

project_dir: Path

project_name property 

project_name: str

__init__ 

__init__(exp: BaseExperiment | None = None) -> None
Source code in src/toolkit/exp/_exp.py 
22
23
def __init__(self, exp: comet.BaseExperiment | None = None) -> None:
    self._exp = exp or comet.get_running_experiment() or unittest.mock.Mock()

get_parameter 

get_parameter(name: str) -> Any
Source code in src/toolkit/exp/_exp.py 
33
34
def get_parameter(self, name: str) -> Any:
    return self._exp.get_parameter(name)

log_asset 

log_asset(path: StrPath, name: str | None = None) -> None
Source code in src/toolkit/exp/_exp.py 
42
43
44
45
46
47
48
49
50
51
52
53
def log_asset(self, path: tp.StrPath, name: str | None = None) -> None:
    path: Path = Path(path)
    if tk.environ.get_bool("EXP_LOG_ASSET_TO_COMET", default=False):
        self._exp.log_asset(path, name)
    if name is None:
        try:
            name = str(path.absolute().relative_to(Path.cwd()))
        except ValueError:
            name = path.name
    target_path: Path = self.exp_dir / name
    target_path.parent.mkdir(parents=True, exist_ok=True)
    shutil.copy2(path, target_path)

log_other 

log_other(key: Any, value: Any) -> None
Source code in src/toolkit/exp/_exp.py 
66
67
def log_other(self, key: Any, value: Any) -> None:
    self._exp.log_other(key, value)

log_others 

log_others(dictionary: Mapping[str, Any]) -> None
Source code in src/toolkit/exp/_exp.py 
69
70
def log_others(self, dictionary: Mapping[str, Any]) -> None:
    self._exp.log_others(dictionary)

log_parameter 

log_parameter(name: str, value: Any) -> None
Source code in src/toolkit/exp/_exp.py 
36
37
def log_parameter(self, name: str, value: Any) -> None:
    self._exp.log_parameter(name, value)

log_parameters 

log_parameters(parameters: dict[str, Any]) -> None
Source code in src/toolkit/exp/_exp.py 
39
40
def log_parameters(self, parameters: dict[str, Any]) -> None:
    self._exp.log_parameters(parameters)

Timer

Methods:

Attributes:

Source code in src/toolkit/logging/_timer.py 
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
class Timer:
    _end: float
    _start: float
    depth: int | None = None
    level: int | str
    name: str | None

    def __init__(
        self,
        name: str | None = None,
        *,
        depth: int | None = None,
        level: int | str = "DEBUG",
    ) -> None:
        self.name = name
        self.depth = depth
        self.level = level

    def __call__(self, fn: Callable[_P, _T]) -> Callable[_P, _T]:
        if self.name is None:
            self.name = fn.__name__ + "()"
        if self.depth is None:
            self.depth = 3

        @functools.wraps(fn)
        def wrapped(*args: _P.args, **kwargs: _P.kwargs) -> _T:
            with self:
                result: _T = fn(*args, **kwargs)
            return result

        return wrapped

    def __enter__(self) -> None:
        self.start()

    def __exit__(
        self,
        exc_type: type[BaseException] | None,
        exc_value: BaseException | None,
        traceback: types.TracebackType | None,
    ) -> None:
        self.stop()

    @property
    def elapsed(self) -> float:
        return self._end - self._start

    def start(self) -> None:
        self._start = time.perf_counter()

    def stop(self) -> None:
        self._end = time.perf_counter()
        logger.opt(depth=self.depth or 2).log(
            self.level, "{} executed in {} sec.", self.name or "Block", self.elapsed
        )

depth class-attribute instance-attribute 

depth: int | None = depth

elapsed property 

elapsed: float

level instance-attribute 

level: int | str = level

name instance-attribute 

name: str | None = name

__call__ 

__call__(fn: Callable[_P, _T]) -> Callable[_P, _T]
Source code in src/toolkit/logging/_timer.py 
35
36
37
38
39
40
41
42
43
44
45
46
47
def __call__(self, fn: Callable[_P, _T]) -> Callable[_P, _T]:
    if self.name is None:
        self.name = fn.__name__ + "()"
    if self.depth is None:
        self.depth = 3

    @functools.wraps(fn)
    def wrapped(*args: _P.args, **kwargs: _P.kwargs) -> _T:
        with self:
            result: _T = fn(*args, **kwargs)
        return result

    return wrapped

__enter__ 

__enter__() -> None
Source code in src/toolkit/logging/_timer.py 
49
50
def __enter__(self) -> None:
    self.start()

__exit__ 

__exit__(exc_type: type[BaseException] | None, exc_value: BaseException | None, traceback: TracebackType | None) -> None
Source code in src/toolkit/logging/_timer.py 
52
53
54
55
56
57
58
def __exit__(
    self,
    exc_type: type[BaseException] | None,
    exc_value: BaseException | None,
    traceback: types.TracebackType | None,
) -> None:
    self.stop()

__init__ 

__init__(name: str | None = None, *, depth: int | None = None, level: int | str = 'DEBUG') -> None
Source code in src/toolkit/logging/_timer.py 
24
25
26
27
28
29
30
31
32
33
def __init__(
    self,
    name: str | None = None,
    *,
    depth: int | None = None,
    level: int | str = "DEBUG",
) -> None:
    self.name = name
    self.depth = depth
    self.level = level

start 

start() -> None
Source code in src/toolkit/logging/_timer.py 
64
65
def start(self) -> None:
    self._start = time.perf_counter()

stop 

stop() -> None
Source code in src/toolkit/logging/_timer.py 
67
68
69
70
71
def stop(self) -> None:
    self._end = time.perf_counter()
    logger.opt(depth=self.depth or 2).log(
        self.level, "{} executed in {} sec.", self.name or "Block", self.elapsed
    )

as_dict_of_numpy 

as_dict_of_numpy(obj: Mapping[str, ArrayLike] | None) -> dict[str, ndarray]
Source code in src/toolkit/array/numpy/_utils/_as_numpy.py 
20
21
22
23
def as_dict_of_numpy(obj: Mapping[str, tp.ArrayLike] | None) -> dict[str, np.ndarray]:
    if obj is None:
        return {}
    return {k: tk.as_numpy(v) for k, v in obj.items()}

as_numpy 

as_numpy(obj: Any) -> ndarray
Source code in src/toolkit/array/numpy/_utils/_as_numpy.py 
12
13
14
15
16
17
def as_numpy(obj: Any) -> np.ndarray:
    if tk.is_numpy(obj):
        return obj
    if tk.is_torch(obj):
        return obj.numpy(force=True)
    return np.asarray(obj)

flatten 

flatten(iterable: _T | Iterable[_T] | Iterable[Iterable[_T]] | Iterable, base_type: tuple[type, ...] = (str, bytes)) -> Iterable[_T]
Source code in src/toolkit/_iter/_sequence.py 
 9
10
11
12
13
14
15
16
17
18
19
20
21
def flatten(
    iterable: _T | Iterable[_T] | Iterable[Iterable[_T]] | Iterable,
    base_type: tuple[type, ...] = (str, bytes),
) -> Iterable[_T]:
    if not tp.is_iterable(iterable, base_type):
        yield iterable  # pyright: ignore [reportReturnType]
        return

    for item in iterable:
        if tp.is_iterable(item, base_type):
            yield from flatten(item)
        else:
            yield item

get_running_experiment 

get_running_experiment() -> Experiment
Source code in src/toolkit/exp/_exp.py 
73
74
def get_running_experiment() -> Experiment:
    return Experiment(comet.get_running_experiment())

is_array_like 

is_array_like(obj: Any) -> bool
Source code in src/toolkit/array/array_like/_utils.py 
4
5
def is_array_like(obj: Any) -> bool:
    return hasattr(obj, "__len__") and not isinstance(obj, str | bytes)

is_jax 

is_jax(obj: Any) -> TypeGuard[Array]
Source code in src/toolkit/array/jax/_utils.py 
11
12
def is_jax(obj: Any) -> TypeGuard[jax.Array]:
    return tp.is_instance_named_partial(obj, "jax.Array")

is_numpy 

is_numpy(obj: Any) -> TypeGuard[ndarray]
Source code in src/toolkit/array/numpy/_utils/_is.py 
10
11
def is_numpy(obj: Any) -> TypeGuard[np.ndarray]:
    return tp.is_instance_named_partial(obj, "numpy.ndarray")

is_subsequence 

is_subsequence(a: Sequence[Any], b: Sequence[Any]) -> bool
Source code in src/toolkit/_iter/_sequence.py 
24
25
26
27
28
29
30
31
def is_subsequence(a: Sequence[Any], b: Sequence[Any]) -> bool:
    i: int = 0
    j: int = 0
    while i < len(a) and j < len(b):
        if a[i] == b[j]:
            i += 1
        j += 1
    return i == len(a)

is_torch 

is_torch(obj: Any) -> TypeGuard[Tensor]
Source code in src/toolkit/array/torch/_utils.py 
11
12
def is_torch(obj: Any) -> TypeGuard[torch.Tensor]:
    return tp.is_instance_named_partial(obj, "torch.Tensor")

load_pydantic 

load_pydantic(fpath: StrPath, cls: type[_C], *, ext: str | None = None) -> _C
Source code in src/toolkit/serialize/_pydantic.py 
11
12
13
def load_pydantic(fpath: tp.StrPath, cls: type[_C], *, ext: str | None = None) -> _C:
    data: Any = tk.serialize.load(fpath, ext=ext)
    return cls.model_validate(data)

log_once cached 

log_once(level: int | str, message: str, *args, **kwargs) -> None
Source code in src/toolkit/logging/once.py 
6
7
8
@functools.cache
def log_once(level: int | str, message: str, *args, **kwargs) -> None:
    logger.opt(depth=1).log(level, message, *args, **kwargs)

main 

main(*, config: dict[str, Any] | None = None, exp_name: str | None = None, log_file: StrPath | None = 'exp.log', log_level: int | str = NOTSET, tags: list[str] | None = None) -> Callable[[Callable[[_C], _T]], Callable[[_C], _T]]
Source code in src/toolkit/exp/_main.py 
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
def main(
    *,
    config: dict[str, Any] | None = None,
    exp_name: str | None = None,
    log_file: tp.StrPath | None = "exp.log",
    log_level: int | str = logging.NOTSET,
    tags: list[str] | None = None,
) -> Callable[[Callable[[_C], _T]], Callable[[_C], _T]]:
    def decorator(fn: Callable[[_C], _T]) -> Callable[[_C], _T]:
        def wrapped(cfg: _C) -> _T:
            tk.logging.init(level=log_level, fpath=log_file)
            exp: tk.Experiment = tk.start(name=exp_name, tags=tags)
            exp.log_parameters(cfg.model_dump())
            exp.log_other("entrypoint", _path_relative_to_git_root())
            result: _T = fn(cfg)
            if log_file:
                exp.log_asset(log_file)
            return result

        if fn.__module__ == "__main__":
            cls: type[_C] = get_type_hints(fn)["cfg"]
            cfg: _C = cls(**(config or {}))
            wrapped(cfg)
        return wrapped

    return decorator

merge_mapping 

merge_mapping(origin: Mapping, update: Mapping) -> dict

Updates the original dict with the new data. Similar to dict.update(), but works with nested dicts.

References
  1. [ConfZ/confz/loaders/loader.py:L10-L28)(https://github.com/Zuehlke/ConfZ/blob/6c99cc2a2938e231590dceeef66749ccf2eb6b4c/confz/loaders/loader.py#L10-L28)
Source code in src/toolkit/_iter/_mapping.py 
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
def merge_mapping(origin: Mapping, update: Mapping) -> dict:
    """Updates the original dict with the new data. Similar to `dict.update()`, but works with nested dicts.

    References:
        1. [ConfZ/confz/loaders/loader.py:L10-L28)(https://github.com/Zuehlke/ConfZ/blob/6c99cc2a2938e231590dceeef66749ccf2eb6b4c/confz/loaders/loader.py#L10-L28)
    """
    original: dict = dict(origin)
    for key, value in update.items():
        if isinstance(value, Mapping) and key in original:
            if not isinstance(original[key], Mapping):
                msg: str = (
                    "Config variables contradict each other: "
                    f"Key {key:!r} is both a value and a nested dict."
                )
                raise ValueError(msg)
            original[key] = merge_mapping(original[key], value)
        else:
            original[key] = value
    return original

save_pydantic 

save_pydantic(fpath: StrPath, data: BaseModel, *, ext: str | None = None) -> None
Source code in src/toolkit/serialize/_pydantic.py 
16
17
18
19
def save_pydantic(
    fpath: tp.StrPath, data: pydantic.BaseModel, *, ext: str | None = None
) -> None:
    tk.serialize.save(fpath, data.model_dump(), ext=ext)

start 

start(*, name: str | None = None, tags: list[str] | None = None) -> Experiment
Source code in src/toolkit/exp/_start.py 
51
52
53
54
55
56
57
58
59
60
def start(*, name: str | None = None, tags: list[str] | None = None) -> tk.Experiment:
    exp_key: str = get_key()
    exp_name: str = get_name(name)
    if tk.environ.get_bool("EXP_AUTO_COMMIT", default=True):
        auto_commit(exp_name, exp_key)
    exp: comet.BaseExperiment = comet.start(
        experiment_key=exp_key,
        experiment_config=comet.ExperimentConfig(name=exp_name, tags=tags),
    )
    return tk.Experiment(exp)

strip_comments 

strip_comments(text: str, comments: str = '#') -> Iterable[str]
Source code in src/toolkit/_text/_strip_comments.py 
4
5
6
7
8
def strip_comments(text: str, comments: str = "#") -> Iterable[str]:
    for line in text.strip().splitlines():
        stripped: str = line.split(comments, 1)[0].strip()
        if stripped:
            yield stripped