Search code examples
pythoncompositionmypypython-3.9enforcement

Enforcing composition in Python with mypy


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.


Solution

  • 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