Generic functions and generic classes in Python

gui commited 9 months ago · 🐍 Python typing

"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)

Sample
Single generic function with proper typing

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:

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:

IDE setting correct types from generic

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()

IDE guessing expected types from generic

👎 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:

Overcoming Python's limitation

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")
Overcoming Python's limitation again but pretty

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:

defaultdict lied all the time to you

Follow me for more Python magic.

  • LinkedIn
  • Tumblr
  • Reddit
  • Google+
  • Pinterest
  • Pocket