Fix plot_allocation's hardcoded PDF output and MB-only unit scaling

src/python/omnimalloc/visualize.py

@@ -8,7 +8,6 @@
 from typing import Final
 
 from omnimalloc.common.optional import require_optional
-from omnimalloc.common.units import MB
 from omnimalloc.primitives import Allocation, BufferKind, IdType, Memory, Pool, System
 
 try:
@@ -39,6 +38,27 @@
 
 logger = logging.getLogger(__name__)
 
+_BYTE_UNITS: Final[tuple[tuple[int, str], ...]] = (
+    (1024**3, "GB"),
+    (1024**2, "MB"),
+    (1024, "KB"),
+)
+
+
+def _byte_unit(value: float) -> tuple[int, str]:
+    """Largest byte unit that keeps `value` at least 1 in that unit."""
+    for divisor, suffix in _BYTE_UNITS:
+        if value >= divisor:
+            return divisor, suffix
+    return 1, "B"
+
+
+def _format_bytes(value: float) -> str:
+    """Human-readable byte count, auto-scaled so small values don't read as 0.0 MB."""
+    divisor, suffix = _byte_unit(value)
+    return f"{value / divisor:.1f}{suffix}"
+
+
 KIND_COLOR_MAP: Final[dict[BufferKind, str]] = {
     BufferKind.WORKSPACE: "C0",
     BufferKind.CONSTANT: "C1",
@@ -155,7 +175,7 @@ def _draw_limit_lines(ax: Axes, limits: dict[str, int]) -> None:
     for name, value in limits.items():
         ax.axhline(value, color="black", linestyle="--", linewidth=1, alpha=0.8)
         ax.annotate(
-            f"{value / MB:.2f} MB\n{name}",
+            f"{_format_bytes(value)}\n{name}",
             xy=(x_max, value),
             xytext=(x_max * 1.02, value * 1.01),
             ha="left",
@@ -176,22 +196,21 @@ def _draw_limit_lines(ax: Axes, limits: dict[str, int]) -> None:
 def _set_axes_ticks(ax: Axes, y_limit: int, num_ticks: int = 8) -> None:
     """Configure axis ticks and formatters."""
     tick_size = y_limit / num_ticks
+    divisor, unit = _byte_unit(y_limit)
     ax.yaxis.set_major_locator(MultipleLocator(tick_size))
-    ax.yaxis.set_major_formatter(FuncFormatter(lambda x, _: f"{x / MB:.1f}"))
+    ax.yaxis.set_major_formatter(FuncFormatter(lambda x, _: f"{x / divisor:.1f}"))
     ax.xaxis.set_major_locator(MaxNLocator(integer=True))
     ax.xaxis.set_major_formatter(FuncFormatter(lambda x, _: f"{int(x):,}"))
     ax.grid(visible=True, alpha=0.5)
+    ax.set_ylabel(f"Memory ({unit})")
 
 
 def _set_axes_labels(
     ax: Axes, memory: Memory, memory_size: int | None, num_pools: int
 ) -> None:
-    size_str = (
-        f"{memory_size / MB:.1f}MB" if memory_size is not None else "Unknown Size"
-    )
+    size_str = _format_bytes(memory_size) if memory_size is not None else "Unknown Size"
     ax.set_title(f"{memory.id} ({size_str}, {num_pools} pools)")
     ax.set_xlabel("Time (Step)")
-    ax.set_ylabel("Memory (MB)")
 
 
 def _set_axes_limits(
@@ -288,7 +307,7 @@ def _visualize_system(
     if file_path is not None:
         file_path = Path(file_path)
         file_path.parent.mkdir(parents=True, exist_ok=True)
-        fig.savefig(file_path, bbox_inches="tight", format="pdf")
+        fig.savefig(file_path, bbox_inches="tight")
 
     plt.close(fig)
 

tests/unit/test_visualize.py

@@ -6,7 +6,13 @@
 
 import pytest
 from omnimalloc.primitives import Allocation, BufferKind, Memory, Pool, System
-from omnimalloc.visualize import HAS_MATPLOTLIB, _canonicalize, plot_allocation
+from omnimalloc.visualize import (
+    HAS_MATPLOTLIB,
+    _byte_unit,
+    _canonicalize,
+    _format_bytes,
+    plot_allocation,
+)
 
 pytestmark = pytest.mark.skipif(not HAS_MATPLOTLIB, reason="matplotlib not installed")
 
@@ -269,3 +275,37 @@ def test_visualize_with_memory_limits(artifacts_dir: Path) -> None:
     assert result == output_path
     assert output_path.exists()
     assert output_path.stat().st_size > 0
+
+
+def test_visualize_saves_png_when_path_has_png_extension(artifacts_dir: Path) -> None:
+    alloc = Allocation(id=1, size=100, start=0, end=10, offset=0)
+    pool = Pool(id=1, allocations=(alloc,), offset=0)
+
+    output_path = artifacts_dir / "test_extension.png"
+    result = plot_allocation(pool, file_path=output_path)
+    assert result == output_path
+    with output_path.open("rb") as f:
+        assert f.read(8) == b"\x89PNG\r\n\x1a\n"
+
+
+def test_visualize_still_saves_pdf_by_default(artifacts_dir: Path) -> None:
+    alloc = Allocation(id=1, size=100, start=0, end=10, offset=0)
+    pool = Pool(id=1, allocations=(alloc,), offset=0)
+
+    output_path = artifacts_dir / "test_extension.pdf"
+    result = plot_allocation(pool, file_path=output_path)
+    assert result == output_path
+    with output_path.open("rb") as f:
+        assert f.read(5) == b"%PDF-"
+
+
+def test_byte_unit_picks_largest_unit_that_keeps_value_above_one() -> None:
+    assert _byte_unit(500) == (1, "B")
+    assert _byte_unit(1024) == (1024, "KB")
+    assert _byte_unit(1024**2) == (1024**2, "MB")
+    assert _byte_unit(1024**3) == (1024**3, "GB")
+
+
+def test_format_bytes_does_not_collapse_small_values_to_zero() -> None:
+    assert _format_bytes(3000) == "2.9KB"
+    assert _format_bytes(200) == "200.0B"