Ensure two fields in a Python class are the same subclass of a given base class using mypy

user_6611672 0
Viewed 153

Let's say I have the following dataclasses:

@dataclass
class Product:
    color: str

@dataclass
class Wrench(Product):
    pass

@dataclass
class Hammer(Product):
    pass

I'm trying to create a new dataclass called Order with two fields, both fields must have the same subclass of Product. I could create the Order class like so:

@dataclass
class Order:
    primary_product: Product
    secondary_product: Product

But this doesn't validate the same Product subclass condition I stated above:

product1 = Wrench(color="Yellow")
product2 = Hammer(color="Black")

order = Order(primary_product=product1, secondary_product=product2)  # NO ERROR

The following implementation of Order with generics gets me some of the way there:

from typing import Generic, TypeVar

T = TypeVar("T")

@dataclass
class Order(Generic[T]):
    primary_product: T
    secondary_product: T

product1 = Wrench(color="Yellow")
product2 = Wrench(color="White")
product3 = Hammer(color="Black")

order1 = Order[Wrench](primary_product=product1, secondary_product=product2)
order2 = Order[Wrench](primary_product=product1, secondary_product=product3)  # error: Argument "secondary_product" to "Order" has incompatible type "Hammer"; expected "Wrench"

It's annoying to pass the target Product subclass (Wrench in the case above) to Order on every object initialization. Moreover, this doesn't ensure both fields are products:

order1 = Order[int](primary_product=1, secondary_product=2)  # NO ERROR

Is there anyway to achieve this or am I pushing the limits of mypy and Python's type hints too far?

mypy python
1 Answers

Further down in the mypy documentation you posted there are two different ways presented. Both are not ideal.

First way: use a type bound

T = TypeVar('T', bound=Product)

Now the generic parameter can only be a Product

order1 = Order[int](primary_product=1, secondary_product=2)
# error: Value of type variable "T" of "Order" cannot be "int"

order1 = Order(primary_product=1, secondary_product=2)
# error: Value of type variable "T" of "Order" cannot be "int"

Unfortunately the generic parameter can now be infered to be exactly Product:

product1 = Wrench(color="Yellow")
product2 = Hammer(color="Black")

order = Order(primary_product=product1, secondary_product=product2)  # no error
reveal_type(order)
# Revealed type is 'Order[Product*]'

So you have to specify the generic type

Second way: value restriction

T = TypeVar('T', Hammer, Wrench)

Now even this is correctly identified as an error

product1 = Wrench(color="Yellow")
product2 = Hammer(color="Black")

order = Order(primary_product=product1, secondary_product=product2)
# error: Value of type variable "T" of "Order" cannot be "Product"

The problem with this approach is obvious: You have to type all the sub classes of Product into the TypeVar constructor.

Third way: factory function

After some experimenting i found a third way that has some odd syntax but combines the advantages of the first two ways. The idea is to force the generic parameter to the type of the first argument.

T = TypeVar('T', bound=Product)

def make_order(primary: P) -> Callable[[P], Order[P]]:
    def inner(secondary: P) -> Order[P]:
        return Order(primary, secondary)
    return inner

make_order(1)(2)
# error: Value of type variable "T" of "make_order" cannot be "int"

product1 = Wrench(color="Yellow")
product2 = Hammer(color="Black")

make_order(product1)(product2)
# Argument 1 has incompatible type "Hammer"; expected "Wrench"

order = make_order(product1)(product1)
reveal_type(order)
# Revealed type is 'Order[Wrench*]'

The downsides are:

  • strange syntax
  • misleading error message ("Argument 1", because it refers to the second call)

To prevent users from directly instantiating Order you could rename the class to _Order.

user_47597260
Related