Type hints and code completion. How to specialize a generic class with types

Question

In C++ I would use templates. I'm checking whether the hints work using the PyCharm auto-completion feature.

I'm working on a "generalized" FSM library (Finite-State Machine). I would like to preserve the type-hints that I had before. They worked when I had the FSM "tailored" to a specific subject.

FSM is just an example. It can be anything that is generic and by specializing it I would like to pass the specializing types down to the generic class.

The FSM looks as follows:

• There is the owner, the subject that uses the FSM
• There are the states that have a common state
• There is the state manager that holds the current state

The lib (simplified):

class BaseState:
    def __init__(self, owner: "???", state_mgr: "StateManager"):
        self.owner = owner
        self.state_mgr = state_mgr

class StateManager:
    def __init__(self, owner, base_state_type):
        self.__owner = owner
        self.__base_state_type = base_state_type
        self.current_state: Type[base_state_type] = base_state_type(self.__owner, self) # Doesn't work

    def transition_to(self, next_state_type):
        assert issubclass(next_state_type, self.__base_state_type)
        self.current_state = next_state_type(self.__owner, self)

The new subject that uses the lib:

class FileSender_BaseState(fsm.BaseState):
    def on_event(self, msg):
        pass

class FileSender_Idle(FileSender_BaseState):
    def on_event(self, msg):
        self.owner.foo()

class FileSender_Sending(FileSender_BaseState):
    def on_event(self, msg):
        pass


class FileSender:
    def __init__(self):
        self._state_mgr = fsm.StateManager(self, FileSender_BaseState)
        self._state_mgr.transition_to(FileSender_Idle)
        self._state_mgr.current_state.on_event()

    def foo(self):
        pass

• How to pass to the BaseState that in this scenario the owner is of type FileSender?
• How to pass to the StateManager that the current_state is of type FileSender_BaseState?
• (optional) How to "type-hint" the next_state_type parameter of the transition_to() method that it is a (sub)type of FileSender_BaseState

For the first point, I have tried to be more precise with the fsm.BaseState initialization, but it doesn't improve.

class FileSender_BaseState(fsm.BaseState):
    def __init__(self, owner: "FileSender", state_mgr: "fsm.StateManager"): # Specifing owner type didn't improve anything
        super().__init__(owner, state_mgr)

For the second point, I've tried to use the typing.Type with the argument that holds the type of the owner, but it also doesn't work, surprisingly

class StateManager:
    def __init__(self, owner, base_state_type):
       self.current_state: Type[base_state_type] = base_state_type()

The third attempt failed miserably because I cannot do

    def transition_to(self, next_state_type: typing.Type["self.__base_state_type"]): # Unresolved reference 'self' of course

I would suspect that PyCharm auto-completion:

class FileSender_Idle(FileSender_BaseState):
    def on_event(self, msg):
        self.owner. # propose the foo() method

self._state_mgr.current_state. # propose the on_event() method

Here is how it was working before: https://pastebin.com/MekPqbnJ

Answer 1

I think we might run into limitations of the currently available type system here.

The way I see it, you essentially want two generic classes here, that depend on each other.

BaseState should be generic in terms of its owner. This would allow the type checker to infer the owner's interface by binding an owner type variable in the BaseState constructor.

And StateManager should be generic in terms of its state (type) and by extension also its owner. This would at least allow to infer the interface of the state after initialization.

But type variables do not allow unspecified generic types as upper bounds/constraints. (see discussion here) And in this case, that would be a proper way to express the that the StateManager is parameterized by its state type, which in turn is parameterized by its owner, and that owner is bound by the owner of the state manager.

If it were possible, I would do something like this:

from __future__ import annotations
from typing import Generic, TypeVar


OwnerT = TypeVar("OwnerT")
StateT = TypeVar("StateT", bound="BaseState[OwnerT]")  # not valid


class BaseState(Generic[OwnerT]):
    def __init__(self, owner: OwnerT, state_mgr: StateManager[OwnerT]) -> None:
        self.owner = owner
        self.state_mgr = state_mgr


class StateManager(Generic[OwnerT, StateT]):
    __owner: OwnerT
    __base_state_type: type[StateT[OwnerT]]
    current_state: StateT[OwnerT]

    def __init__(self, owner: OwnerT, base_state_type: type[StateT[OwnerT]]) -> None:
        self.__owner = owner
        self.__base_state_type = base_state_type
        self.current_state = base_state_type(self.__owner, self)
...

Alas, the type system has no expression for this.

---

Workaround

I think the closest thing to what you are trying to achieve could be something like this:

from __future__ import annotations
from typing import Generic, TypeVar


OwnerT = TypeVar("OwnerT")
StateT = TypeVar("StateT")


class BaseState(Generic[OwnerT]):
    def __init__(self, owner: OwnerT, state_mgr: StateManager[OwnerT, StateT]) -> None:
        self.owner = owner
        self.state_mgr = state_mgr


class StateManager(Generic[OwnerT, StateT]):
    _owner: OwnerT
    _base_state_type: type[StateT]
    current_state: StateT

    def __init__(self, owner: OwnerT, base_state_type: type[StateT]) -> None:
        self._owner = owner
        self._base_state_type = base_state_type
        self.current_state = base_state_type(self._owner, self)  # type: ignore[call-arg]

    def transition_to(self, next_state_type: type[StateT]) -> None:
        assert issubclass(next_state_type, self._base_state_type)
        self.current_state = next_state_type(self._owner, self)  # type: ignore[call-arg]


class FileSenderBaseState(BaseState["FileSender"]):
    def on_event(self, msg: object) -> None:
        pass


class FileSenderIdle(FileSenderBaseState):
    def on_event(self, msg: object) -> None:
        self.owner.foo()


class FileSenderSending(FileSenderBaseState):
    def on_event(self, msg: object) -> None:
        pass


class FileSender:
    def __init__(self) -> None:
        self._state_mgr = StateManager(self, FileSenderBaseState)
        self._state_mgr.transition_to(FileSenderIdle)
        # Type checker still infers current state as `FileSenderBaseState`
        assert isinstance(self._state_mgr.current_state, FileSenderIdle)
        self._state_mgr.current_state.on_event(...)

    def foo(self) -> None:
        pass

---

Caveats

Since we have no upper bound on StateT, type[StateT] is treated by the type checker as object, which raises an error, when we initialize it with arguments (because the object constructor does not take arguments). Thus, we utilize a specific type: ignore directive, whenever we initialize type[StateT].

The second limitation is that there is not really a way to let the type checker know that the current_state type on an instance of StateManager has changed without explicitly setting that attribute on that instance or asserting the type explicitly. Since your transition_to method changes the type, this will remain opaque to the outside. The only workaround I can think of is that assert statement.

Other than that, this setup satisfies mypy and should give you the desired auto-suggestions from your IDE.

If you add reveal_type(self._state_mgr._owner) and reveal_type(self._state_mgr.current_state) at the end in that FileSender.__init__ method, mypy will output the following:

Revealed type is "FileSender"
Revealed type is "FileSenderIdle"

---

Depending on what your goals are, it may be worth reconsidering the entire design. It is hard to comment on that without more context though. I hope this illustrates at least a few concepts and tricks you can use for type safety.

I would suggest keeping this question open, i.e. not accepting my answer (if you were so inclined in the first place) in the hopes that someone comes along with a better idea.