Search code examples
pythonmypypython-typing

mypy: how to best deal with random.choice


What's the best way to handle the following mypy error -

Incompatible types in assignment (expression has type "object", variable has type "Union[ClassOne, ClassTwo, ClassThree, ClassFour]")

for the below loop:

def func(self) -> List:
        thing = [x for x in self.Bar if x.condition]
        Foo: Union[ClassOne, ClassTwo, ClassThree, ClassFour]
        if len(thing) == 1:
            Foo = ClassOne(self.Bar)
        elif len(thing) == 2:
            Foo = random.choice([ClassTwo(self.Bar), ClassThree(self.Bar)])
        elif len(thing) == 3:
            Foo = ClassFour(self.Bar)
        result = Foo.mymethod()

        return result

As you can see I'm choosing an initialised class based on a condition and running a method on it.

When random.choice is used though, mypy is insisting the type is object.

Adding "object" or "Any" in the type declaration for Foo doesn't help cos then mypy complains object doesn't have the mymethod attribute.


Solution

  • Mypy infers that the type of [ClassTwo(self.Bar), ClassThree(self.Bar)] is List[object], which is reasonable since object is the last common base class for ClassTwo and ClassThree.

    But you can make its type more restrictive by setting its type explcitly:

    choices: List[Union[ClassTwo, ClassThree]] = [ClassTwo(self.Bar), ClassThree(self.Bar)]
    Foo = random.choice(choices)