diff --git a/src/dopt_basics/system.py b/src/dopt_basics/system.py new file mode 100644 index 0000000..41f89e1 --- /dev/null +++ b/src/dopt_basics/system.py @@ -0,0 +1,38 @@ +from __future__ import annotations + +import sys +from dataclasses import fields, is_dataclass +from typing import Any, Literal + + +def obj_size( + obj: Any, + seen: set[int] | None = None, + unit: Literal["b", "kb", "mb", "gb"] = "b", +) -> float: + size = sys.getsizeof(obj) + if seen is None: + seen = set() + + obj_id = id(obj) + if obj_id in seen: + return 0 + seen.add(obj_id) + + if is_dataclass(obj): + for f in fields(obj): + size += obj_size(getattr(obj, f.name), seen) + elif isinstance(obj, dict): + size += sum(obj_size(v, seen) + obj_size(k, seen) for k, v in obj.items()) + elif isinstance(obj, (list, tuple, set)): + size += sum(obj_size(i, seen) for i in obj) + + match unit: + case "kb": + size /= 1024 + case "mb": + size /= 1024 * 1024 + case "gb": + size /= 1024 * 1024 * 1024 + + return size diff --git a/tests/test_system.py b/tests/test_system.py new file mode 100644 index 0000000..de66cb5 --- /dev/null +++ b/tests/test_system.py @@ -0,0 +1,57 @@ +import dataclasses as dc +import sys + +import pytest + +from dopt_basics import system + + +@dc.dataclass() +class _TData: + dic: dict[str, str] + dic2: dict[str, dict[str, int]] + lst: list[int] + tup: tuple[str, ...] + string: str + id_: int + + +@pytest.fixture(scope="module") +def TData() -> _TData: + return _TData( + {"test": "test", "test2": "test"}, + {"test": {"prop1": 3, "prop2": 500}}, + [1, 2, 3, 4, 5, 6, 7, 8, 9], + ("test", "test", "test", "test", "test", "test", "test"), + "This is one test string", + 1234, + ) + + +@pytest.fixture(scope="module") +def TData_real_size(TData) -> int: + return 1435 + + +def test_obj_size(TData, TData_real_size): + ret = system.obj_size(TData) + + assert ret == TData_real_size + + +def test_obj_size_kb(TData, TData_real_size): + ret = system.obj_size(TData, unit="kb") + + assert ret == pytest.approx(TData_real_size / 1024) + + +def test_obj_size_mb(TData, TData_real_size): + ret = system.obj_size(TData, unit="mb") + + assert ret == pytest.approx(TData_real_size / 1024**2) + + +def test_obj_size_gb(TData, TData_real_size): + ret = system.obj_size(TData, unit="gb") + + assert ret == pytest.approx(TData_real_size / 1024**3)