AnyImage/tools/validate_cutout_asset.py

from pathlib import Path
import sys

import bpy


PROJECT_ROOT = Path(__file__).parents[1]
sys.path.insert(0, str(PROJECT_ROOT / "src"))
sys.path.insert(
    0,
    str(next((PROJECT_ROOT / "src" / "anyimage" / "wheels").glob("blendjob-*.whl"))),
)

from anyimage.operators import mesh
from anyimage.operators.material import modifier_input_identifier, set_modifier_input


def evaluated_mesh(obj):
    evaluated = obj.evaluated_get(bpy.context.evaluated_depsgraph_get())
    result = evaluated.to_mesh()
    return evaluated, result


def grid_object(name):
    vertices = [
        (0.0, float(y), float(z))
        for z in range(3)
        for y in range(3)
    ]
    faces = []
    for z in range(2):
        for y in range(2):
            a = z * 3 + y
            faces.extend(((a, a + 4, a + 1), (a, a + 3, a + 4)))
    data = bpy.data.meshes.new(name)
    data.from_pydata(vertices, (), faces)
    obj = bpy.data.objects.new(name, data)
    bpy.context.scene.collection.objects.link(obj)
    return obj


def add_modifier(obj, shape, thickness):
    group = mesh._mesh_node_group(shape)
    modifier = obj.modifiers.new(group.name, "NODES")
    modifier.node_group = group
    set_modifier_input(
        modifier,
        modifier_input_identifier(group, "Cleanup"),
        0.0,
    )
    set_modifier_input(
        modifier,
        modifier_input_identifier(group, "Thickness"),
        thickness,
    )
    return modifier


def main():
    bpy.ops.wm.read_factory_settings(use_empty=True)
    groups = {
        shape: mesh._mesh_node_group(shape)
        for shape in (
            mesh.CUTOUT,
            mesh.BALLOON,
            mesh.DEPTH_BALLOON,
            mesh.DEPTH_SURFACE,
        )
    }
    assert [groups[shape].name for shape in groups] == [
        "O Image Surface",
        "O Image Balloon",
        "O Image Depth Balloon",
        "O Image Depth Surface",
    ]
    repeat_counts = {
        shape: sum(
            node.bl_idname == "GeometryNodeRepeatInput"
            for node in group.nodes
        )
        for shape, group in groups.items()
    }
    assert repeat_counts == {
        mesh.CUTOUT: 1,
        mesh.BALLOON: 0,
        mesh.DEPTH_BALLOON: 0,
        mesh.DEPTH_SURFACE: 2,
    }
    for group in groups.values():
        assert group.is_modifier
        assert all(not node.label for node in group.nodes)
        assert all(
            link.from_node.location.x <= link.to_node.location.x
            for link in group.links
        )
        group_nodes = list(group.nodes)
        assert len({tuple(node.location) for node in group_nodes}) == len(group_nodes)
        for index, node in enumerate(group_nodes):
            for other in group_nodes[index + 1:]:
                assert not (
                    abs(node.location.x - other.location.x) < 150.0
                    and abs(node.location.y - other.location.y) < 100.0
                )
        for group_input in (
            node for node in group_nodes if node.bl_idname == "NodeGroupInput"
        ):
            assert all(
                output.is_linked or output.hide
                for output in group_input.outputs
                if output.name
            )
        for item in group.interface.items_tree:
            if item.item_type == "SOCKET" and item.in_out == "INPUT":
                assert item.structure_type == "SINGLE"

    surface = grid_object("Surface Validation")
    add_modifier(surface, mesh.CUTOUT, 2.4)
    evaluated, result = evaluated_mesh(surface)
    try:
        assert len(result.vertices) == 34
        assert len(result.polygons) == 40
        side_faces = [
            polygon for polygon in result.polygons
            if abs(polygon.normal.x) < 0.1
        ]
        assert len(side_faces) == 24
    finally:
        evaluated.to_mesh_clear()

    balloon = grid_object("Balloon Validation")
    profile = balloon.data.attributes.new(mesh.BALLOON_ATTRIBUTE, "FLOAT", "POINT")
    profile.data.foreach_set(
        "value",
        (0.0, 0.0, 0.0, 0.0, 1.0, 0.0, 0.0, 0.0, 0.0),
    )
    modifier = add_modifier(balloon, mesh.BALLOON, 1.0)
    evaluated, result = evaluated_mesh(balloon)
    try:
        initial_faces = len(result.polygons)
        initial_extent = max(abs(vertex.co.x) for vertex in result.vertices)
    finally:
        evaluated.to_mesh_clear()
    set_modifier_input(
        modifier,
        modifier_input_identifier(modifier.node_group, "Thickness"),
        2.0,
    )
    balloon.update_tag(refresh={"DATA"})
    bpy.context.view_layer.update()
    evaluated, result = evaluated_mesh(balloon)
    try:
        assert len(result.polygons) == initial_faces
        extent = max(abs(vertex.co.x) for vertex in result.vertices)
        assert abs(extent - initial_extent * 2.0) < 1e-5
    finally:
        evaluated.to_mesh_clear()

    print(
        f"Validated Blender {bpy.app.version_string}: "
        f"{', '.join(group.name for group in groups.values())}"
    )


if __name__ == "__main__":
    main()