"typestate" with Literal types in Python 3.8

typestate.py

# check with mypy typestate.py --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 = door.open() # open it
    door = door.close() # close it again
    door = door.lock() # lock it
    door = door.unlock() # unlock it
    door = door.open() # 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

typestate2.py

# check with mypy typestate.py --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):
        self.name = name

    def fool(self: 'Me[Literal[MeState.UNFOOLED]]'
            ) -> 'Me[Literal[MeState.FOOLED]]':
        print(f'shame on {self.name}')
        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!

typestate3.py

# check with mypy typestate.py --allow-redefinition

from enum import Enum, auto
from typing import cast, Generic, Literal, TypeVar

# natural numbers, as types

class Nat:
    pass

class Z(Nat):
    pass

_N = TypeVar('_N', bound = Nat)
class S(Nat, Generic[_N]):
    pass

# proofs of less-or-equal between type-level nats, as types

_M = TypeVar('_M', bound = Nat)
class _LTE(Generic[_N, _M]):
    pass

# 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):
        self.name = name

    def fool(self: 'Me[_N]', prf: _LTE[_N, FOUR]) -> 'Me[S[_N]]':
        print(f'shame on {self.name}')
        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!

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.