Generic functions and generic classes in Python
"Generic" is a term used for any typing that might change based on the context.
If we have a function that may take either strings or ints and return the sum or concatenation of both values, without generic we would have to define two distinct functions:
def sum_numbers(v1: int, v2: int) -> int:
return v1 + v2
def concat_strs(v1: str, v2: str) -> str:
return v1 + v2
numbers = sum_numbers(10, 20)
strs = concat_strs("app", "le")
print(numbers)
print(strs)
Even though we got the proper typing, this isn't good. My code is mostly duplicated just for the sake of types - that's not really what we want here.
That's where generics can make our life easier. We reuse the code snippet while keeping the dynamic typing based on context.
This is quite popular in Typescript but doesn't seem as popular in Python.
🏷️ Python generic in functions
To make this happen we need to define a typing.TypeVar
to be used as the type of each argument and output.
import typing as t
T = t.TypeVar("T") # Defines the TypeVar
# v1, v2 and outcome should all be of type T
def sum(v1: T, v2: T) -> T:
return v1 + v2
numbers = sum(10, 20)
strs = sum("app", "le")
print(numbers)
print(strs)
Now your IDE infers the output based on the argument type:
🤌 Narrow down TypeVar types
This will work with any value that can be "summed". This is problematic though because someone may attempt to pass an invalid type that wouldn't make sense.
Like:
will_fail = sum(Exception("what?"), Exception("lol"))
Even though the IDE will resolve the outcome to another Exception
, the execution will raise an exception: TypeError: unsupported operand type(s) for +: 'Exception' and 'Exception'
.
To resolve this we can limit allowed types to work properly using TypeVar
's bound
argument.
T = t.TypeVar("T", bound=str | int | float)
This will ensure your type checker catches if any other type that is non str, int, or float is passed as an argument.
🎩 Using generics in classes
Generic classes also exist, and they allow more complex configurations.
Imagine a data-layer class that reads data from a data source and parses to some Python model for our scenario.
typing.TypeVar
is not enough anymore, we also need typing.Generic
.
Each data-layer class will be responsible for:
- Returning one entity based on id
- Listing all entities
- Creating an entity
Note we're not implementing this functionality as our purpose is only to understand how to use generics to define complex classes.
import typing as t
from datetime import datetime
from abc import ABC, abstractmethod
T = t.TypeVar("T") # 👈 We still use the TypeVar
# 👇 Now we must also rely on Generic to say the class
# accepts a type
class BaseDatabase(t.Generic[T], ABC):
@abstractmethod
def get_by_id(self, id: int) -> T:
...
@abstractmethod
def list_all(self) -> list[T]:
...
@abstractmethod
def create(self, entity: T) -> None:
...
We can start with the CompanyDatabase
class:
class Company: # Model sample for the company DB
name: str
phone: str
address: str
class CompanyDatabase(BaseDatabase[Company]):
def get_by_id(self, id: int) -> Company:
return super().get_by_id(id)
def list_all(self) -> list[Company]:
return super().list_all()
def create(self, entity: Company) -> None:
return super().create(entity)
Note the IDE is capable of defining the correct typing by itself:
and even if we don't create the methods, it also guesses everything correctly:
class EmployeeDatabase(BaseDatabase[Employee]):
pass
employee_db = EmployeeDatabase()
found_employee = employee_db.get_by_id(1)
all_employees = employee_db.list_all()
👎 Overcoming Python generic limitations
Unfortunately, Python is not as strong as TypeScript regarding typing.
Functions, differently from typing, can't use the Generic
which means we can't:
# ⚠️⚠️⚠️ Broken code for concept only:
import typing as T
T = t.TypeVar("T")
def get_something(t.Generic[T], v: str) -> T:
...
some_str = get_something[str]("str")
some_int = get_something[int]("10")
This won't work.
PEP 0484 suggests we pass the actual type as an argument to allow proper inference.
import typing as t
T = t.TypeVar("T")
# 👇 We must define whether we're receiving a TYPE of T
def get_something(output_type: t.Type[T], v: str) -> T:
...
# 👇 Now it works
some_str = get_something(str, "str")
some_int = get_something(int, "10")
The IDE recognizes it correctly:
But I still don't like it 😅 is it just me or it seems a bit hacky?
Using generic classes to behave as functions
I don't know you, but I'd like to keep the syntax we already follow using brackets, like:
x = list[str] # [str]
x = set[int] # [int]
class EmployeeDatabase(BaseDatabase[Employee]): # [Employee]
...
# WTH? This feels weird
x = get_something(int, 10)
I must implement something to feel more natural like we do on TypeScript.
Typescript relies consistently on <>
e.g. <int>
and <str>
. I do believe Python should follow the same logic for []
.
To make this happen we must define a class and override its __new__
magic method to behave like a function:
import typing as t
T = t.TypeVar("T")
# 👇 Keep class name lowercase so it feels like a function, name it as you want your '''function''' to be named:
class get_something(t.Generic[T]):
# 👇 This is the secret
def __new__(
cls,
v: str, # Add here as many args you think your function should take
):
generated_instance = super().__new__(cls)
return generated_instance.execute(v)
# 👇 Pretend this is your actual function implementation, name it anything you wish
def execute(self, v: str) -> T: # 👈 Define T as the return type
... # Do whatever you want
return t.cast(T, v) # 👈 Ensure returned type is T
# 💁♂️🪄🐰 It just works
some_str = get_something[str]("str")
some_int = get_something[int]("10")
It feels better to me.
Note this is not something "new" I'm coming up with. Some standard Python ""functions"" (that are not functions) do the same as defaultdict
:
Follow me for more Python magic.