Enforcing composition in Python with mypy

296 Views Asked by At

Say I have some code where I expect a class possessing some characteristics like attributes or methods, but I don't want to use inheritance to enforce this, preferring composition instead.

How would I go about enforcing that the class/object received as argument actually presents these characteristics?

Example with inheritance:

from abc import ABC, abstractmethod

class FooBase(ABC):
    @abstractmethod
    def do_stuff(self):
        pass

class Foo(FooBase):
    def do_stuff(self):
        return 'hello'

def some_function(obj: FooBase):
    print(obj.do_stuff())

In this above example, some_function is enforcing (statically through MyPy) that the obj parameter is a subclass of FooBase.

Instead, I'd like that some_function enforces (statically, via MyPy) that obj has a do_stuff method, without it necessarily being a subclass of FooBase.

1

There are 1 best solutions below

14
Samwise On BEST ANSWER

Use typing.Protocol:

from typing import Protocol

class Foo(Protocol):
    def do_stuff(self) -> str: ...

class Bar:
    def do_stuff(self) -> str:
        return "hello"

def some_function(obj: Foo):
    print(obj.do_stuff())


some_function(Bar())  # ok

The difference between a Protocol and an ABC is that you can implement a protocol without explicitly inheriting from it, so Bar is considered a Foo by mypy even though it's not a subclass of Foo.

If you pass a type that doesn't implement the protocol, you get an error, e.g.:

class Baz:
    def do_stuff(self) -> int:
        return 42

some_function(Baz())  
# error: Argument 1 to "some_function" has incompatible type "Baz"; expected "Foo"
# note: Following member(s) of "Baz" have conflicts:
# note:     Expected:
# note:         def do_stuff(self) -> str
# note:     Got:
# note:         def do_stuff(self) -> int