Input and Adapter Types

types

BatchParam dataclass

Bases: Generic[_BatchT]

Explicit shared/separate wrapper for batch inputs.

Source code in src/mifrost/encoders/types.py

@dataclass(frozen=True)
class BatchParam(Generic[_BatchT]):
    """Explicit shared/separate wrapper for batch inputs."""

    kind: Literal["shared", "separate", "none"]
    value: Any | None = None

    @classmethod
    def shared(cls, value: _BatchT) -> "BatchParam[_BatchT]":
        return cls(kind="shared", value=value)

    @classmethod
    def separate(cls, values: Iterable[_BatchT | None]) -> "BatchParam[_BatchT]":
        return cls(kind="separate", value=list(values))

    @classmethod
    def none(cls) -> "BatchParam[_BatchT]":
        return cls(kind="none", value=None)

BatchEncodingLike

Bases: Protocol

Structural type for native C++ BatchEncoding-like objects.

Concrete bindings expose this shape from C++ while Python code can type against the protocol without importing binding internals. The protocol mirrors PyG-level structural metadata (graph/node/edge counts and type lists) but intentionally does not model dynamic node storage attributes like data["atom"].x.

Source code in src/mifrost/encoders/types.py

@runtime_checkable
class BatchEncodingLike(Protocol):
    """
    Structural type for native C++ ``BatchEncoding``-like objects.

    Concrete bindings expose this shape from C++ while Python code can
    type against the protocol without importing binding internals.
    The protocol mirrors PyG-level structural metadata (graph/node/edge
    counts and type lists) but intentionally does not model dynamic node
    storage attributes like ``data["atom"].x``.
    """

    num_graphs: int
    num_nodes: int
    num_edges: int
    graph_kind: str
    node_types: Sequence[str]
    edge_types: Sequence[tuple[str, str, str]]

    def as_dict(self) -> EncodingDict: ...

    def as_pyg(self, *, as_batch: bool | None = None) -> PygDataLike: ...

    def as_hetero(self) -> "HeteroBatchEncodingView": ...

    def as_homo(self) -> "HomoBatchEncodingView": ...

    def to(self, device: Any) -> "BatchEncodingLike": ...

    def schema_fingerprint(self) -> int: ...

HeteroEncoding

Bases: BatchEncodingLike, Protocol

Refined protocol for hetero native encodings.

Source code in src/mifrost/encoders/types.py

@runtime_checkable
class HeteroEncoding(BatchEncodingLike, Protocol):
    """Refined protocol for hetero native encodings."""

    graph_kind: Literal["hetero"]

    def as_pyg(self, *, as_batch: bool | None = None) -> PygDataLike: ...

HomoEncoding

Bases: BatchEncodingLike, Protocol

Refined protocol for homo native encodings.

Source code in src/mifrost/encoders/types.py

@runtime_checkable
class HomoEncoding(BatchEncodingLike, Protocol):
    """Refined protocol for homo native encodings."""

    graph_kind: Literal["homo"]

    def as_pyg(self, *, as_batch: bool | None = None) -> PygDataLike: ...

FlatEncoding

Bases: BatchEncodingLike, Protocol

Refined protocol for flat native encodings.

Flat encodings use graph_kind == "flat" even though their tensor layout is still homo-shaped and converts through the flat PyG carrier.

Source code in src/mifrost/encoders/types.py

@runtime_checkable
class FlatEncoding(BatchEncodingLike, Protocol):
    """
    Refined protocol for flat native encodings.

    Flat encodings use `graph_kind == "flat"` even though their tensor layout is
    still homo-shaped and converts through the flat PyG carrier.
    """

    graph_kind: Literal["flat"]

    def as_pyg(self, *, as_batch: bool | None = None) -> PygDataLike: ...

register_state_adapter(state_type, adapter)

Register a custom state adapter.

The adapter must return a pymimir.advanced.search.State instance. Adapters are matched by exact concrete type.

Source code in src/mifrost/encoders/types.py

def register_state_adapter(state_type: type[object], adapter: StateAdapter) -> None:
    """
    Register a custom state adapter.

    The adapter must return a ``pymimir.advanced.search.State`` instance.
    Adapters are matched by exact concrete type.
    """
    _STATE_ADAPTERS[state_type] = adapter

unregister_state_adapter(state_type)

Remove a previously registered custom state adapter.

Source code in src/mifrost/encoders/types.py

def unregister_state_adapter(state_type: type[object]) -> None:
    """Remove a previously registered custom state adapter."""
    _STATE_ADAPTERS.pop(state_type, None)

is_state_input(value)

Return whether value is accepted as a state input.

Source code in src/mifrost/encoders/types.py

def is_state_input(value: object) -> bool:
    """Return whether ``value`` is accepted as a state input."""
    return isinstance(value, STATE_TYPES) or _resolve_state_adapter(value) is not None

register_domain_adapter(domain_type, adapter)

Register a custom domain adapter matched by exact concrete type.

Source code in src/mifrost/encoders/types.py

def register_domain_adapter(domain_type: type[object], adapter: DomainAdapter) -> None:
    """Register a custom domain adapter matched by exact concrete type."""
    _DOMAIN_ADAPTERS[domain_type] = adapter

unregister_domain_adapter(domain_type)

Remove a previously registered custom domain adapter.

Source code in src/mifrost/encoders/types.py

def unregister_domain_adapter(domain_type: type[object]) -> None:
    """Remove a previously registered custom domain adapter."""
    _DOMAIN_ADAPTERS.pop(domain_type, None)

is_domain_input(value)

Return whether value is accepted as a domain input.

Source code in src/mifrost/encoders/types.py

def is_domain_input(value: object) -> bool:
    """Return whether ``value`` is accepted as a domain input."""
    return isinstance(value, DOMAIN_TYPES) or _resolve_domain_adapter(value) is not None

register_literal_adapter(literal_type, adapter)

Register a custom goal-literal adapter matched by exact concrete type.

Source code in src/mifrost/encoders/types.py

def register_literal_adapter(
    literal_type: type[object], adapter: LiteralAdapter
) -> None:
    """Register a custom goal-literal adapter matched by exact concrete type."""
    _LITERAL_ADAPTERS[literal_type] = adapter

unregister_literal_adapter(literal_type)

Remove a previously registered custom literal adapter.

Source code in src/mifrost/encoders/types.py

def unregister_literal_adapter(literal_type: type[object]) -> None:
    """Remove a previously registered custom literal adapter."""
    _LITERAL_ADAPTERS.pop(literal_type, None)

is_goal_literal_input(value)

Return whether value is accepted as a goal literal input.

Source code in src/mifrost/encoders/types.py

def is_goal_literal_input(value: object) -> bool:
    """Return whether ``value`` is accepted as a goal literal input."""
    return (
        isinstance(value, GOAL_LITERAL_TYPES)
        or _resolve_literal_adapter(value) is not None
    )

register_action_adapter(action_type, adapter)

Register a custom ground-action adapter matched by exact concrete type.

Source code in src/mifrost/encoders/types.py

def register_action_adapter(action_type: type[object], adapter: ActionAdapter) -> None:
    """Register a custom ground-action adapter matched by exact concrete type."""
    _ACTION_ADAPTERS[action_type] = adapter

unregister_action_adapter(action_type)

Remove a previously registered custom action adapter.

Source code in src/mifrost/encoders/types.py

def unregister_action_adapter(action_type: type[object]) -> None:
    """Remove a previously registered custom action adapter."""
    _ACTION_ADAPTERS.pop(action_type, None)

is_action_input(value)

Return whether value is accepted as an action input.

Source code in src/mifrost/encoders/types.py

def is_action_input(value: object) -> bool:
    """Return whether ``value`` is accepted as an action input."""
    return (
        isinstance(value, GROUND_ACTION_TYPES)
        or _resolve_action_adapter(value) is not None
    )