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

You might like https://github.com/alexmojaki/friendly_states

@alexmojaki thanks - looks like a nice (and much more complete) runtime implementation of similar ideas. I like the ability to statically check that transition functions are used correctly, but it'd be great to find a way to have both static and dynamic checking.

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.

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.