Source code for eradiate.test_tools.types

from __future__ import annotations

import mitsuba as mi
import pytest

from ..contexts import KernelContext
from ..kernel import MitsubaObjectWrapper, mi_traverse
from ..scenes.core import CompositeSceneElement, NodeSceneElement, Scene, traverse




def check_scene_element(
    instance: NodeSceneElement | CompositeSceneElement,
    mi_cls=None,
    ctx: KernelContext = None,
    drop_parameters: bool = True,
) -> MitsubaObjectWrapper:
    """
    Perform kernel dictionary checks on a scene element.

    This function checks if the scene element can produce a valid kernel
    dictionary template, as well as an appropriate parameter table.

    The returned Mitsuba object and parameter table can be used to perform
    additional checks.

    Parameters
    ----------
    instance : :class:`.NodeSceneElement` or :class:`.CompositeSceneElement`
        Node scene element to check.

    mi_cls : :class:`mitsuba.Object`
        Mitsuba class the node scene element expands to. Must be set if
        `instance` is a :class:`.NodeSceneElement`; ignored otherwise.

    ctx : .KernelContext, optional
        If provided, the kernel dictionary context to use. Otherwise, a default
        context is created.

    drop_parameters : bool, default: True
        If ``True``, the Mitsuba scene parameter table will be stripped off from
        untracked parameters.

    Returns
    -------
    mi_obj : :class:`mitsuba.Object`
        Mitsuba object the scene element was expanded to. See notes for details.

    mi_params : dict
        Parameter table of the Mitsuba objects generated by the tested scene
        element.

    Notes
    -----
    * If `instance` is a :class:`.NodeSceneElement`, the corresponding Mitsuba
      object type is checked against `mi_cls`.
    * If `instance` is a :class:`.CompositeSceneElement`, the corresponding
      Mitsuba objects are automatically encapsulated into a
      :class:`mitsuba.Scene` object.
    """
    if isinstance(instance, NodeSceneElement):
        if mi_cls is None:
            raise ValueError("Expected Mitsuba class must be set")
        kdict_template, umap_template = traverse(instance)

    elif isinstance(instance, CompositeSceneElement):
        mi_cls = mi.Scene
        kdict_template, umap_template = traverse(Scene(objects={"composite": instance}))

    else:
        raise RuntimeError(f"Cannot test type '{instance.__class__}'")

    # Check if the template can be instantiated
    ctx = KernelContext() if ctx is None else ctx
    kernel_dict = kdict_template.render(ctx)

    try:
        mi_obj = mi.load_dict(kernel_dict)
    except RuntimeError as e:
        pytest.fail(
            reason=f"could not load scene dictionary, got RuntimeError: {e}\n"
            f"{kernel_dict = }"
        )

    assert isinstance(mi_obj, mi_cls)

    # Collect Mitsuba parameters, resolve update map parameter paths
    mi_wrapper = mi_traverse(mi_obj, umap_template)
    mi_params = mi_wrapper.parameters

    # Check that parameters can all be set
    umap = umap_template.render(ctx)
    for key, value in umap.items():
        try:
            mi_params[key] = value
        except KeyError as e:
            pytest.fail(
                reason=f"could not set parameter, got KeyError: {e}\n{mi_params = }"
            )

    mi_params.update()

    # Drop untracked parameters (this will detect param lookup failures)
    if drop_parameters:
        mi_wrapper.drop_parameters()

    return mi_wrapper