Function with multiple generic parameters of same type

[sh0ck-wave]

I want to create a function with two parameters (P1, P2) with the following constraints

1. P1 and P2 are bound to the same base type (B1)
2. Type[P1] == Type[P2] == ReturnType

An example to demonstrate what I mean :

from typing import Dict, Type, TypeVar

K = TypeVar("K", bound="Food")


class Food:
    def __init__(self, quantity: int) -> None:
        self._quantity = quantity
        super().__init__()

    def quantity(self):
        return self._quantity

    def mix(self: K, d: K) -> K:
        return type(self)(self.quantity() + d.quantity())


class Rice(Food):
    ...


class Sandwich(Food):
    ...


def food_mixer(f1: K, f2: K) -> K:
    return f1.mix(f2)


mixed_food1 = Rice(10).mix(Sandwich(5))  # this line shows an error as expected
mixed_food2 = food_mixer(Rice(10), Sandwich(5))  # this line does not show an error

The reason the second line does not throw an error seems to be that the type of K@food_mixer is assumed to be Union[Rice, Sandwich]
Is there a way to constrain the food_mixer function to disallow Union types while still allowing any subtype of Food to be passed in ? So that

rice + rice = allowed
sandwich + sandwich = allowed
curry + curry = allowed
etc.

but
rice + sandwich = disallowed
sandwich + curry = disallowed
etc.

I was hoping for a solution which does not require creating a constraint like V = TypeVar("V", Sandwich, Rice) since that would require modification every time a new subtype is created.

[erictraut]

No, you can't accomplish this with bound TypeVars. The constraint solver will find always find a valid solution using either a union (Sandwich | Rice) or a join (Food). Both solutions are valid.

As you pointed out, you can use a constrained TypeVar to accomplish this, but that requires each subtype to be listed in the TypeVar definition.

Here's another variant that does work with a bound TypeVar, but it requires another class.

class FoodMixer(Generic[K]):
    def __init__(self, f1: K):
        self._f1 = f1

    def mix(self, f2: K) -> K:
        return self._f1.mix(f2)

FoodMixer(Rice(10)).mix(Sandwich(5)) # Type error

And here's another variant, but it requires you to explicitly specialize the class when you use it:

class FoodMixer(Generic[K]):
    @staticmethod
    def mix(f1: K, f2: K) -> K:
        return f1.mix(f2)

FoodMixer[Rice].mix(Rice(10), Sandwich(4)) # Type error

# A type alias makes this a little nicer.
RiceMixer = FoodMixer[Rice]
RiceMixer.mix(Rice(10), Sandwich(4)) # Type error

[sh0ck-wave]

For anyone looking at this discussion in the future, there is a more hacky solution that depends on circular forward references here in case its not possible to wrap the function in a class.

[sh0ck-wave]

@erictraut Thank you for your alternatives. I tried using the Invariant property of TypeVar to force the constraint checker to disregard unions.
I am still new to python typing so would like some feedback on if this looks okay or if its a hacky workaround which might break in the future. It would be nice if there was an easier way to describe this kind of "Invariant" relationship between multiple generic parameters of a function.

from abc import abstractmethod
from typing import Generic, TypeVar

K = TypeVar("K", bound="Food")


class AnyFood(Generic[K]):
    @abstractmethod
    def unwrap(self) -> K:
        pass


class Food:
    def __init__(self, quantity: int) -> None:
        self._quantity = quantity
        super().__init__()

    def quantity(self):
        return self._quantity


class Rice(Food, AnyFood["Rice"]):
    def unwrap(self):
        return self


class Sandwich(Food, AnyFood["Sandwich"]):
    def unwrap(self):
        return self

# does not inherit from Food
class FakeFood(AnyFood["FakeFood"]):
    def unwrap(self):
        return self


def food_mixer(w1: AnyFood[K], w2: AnyFood[K]) -> K:
    f1 = w1.unwrap()
    f2 = w2.unwrap()
    return type(f1)(f1.quantity() + f2.quantity())


mixed_food1 = food_mixer(Rice(10), Rice(10))  # this line does not show an error
mixed_food2 = food_mixer(Sandwich(10), Sandwich(10))  # this line does not show an error
mixed_food3 = food_mixer(Rice(10), Sandwich(10))  # this line does show an error as expected!!
mixed_food4 = food_mixer(FakeFood(10), FakeFood(10))  # this line does show an error as expected!!

It is giving the expected outcomes in both pyright and mypy. And pylance is able to derive the type of mixed_food1 and mixed_food2 to be Rice and Sandwich respectively.

[erictraut]

That's clever, but perhaps too clever. You're creating a circular dependency in the class hierarchy (e.g. Rice derives from a class that is parameterized by Rice). Both mypy and pyright can handle this type of circularity, but other type checkers may not. This may be OK for internal code, but I'd recommend against using this trick in a public interface for a library.

This approach also requires a lot of boilerplate that doesn't scale well when adding new subclasses. A constrained TypeVar would probably be easier to maintain. Constrained TypeVars were designed to handle cases like this.

[hauntsaninja]

(Fwiw, I'm optimistic that most type checkers have some support for circular definitions, largely because str is a Sequence[str])

[hmc-cs-mdrissi]

Some of the unwrap boilerplate you can eliminate by doing,

class AnyFood(Generic[K]):
    @final
    def unwrap(self) -> K:
        return self  # type: ignore

I still prefer non cyclic version mostly because needing to call unwrap in food_mixer makes code more confusing.

[sh0ck-wave]

Thanks for the feedback. I agree its probably too much of a workaround for any public api. Just out of curiosity I wanted to see how much I could reduce the boiler plate so I took @hmc-cs-mdrissi suggestion and also moved the unwrapping code to a decorator.
Invariant code:

from typing import Callable, Generic, TypeVar, final

U = TypeVar("U")


class InvariantParam(Generic[U]):
    @final
    def unwrap(self) -> U:
        return self  # type: ignore


def invariant_decorator(
    func: Callable[[U, U], U]
) -> Callable[[InvariantParam[U], InvariantParam[U]], U]:
    def constrained_func(w1: InvariantParam[U], w2: InvariantParam[U]):
        return func(w1.unwrap(), w2.unwrap())
    return constrained_func

Rest of the code:

class Food:
    def __init__(self, quantity: int) -> None:
        self._quantity = quantity
        super().__init__()

    def quantity(self):
        return self._quantity


class Rice(Food, InvariantParam["Rice"]):
    ...


class Sandwich(Food, InvariantParam["Sandwich"]):
    ...


# does not inherit from Food
class FakeFood(InvariantParam["FakeFood"]):
    def __init__(self, quantity: int) -> None:
        super().__init__()


K = TypeVar("K", bound="Food")


@invariant_decorator
def food_mixer(f1: K, f2: K) -> K:
    return type(f1)(f1.quantity() + f2.quantity())


mixed_food1 = food_mixer(Rice(10), Rice(10))  # this line does not show an error
mixed_food2 = food_mixer(Sandwich(10), Sandwich(10))  # this line does not show an error
mixed_food3 = food_mixer(Rice(10), Sandwich(10))  # this line does show an error!!
mixed_food4 = food_mixer(FakeFood(10), FakeFood(10))  # this line does show an error!!

[saulshanabrook]

FYI I have opened a discussion on the typing forum to see if we could add a solution to this issue.