Source code for pymonad.list

# --------------------------------------------------------
# (c) Copyright 2014, 2020 by Jason DeLaat.
# Licensed under BSD 3-clause licence.
# --------------------------------------------------------
""" Implements the List monad.

The List monad is frequently used to represent calculations with
non-deterministic results, that is: functions which return more than
one (possible) result. For example, calculating how chess pieces might
move.

Example:
    >>> def knight_move(position):
    >>>    # calculates a list of every possible square a knight could move
    >>>    # to from it's current position
    >>>    return ListMonad(position_1, position_2, ..., position_N)

    >>> # A list containing every square a knight could reach after 3 moves.
    >>> three_moves = (
    >>>     List
    >>>     .insert(initial_position) # However positions are defined.
    >>>     .then(knight_move)
    >>>     .then(knight_move)
    >>>     .then(knight_move)
    >>> )
"""
from typing import (
    Any,
    Callable,
    Generic,
    List,
    TypeVar,
    Union,
)  # pylint: disable=unused-import

import pymonad.monad

S = TypeVar("S")  # pylint: disable=invalid-name
T = TypeVar("T")  # pylint: disable=invalid-name


class _List(pymonad.monad.Monad, Generic[T]):
    @classmethod
    def insert(cls, value: T) -> "_List[T]":
        return ListMonad(value)

    def amap(self: "_List[Callable[[S], T]]", monad_value: "_List[S]") -> "_List[T]":
        result = []
        for function in self:
            for value in monad_value:
                result.append(function(value))
        return ListMonad(*result)

    def bind(
        self: "_List[S]", kleisli_function: Callable[[S], "_List[T]"]
    ) -> "_List[T]":
        return self.map(kleisli_function).join()

    def join(self: "_List[_List[T]]") -> "_List[T]":
        """ Flattens a nested ListMonad instance one level. """
        return ListMonad(*[element for lists in self for element in lists])

    def map(self: "_List[S]", function: Callable[[S], T]) -> "_List[T]":
        return ListMonad(*[function(x) for x in self])

    def then(
        self: "_List[S]", function: Union[Callable[[S], T], Callable[[S], "_List[T]"]]
    ) -> "_List[T]":
        try:
            return self.bind(function)
        except TypeError:
            return self.map(function)

    def __eq__(self, other):
        return self.value == other.value

    def __getitem__(self, index):
        try:
            return ListMonad(*self.value.__getitem__(index))
        except TypeError:
            return self.value[index]

    def __iter__(self):
        return iter(self.value)

    def __len__(self):
        return len(self.value)

    def __repr__(self):
        return str(self.value)


def ListMonad(*elements: List[T]) -> _List[T]:  # pylint: disable=invalid-name
    """ Creates an instance of the List monad.

    Args:
        elements: any number of elements to be inserted into the list

    Returns:
        An instance of the List monad.
    """

    return _List(list(elements), None)


ListMonad.insert = _List.insert
ListMonad.apply = _List.apply