imodfile 0.2.2

A pure-Rust IMOD model file decoder/encoder — binary & ASCII, with Python bindings
Documentation
"""Tests for imodfile Python bindings.

Requires: pytest, numpy
Run with: python -m pytest tests/test_python.py -v
"""

import importlib.util
import os
import sys

import numpy as np

# ── locate the compiled .so ──────────────────────────────────────────────

_so_paths = [
    # typical maturin develop location
    "target/debug/libimodfile.so",
    # typical maturin build location
    "target/release/libimodfile.so",
]

_mod = None
for _p in _so_paths:
    if os.path.exists(_p):
        _spec = importlib.util.spec_from_file_location("imodfile", _p)
        if _spec:
            _mod = importlib.util.module_from_spec(_spec)
            _spec.loader.exec_module(_mod)
            break

if _mod is None:
    raise ImportError(
        "imodfile not found. Build with: maturin build --features python && pip install target/wheels/imodfile-*.whl"
    )
imodfile = _mod


# ── helpers ──────────────────────────────────────────────────────────────

_MODEL_PATH = os.path.join(os.path.dirname(__file__), "test_001.mod")


def _load_test_model():
    assert os.path.exists(_MODEL_PATH), f"test file not found: {_MODEL_PATH}"
    return imodfile.load(_MODEL_PATH)


# ── tests ────────────────────────────────────────────────────────────────


def test_load_model():
    """Load a real IMOD file and check top-level metadata."""
    model = _load_test_model()
    assert model.name == "IMOD-NewModel"
    assert model.image_size == (1022, 1440, 750)
    assert abs(model.pixel_size - 0.868) < 0.001
    assert len(model) == 1  # __len__ → number of objects
    assert len(model.objects) == 1


def test_read_objects():
    """Read objects and their contours."""
    model = _load_test_model()
    obj = model.objects[0]

    assert obj.name == ""
    assert obj.color == (0.0, 1.0, 0.0)  # green
    assert obj.flags == 0x18040000
    assert len(obj) == 7  # 7 contours

    # Check all contours
    for i, cont in enumerate(obj.contours):
        pts = cont.points
        assert pts.shape == (2, 3), f"contour {i}: expected (2,3), got {pts.shape}"
        assert pts.dtype == np.float32
        assert cont.psize == 2
        assert cont.surface == 0
        assert cont.time == 1
        assert cont.flags == 0x10


def test_contour_points_values():
    """Verify specific point coordinates."""
    model = _load_test_model()
    cont = model.objects[0].contours[0]
    pts = cont.points

    assert abs(pts[0, 0] - 632.3413) < 0.01
    assert abs(pts[0, 1] - 442.1411) < 0.01
    assert abs(pts[0, 2] - 404.1057) < 0.01
    assert abs(pts[1, 0] - 649.4098) < 0.01
    assert abs(pts[1, 1] - 558.5669) < 0.01
    assert abs(pts[1, 2] - 420.9811) < 0.01


def test_all_points():
    """model.points() returns all points concatenated."""
    model = _load_test_model()
    all_pts = model.points()
    assert all_pts.shape == (14, 3)  # 7 contours × 2 points
    assert all_pts.dtype == np.float32

    # First point matches first contour
    first_cont = model.objects[0].contours[0].points
    np.testing.assert_array_almost_equal(all_pts[0], first_cont[0])


def test_create_object():
    """Create a new object, modify it, read back."""
    model = _load_test_model()
    n_before = len(model)

    obj = model.add_object()
    obj.name = "test"
    obj.color = (1.0, 0.0, 0.0)  # red
    obj.flags = 0x100  # off

    assert obj.name == "test"
    assert obj.color == (1.0, 0.0, 0.0)
    assert obj.flags == 0x100
    assert len(model) == n_before + 1


def test_create_contour():
    """Create contour with numpy points."""
    model = _load_test_model()
    obj = model.add_object()

    cont = obj.add_contour()
    cont.points = np.array([[0, 0, 0], [1, 0, 0], [1, 1, 0]], dtype=np.float32)
    cont.surface = 1
    cont.time = 2
    cont.flags = 0x10

    assert cont.psize == 3
    assert cont.surface == 1
    assert cont.time == 2
    assert cont.flags == 0x10
    np.testing.assert_array_almost_equal(cont.points[0], [0, 0, 0])
    np.testing.assert_array_almost_equal(cont.points[2], [1, 1, 0])


def test_create_mesh():
    """Create mesh with vertices and indices."""
    model = _load_test_model()
    obj = model.add_object()

    mesh = obj.add_mesh()
    mesh.vertices = np.array(
        [[0, 0, 0], [1, 0, 0], [1, 1, 0], [0, 1, 0]], dtype=np.float32
    )
    mesh.indices = [-23, 0, 0, 1, 0, 2, 0, -22]

    assert mesh.vertices.shape == (4, 3)
    assert len(mesh.indices) == 8
    assert mesh.indices[0] == -23
    np.testing.assert_array_almost_equal(mesh.vertices[0], [0, 0, 0])


def test_modify_contour_points():
    """Modify an existing contour's points."""
    model = _load_test_model()
    cont = model.objects[0].contours[0]
    old_shape = cont.points.shape

    cont.points = np.array([[100, 200, 300], [400, 500, 600]], dtype=np.float32)
    assert cont.psize == 2
    np.testing.assert_array_almost_equal(cont.points[0], [100, 200, 300])
    np.testing.assert_array_almost_equal(cont.points[1], [400, 500, 600])


def test_remove_contour():
    """Remove a contour by index."""
    model = _load_test_model()
    obj = model.objects[0]
    n_before = len(obj)

    obj.remove_contour(0)
    assert len(obj) == n_before - 1

    # First contour is now the old second one
    assert abs(obj.contours[0].points[0, 0] - 431.695) < 0.01


def test_remove_object():
    """Remove an object by index."""
    model = _load_test_model()
    model.add_object()
    n_before = len(model)

    model.remove_object(1)
    assert len(model) == n_before - 1


def test_save_and_reload(tmp_path):
    """Save model to binary file and reload."""
    model = _load_test_model()

    # Add content to verify it persists
    obj = model.add_object()
    obj.name = "persist_test"
    cont = obj.add_contour()
    cont.points = np.array([[42, 43, 44]], dtype=np.float32)

    path = str(tmp_path / "test_save.mod")
    model.save(path)

    model2 = imodfile.load(path)
    assert len(model2) == 2
    assert model2.objects[1].name == "persist_test"
    assert model2.objects[1].contours[0].psize == 1
    np.testing.assert_array_almost_equal(
        model2.objects[1].contours[0].points[0], [42, 43, 44]
    )


def test_save_ascii(tmp_path):
    """Save to ASCII (.txt) and re-load."""
    model = _load_test_model()
    path = str(tmp_path / "test_ascii.txt")
    model.save(path)

    with open(path) as f:
        content = f.read()
    assert content.startswith("# imod ascii file version 2.0")
    assert "imod 1" in content
    assert "contour 0 0 2" in content

    # Re-load ASCII file
    model2 = imodfile.load(path)
    assert len(model2) == 1
    assert len(model2.objects[0]) == 7


def test_save_points(tmp_path):
    """save_points writes correct text format."""
    model = _load_test_model()
    path = str(tmp_path / "points.txt")
    model.save_points(path)

    with open(path) as f:
        lines = f.read().strip().split("\n")

    # 7 contours × 2 points = 14 coordinate lines + 6 blank-line separators
    coord_lines = [l for l in lines if l.strip()]
    assert len(coord_lines) == 14, f"expected 14 coordinate lines, got {len(coord_lines)}"


def test_repr():
    """__repr__ returns meaningful strings."""
    model = _load_test_model()
    assert "Model" in repr(model)

    obj = model.objects[0]
    assert "Object" in repr(obj)

    cont = obj.contours[0]
    assert "Contour" in repr(cont)

    mesh = obj.add_mesh()
    assert "Mesh" in repr(mesh)


def test_invalid_points_shape():
    """Setting points to wrong shape raises ValueError."""
    model = _load_test_model()
    cont = model.objects[0].add_contour()

    import pytest

    with pytest.raises((ValueError, TypeError)):
        cont.points = np.array([1, 2, 3], dtype=np.float32)  # 1D, not (N,3)