Skip to content

Instantly share code, notes, and snippets.

Last active July 24, 2020 17:14
What would you like to do?
"typestate" with Literal types in Python 3.8
# check with mypy --allow-redefinition
# (this option should be called --no-demented-nazi-mode)
from enum import Enum, auto
from typing import cast, Generic, Literal, TypeVar
class DoorState(Enum):
CLOSED = auto()
LOCKED = auto()
OPEN = auto()
_S = TypeVar('_S', bound = DoorState)
class Door(Generic[_S]):
# these methods are "privileged"; they can cast because the invariants
# are guaranteed as long as calling code typechecks
def open(self: 'Door[Literal[DoorState.CLOSED]]'
) -> 'Door[Literal[DoorState.OPEN]]':
return cast(Door[Literal[DoorState.OPEN]], self)
def close(self: 'Door[Literal[DoorState.OPEN]]'
) -> 'Door[Literal[DoorState.CLOSED]]':
return cast(Door[Literal[DoorState.CLOSED]], self)
def lock(self: 'Door[Literal[DoorState.CLOSED]]'
) -> 'Door[Literal[DoorState.LOCKED]]':
return cast(Door[Literal[DoorState.LOCKED]], self)
def unlock(self: 'Door[Literal[DoorState.LOCKED]]'
) -> 'Door[Literal[DoorState.CLOSED]]':
return cast(Door[Literal[DoorState.CLOSED]], self)
def new_door() -> Door[Literal[DoorState.CLOSED]]:
return Door()
if __name__ == '__main__':
# a valid sequence
door = new_door() # a closed door
door = # open it
door = door.close() # close it again
door = door.lock() # lock it
door = door.unlock() # unlock it
door = # open it again
# various forbidden operations
# door = door.lock() # can't lock an open door
# door = door.unlock() # or unlock it
door = door.close()
# door = door.close() # or close a closed door
# check with mypy --allow-redefinition
from enum import Enum, auto
from typing import cast, Generic, Literal, TypeVar
class MeState(Enum):
UNFOOLED = auto()
FOOLED = auto()
_S = TypeVar('_S', bound = MeState)
class Me(Generic[_S]):
def __init__(self, name: str): = name
def fool(self: 'Me[Literal[MeState.UNFOOLED]]'
) -> 'Me[Literal[MeState.FOOLED]]':
print(f'shame on {}')
return cast(Me[Literal[MeState.FOOLED]], self)
if __name__ == '__main__':
me: Me[Literal[MeState.UNFOOLED]] = Me('George')
me = me.fool()
# me = me.fool() # you don't get fooled again!
# check with mypy --allow-redefinition
from enum import Enum, auto
from typing import cast, Generic, Literal, TypeVar
# natural numbers, as types
class Nat:
class Z(Nat):
_N = TypeVar('_N', bound = Nat)
class S(Nat, Generic[_N]):
# proofs of less-or-equal between type-level nats, as types
_M = TypeVar('_M', bound = Nat)
class _LTE(Generic[_N, _M]):
# lock LTE construction behind functions which ensure well-typedness
def duh() -> _LTE[Z, _N]:
return _LTE()
def well(prf: _LTE[_N, _M]) -> _LTE[S[_N], S[_M]]:
return _LTE()
# we'll parameterize Me by a Nat indicating the number of times fooled thus far
# since we can be fooled at most 5 times, we'll want a handy type-level four
FOUR = S[S[S[S[Z]]]]
class Me(Generic[_N]):
def __init__(self, name: str): = name
def fool(self: 'Me[_N]', prf: _LTE[_N, FOUR]) -> 'Me[S[_N]]':
print(f'shame on {}')
return cast(Me[S[_N]], self)
def naif(name: str) -> Me[Z]:
return Me(name)
if __name__ == '__main__':
me = naif('George')
me = me.fool(duh())
# unfortunately, this isn't Idris, so we have to bring our own proofs
# each time we get fooled, attesting that we've been fooled no more
# than four times already
me = me.fool(well(duh()))
me = me.fool(well(well(duh())))
me = me.fool(well(well(well(duh()))))
me = me.fool(well(well(well(well(duh())))))
# me = me.fool(no_possible_proof) # you don't get fooled again!
Copy link

There is both static and runtime checking. For example one would typically write Closed(door).open(), which checks that door is in state Closed at runtime and also checks statically that open is a valid transition - Open(door).open() will raise a warning in a type checker about the method not existing.

Copy link

Ah, that makes sense - I read the example a little backwards at first glance. Each state is a different type with a common base class. Nice.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment