pythonooppolymorphism

Python function change behaviour depending on type of argument


I'm playing around with OOP in Python, and I have something like this:

class Person:
    def __init__(self,name, age):
        self.name = name
        self.age = age
        self.hobbies = []

    def __add__(self,hobby):
        self.hobbies.append(hobby)
        return self 

    def get_hobbies(self):
        print(self.hobbies)

tom = Person("Tom", 25) 
tom = tom + "Skiing"

tom.get_hobbies()
##["Skiing"]

Is it possible to do something to the __add__ function so that if the hobby parameter is a list, it adds the elements to the hobbies list?

tom = tom + ["cooking", "swimming"]
tom.get_hobbies()
##["Skiing","cooking","swimming"]
##instead of ["Skiing", ["cooking","swimming"]]

Solution

  • I would suggest you to use a bit of typing, to manage if the dunder method is going to accept a list a string, or raise an error if you try to feed it with something else.

    from typing import Union
    
    class Person:
        def __init__(self, name: str, age: int) -> None:
            self.name = name
            self.age = age
            self.hobbies: list[str] = []
    
        def __add__(self, hobby: Union[str, list[str]]) -> "Person":
            if isinstance(hobby, str):
                self.hobbies.append(hobby)
            elif isinstance(hobby, list) and all(isinstance(h, str) for h in hobby):
                self.hobbies.extend(hobby)
            else:
                raise TypeError(
                    f"Unsupported operand type for +: 'Person' and '{type(hobby).__name__}' "
                    "(expected str or list[str])"
                )
            return self  # allows chaining
    
        def get_hobbies(self) -> None:
            print(self.hobbies)