filter_domain.py

from typing import Protocol, TypeVar
from domain import Domain
from map_domain import map_domain
from solve import Problem

T = TypeVar("T")


def __identity(_: Problem, candidate: T) -> T:
    return candidate


class Predicate(Protocol[Problem, T]):
    def __call__(self, problem: Problem, candidate: T) -> bool:
        pass


class Filter(Protocol[Problem, T]):
    def __call__(self, domain: Domain[Problem, T]) -> Domain[Problem, T]:
        pass


def filter_domain(predicate: Predicate[Problem, T]) -> Filter[Problem, T]:
    def __filter(domain: Domain[Problem, T]) -> Domain[Problem, T]:
        def next_matching(problem: Problem, candidate: T | None) -> T | None:
            while candidate is not None and not predicate(problem, candidate):
                candidate = domain.next(problem, candidate)
            return candidate

        return map_domain(next_matching, __identity)(domain)

    return __filter