Search code examples
pythonenumsoverloadingmetaclass

How do you overload python class method based on Enum type


I am trying to create a python method for a class which uses overload typing decorator. I want the method to be overloaded based on the Enum type being passed. For example,

class GetPortType(enum.Enum):
    DIFF_1: "Diff1"
    DIFF_2: "Diff2"
    SAME_1: "Same1"
    SAME_2: "Same2"

class Port:
...
    @typing.overload
    def get_port_on_type(self, type: GetPortType.DIFF_1, num: int, skip: Optional[List]):
        statement1
        statement2
        return [values]

    @typing.overload
    def get_port_on_type(self, type: GetPortType.DIFF_2, num: int, skip: Optional[List]):
        statement3
        statement4
        return [values]

    @typing.overload
    def get_port_on_type(self, type: GetPortType.SAME_1, num: int, skip: Optional[List]):
        statement1
        statement3
        return [values]

    @typing.overload
    def get_port_on_type(self, type: GetPortType.SAME_2, num: int, skip: Optional[List]):
        statement2
        statement4
        return [values]

How do I achieve this? Should I create a "def __ new __(self, value):" method for the Enum and return type for all these cases?

EDIT: I was wondering if there was something we could do with Metaclass and like value return type maybe? Which could be a property of this metaclass?

@property
def type(self):
    return type(self.value)

""" changing the type annotation """
type: GetPortType.SAME_2.type

Solution

  • It appears you've misunderstood what typing.overload does. It does not let you define different versions of your code that get run in different situations. Rather, it's used for type hinting, to indicate that several combinations of types are supported by a single implementation. None of the overload definitions of the function will ever be run, they only add better type hinting to the real version of the function that comes later.

    Here's an example given in the documentation:

    @overload
    def process(response: None) -> None:
        ...
    @overload
    def process(response: int) -> tuple[int, str]:
        ...
    @overload
    def process(response: bytes) -> str:
        ...
    def process(response):
        <actual implementation>
    

    Note that the ... elipsis literal is something you might actually put in this code, it's not standing in for something only left out for the documentation. There's no body needed in those @overload decorated functions, since they never run. The ... literal seems to have emerged among type-hinting aficionados as a preferred body for function stubs, rather than pass (which at least used to be the preferred "do nothing" body for a function).

    If you actually need a decorator that will run a different version of a function depending on an argument's type, you might be able to use functools.singledispatch. But it only dispatches on actual types, not literal values like you're wanting (e.g. specific instances of an Enum).

    The simple solution to your problem is just to write a set of if/elif/else blocks to separate out the calls once you're inside the function:

    def get_port_on_type(self, type: GetPortType, num: int, skip: Optional[List]):
        if type == GetPortType.DIFF_1:
            do_stuff_1()
        elif type == GetPortType.DIFF_2:
            do_stuff_2()
        elif type == GetPortType.SAME_1:
            do_stuff_3()
        else: # type == GetPortType.SAME_2:
            do_stuff_4()
    

    Starting in Python 3.10, you can use the new match and case statements to do essentially the same thing as the chain of if/elif/else code above, with a very slightly nicer syntax, but I'm still using 3.9 and I don't feel confident in writing an example for you that I can't test (see PEP 636 for a tutorial on the new statement types).