Differentiating custom python types that have the same fields in the type checking stage

Question

I have a custom python type like this:

class Pointer(TypedDict):
  id: str

id is a pointer to other data structures, but the id field doesn't contain information about which data structure it is pointing to.

Ideally, I would be able to type the Pointer class as to which data structure the Pointer is pointing to, so that I don't accidentally confuse Pointers. For example:

# I know the [subscript] syntax isnt right, but this is what I imagine
# a solution might look like.
def get_pointer_to_bar() -> Pointer[Bar]:
  return { 'id': 'abc123' }

def get_pointer_to_baz() -> Pointer[Baz]:
  return { 'id': 'xyz789' }

def use_bar_pointer(p: Pointer[Bar]) -> None:
  do_something_with_p(p)

This would allow me to make sure I'm passing in the right Pointers at the type checking stage, while also letting me create functions that operate on specific kinds of pointers.

Is 'subscripting' custom types possible? Or in general is there a way to tell the type checker how to differentiate types that are ostensibly the same?

Answer 1

You can't seem to use TypedDict since it imposes its own restrictions on subclassing, but something like

from typing import Generic, TypeVar

T = TypeVar("T")


class Pointer(Generic[T]):
    __slots__ = ("id",)

    def __init__(self, id: str):
        self.id = id


class Bar:
    pass


class Baz:
    pass


def get_pointer_to_bar() -> Pointer[Bar]:
    return Pointer(id="abc123")


def get_pointer_to_baz() -> Pointer[Baz]:
    return Pointer(id="xyz456")


def use_bar_pointer(p: Pointer[Bar]) -> None:
    print(p)


p: Pointer[Baz] = Pointer(id="xyz456")
use_bar_pointer(p)

does the trick – i.e. Mypy yells:

xx.py:34: error: Argument 1 to "use_bar_pointer" has incompatible type "Pointer[Baz]"; expected "Pointer[Bar]"