object size measurement function, related to #7

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

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)