Feature Request: type safe function composition

[mattdornfeld]

There are many third party libraries, which aim to provide functional programming tools. Examples of this are fn.py and and funcy. They both provide a compose function that allows you to chain together an arbitrary number of functions. However none of the implementations are compatible with MyPy. Has there been any thought on the MyPy team about supporting such a feature, possibly in the typing submodule? I was able to get such a feature working with composing two functions, similar to this blog post https://www.fabianism.us/blog/2015/11/18/function-composition-with-types.html. But this doesn't scale to an arbitrary number of functions.

[JukkaL]

I'm not quite sure about the exact thing you are trying to do. Would you like to a function compose(f1, f2, ...) that can compose an arbitrary number of functions? If yes, that is too complicated for mypy to understand. If you really care about this, you may be able to support this by writing a mypy plugin.

[mattdornfeld]

Hi @JukkaL, yes that’s exactly what I’m trying to do. I tried writing an extension and I came to the same conclusion. It’s too complicated for mypy to understand. I was hoping the mypy devs maybe knew something I didn’t :)

[jaycunningham-8451]

Is it possible to statically type a generic compose function that only takes two arguments? E.g.

A = TypeVar('A')
X = TypeVar('X')
Y = TypeVar('Y')
Z = TypeVar('Z')

def compose(f: Callable[[X], Y], g: Callable[[Y], Z]) -> Callable[[X], Z]:
    def _compose(x: X) -> Z:
        return g(f(x))
    return _compose

def f(a: A) -> A:
    return a

def g(a: int) -> int:
    return a

compose(f, g)(3)  # error: Cannot infer type of argument 2 of "compose"

[MartinSeeler]

It is possible, I think you mixed up some types in your compose function. Since f is after g, f has to be an arrow from Y => Z, e.g.

A = TypeVar("A")
B = TypeVar("B")
X = TypeVar("X")
Y = TypeVar("Y")
Z = TypeVar("Z")

def id(x: A) -> A:
    return x

def compose(f: Callable[[Y], Z], g: Callable[[X], Y]) -> Callable[[X], Z]:
    def _compose(x: X) -> Z:
        return f(g(x))
    return _compose

def to_str(x: B) -> str:
    return str(x)

res: str = id("foo")
foo: str = to_str(id(4))
my_fn: Callable[[B], str] = compose(id, to_str)
my_fn2: Callable[[A], str] = compose(to_str, id)
my_res: str = my_fn(4)

[janluke]

I expected the following to work (composition of functions T -> T):

from typing import Callable, TypeVar

T = TypeVar('T')

def compose(*functions: Callable[[T], T]) -> Callable[[T], T]:
    def composition(x):
        y = x
        for f in functions:
            y = f(y)
        return y
    return composition

def identity(x: T) -> T:
    return x

f = compose(identity, identity)
reveal_type(f)  # Revealed type is "def (T`-1) -> T`-1"
f(1)  # error: Argument 1 has incompatible type "int"; expected "T"

Is this just a limitation or a bug?

EDIT: note that using a different TypeVar for the return type works:

from typing import Callable, TypeVar

A = TypeVar('A')
B = TypeVar('B')

def compose(*functions: Callable[[A], A]) -> Callable[[B], B]:
    def composition(x):
        y = x
        for f in functions:
            y = f(y)
        return y
    return composition

def identity(x: A) -> A:
    return x

f = compose(identity, identity)
reveal_type(f)  # Revealed type is "def (T`-1) -> T`-1"
res = f(1)
reveal_type(res)  # Revealed type is "builtins.int*"

[cardoso-neto]

Relevant discussion regarding compose as an operator (or maybe a classmethod) and variadic functions is also being discussed on pytoolz/toolz#523. Would love to hear more opinions.

[mentalisttraceur]

Couple notes that it would be nice to keep in mind when implementing type-checking support for function composition:

1. It's sometimes useful to compose async functions or a mix of sync and async functions (so consider what typing would look like for the acompose provided by my compose library).

2. One thing I realized while responding to mentalisttraceur/compose#1 is that beause operators are ergonomics/convenience for writing code, you probably want composition-as-operator to automatically notice if at least one operand is async and have the result be an async callable as well in that case (this is a more involved case for static typing than the acompose case in the first point, because acompose(f, g) is always an async callable but f + g could be either a sync or async callable - but on the other hand it's simpler in that the operator overload only has two parameters to deal with, and can probably be handled outside of MyPy the same way that they're discussed doing in the issue linked in the last comment).

[hauntsaninja]

This example type checks with modern mypy: #8449 (comment)

[mentalisttraceur]

@hauntsaninja this issue is about

"a compose function that allows you to chain together an arbitrary number of functions"

The example you linked to only works for a compose that's hard-coded to take two functions, which is inadequate for basically every compose implementation in the wild.

That technique can be adjusted to other hard-coded amounts of functions, and can be scaled with typing.overload, as I do in compose-stubs, but that's

1. awful to read/maintain, and
2. perceptibly slows type checkers.

[mentalisttraceur]

Also, type checking for compose implementations which support async is still broken.

Consider this type hint for my acompose with two arguments:

P = ParamSpec('P')
R1 = TypeVar('R1')
R2 = TypeVar('R2')


def acompose(f2: Callable[[R1], Union[R2, Awaitable[R2]]], f1: Callable[P, Union[R1, Awaitable[R1]]], /) -> Callable[P, Awaitable[R2]]:
    ...

As of this morning, a fresh install of MyPy on Python 3.11 errors out (something about expecting a return type of Awaitable[Never]) when my acompose is called with an async function. (Worked fine in Pyre and Pyright last I checked.)

[antonagestam]

I wonder though, is this possible to express without changes to introduce some new type level construct? If not, perhaps the issue should be closed anyway, with encouragement to propose adding such a construct with a PEP.

(I think this conceptually needs a type-level for loop, to be able to express each function is compatible with the next)

[antonagestam]

@mentalisttraceur I would encourage opening a separate bug report for the issue you are seeing with async functions.

[mentalisttraceur]

@antonagestam I agree that this is best solved by a new typing construct through a PEP.

I probably don't have the time and energy to pursue a PEP right now, but I'll jot down an idea for now:

Idea

1. The generic need here is a way to specify a relationship between types in a variadict tuple of types: T1 is somehow related to T2, T2 is related to T3, and so on.

2. I think something inspired by "base case and induction step" approaches of recursive functions and inductive proofs might work really well here.

3. So as an initial building block, I imagine typing.Relation[T1, T2], where T1 and T2 are type hints which contain one or more of the same typevars.

4. Then, I imagine a typing.RelatedTypes[*types_or_relations], which represents multiple types (similar to typing.TypeVarTuple), and takes one or more type hints and relations.

Example

The type hint for compose (two or more arguments and no async support for simplicity) can now be expressed like this:

P = ParamSpec('P')
R = TypeVar('R')
R1 = TypeVar('R1')
R2 = TypeVar('R2')

@override
def compose(*functions: RelatedTypes[Callable[[Any], R], Relation[Callable[[R1], R2], Callable[..., R1]], Callable[P, R1]) -> Callable[P, R]:
    ...

Details

I would love to hear thoughts from implementers of type checkers on this! Did I miss some reason why this is infeasible/horrible/inefficient to implement?

1. If I didn't miss anything, a type checker could match RelatedTypes to values (in arguments, but maybe also in lists, etc) in one left-to-right pass (no backtracking, O(number_of_values) time and O(1) space).

2. Types match exactly one value.

3. Relations match two or more values. (An override or union can cover the one/zero value cases. This avoids some problems - consider how tricky the spec would be if we had to define a sensible thing for the RelatedTyped of compose to do when called with just one argument!)

4. Relations overlap for one value on each side if there's another type or relation specified on that side within the RelatedTypes.

In other words, given a signature like def f(*args: RelatedTypes[T0, Relation[T1, T2], T3]): ...:

• a call like f(foo, qux), foo must pass for both T0 and T1, and qux must pass T2 and T3, and

• a call like f(foo, bar, qux), foo must pass for both T0 and T1, bar must pass both T1 and T2, and qux must pass T2 and T3.

[antonagestam]

@mentalisttraceur I think the post you just made here would be great to take to the typing forum, you can likely get some feedback from the community and hopefully also type checker maintainers there.

https://discuss.python.org/c/typing/32

[mentalisttraceur]

@antonagestam Cheers, done.