title: Typing generic state functions in Python
body: <p>I haven't been able to make either pylance or mypy happy with typing a generic state function. If I hardcode the state machine (not using a generic base class), I get no errors. In my actual implementation, the generic class is not also an abstract class, but I get the same problem, so the issue is unrelated to that.</p>
  <p>Here's the code. Mind you, the implementation here is unrelated to what my state machine does. I just created something minimal that shows the same problems. I've also struggled to make type checking happy having a particular state function designate the end of the stream (e.g. <code>state_done</code>), instead of looking for the value <code>None</code>.</p>
  <pre><code>from __future__ import annotations
  from abc import ABC, abstractmethod
  from typing import TypeVar, Generic, Optional, Callable
  import random
  
  
  _T = TypeVar(&quot;_T&quot;)
  StateFunc = Optional[Callable[[_T], Optional[&quot;StateFunc&quot;]]]
  
  
  class GenericMachine(Generic[_T], ABC):
      def __init__(self):
          self._state: StateFunc[_T]
  
      def run(self) -&gt; None:
          while self._state is not None:
              self._state = self._state(self.get_next_event())
  
      @abstractmethod
      def get_next_event(self) -&gt; _T:
          pass
  
  
  class Catch(GenericMachine[float]):
      _MyStateFunc = StateFunc[float]
  
      def __init__(self):
          print(&quot;person has ball&quot;)
          self._state: Catch._MyStateFunc = self._state_person_with_ball  # ERROR 1
  
      def get_next_event(self) -&gt; float:
          return random.random()
  
      def _state_person_with_ball(self, r: float) -&gt; Catch._MyStateFunc:
          if r &lt; 0.2:
              return None
          print(&quot;ball in air&quot;)
          return self._state_ball_in_air  # ERROR 2
  
      def _state_ball_in_air(self, r: float) -&gt; Catch._MyStateFunc:
          print(&quot;dog has ball&quot;)
          return self._state_dog_with_ball  # ERROR 2
  
      def _state_dog_with_ball(self, r: float) -&gt; Catch._MyStateFunc:
          if r &lt; 0.2:
              return None
          print(&quot;person has ball&quot;)
          return self._state_person_with_ball  # ERROR 2
  
  
  if __name__ == &quot;__main__&quot;:
      Catch().run()
  </code></pre>
  <p>EDIT: I realize that mypy won't get past the recursive type definition. I get similar errors on mypy when I just declare <code>StateFunc</code> to return a <code>Callable</code>.</p>
  <p>Pylance reports the following for ERROR 1 (via VS Code):</p>
  <pre><code>Cannot assign member &quot;_state&quot; for type &quot;Catch&quot;
    Expression of type &quot;(r: float) -&gt; (_p0: float) -&gt; StateFunc | None | None&quot; cannot be assigned to member &quot;_state&quot; of class &quot;Catch&quot;
      Type &quot;(r: float) -&gt; (_p0: float) -&gt; StateFunc | None | None&quot; cannot be assigned to type &quot;(_p0: float) -&gt; StateFunc | None | None&quot;
        Type &quot;(r: float) -&gt; (_p0: float) -&gt; StateFunc | None | None&quot; cannot be assigned to type &quot;(_p0: float) -&gt; StateFunc | None&quot;
        Type cannot be assigned to type &quot;None&quot;
          Function return type &quot;(_p0: float) -&gt; StateFunc | None | None&quot; is incompatible with type &quot;StateFunc | None&quot;
            Type &quot;(_p0: float) -&gt; StateFunc | None | None&quot; cannot be assigned to type &quot;StateFunc | None&quot;PylancereportGeneralTypeIssues
  (variable) _state: (r: float) -&gt; Unknown
  </code></pre>
  <p>And it reports the following for ERROR 2 (via VS Code):</p>
  <pre><code>Expression of type &quot;(r: float) -&gt; (_p0: float) -&gt; StateFunc | None | None&quot; cannot be assigned to return type &quot;(_p0: float) -&gt; StateFunc | None | None&quot;
    Type &quot;(r: float) -&gt; (_p0: float) -&gt; StateFunc | None | None&quot; cannot be assigned to type &quot;(_p0: float) -&gt; StateFunc | None | None&quot;
      Type &quot;(r: float) -&gt; (_p0: float) -&gt; StateFunc | None | None&quot; cannot be assigned to type &quot;(_p0: float) -&gt; StateFunc | None&quot;
      Type cannot be assigned to type &quot;None&quot;
        Function return type &quot;(_p0: float) -&gt; StateFunc | None | None&quot; is incompatible with type &quot;StateFunc | None&quot;
          Type &quot;(_p0: float) -&gt; StateFunc | None | None&quot; cannot be assigned to type &quot;StateFunc | None&quot;
            Type &quot;(_p0: float) -&gt; StateFunc | None&quot; cannot be assigned to type &quot;StateFunc | None&quot;PylancereportGeneralTypeIssues
  (parameter) self: Catch
  </code></pre>
  <p>Any idea how to make Pylance happy?</p>
  <p>UPDATE: The problem goes away if I define StateFunc as follows:</p>
  <p><code>StateFunc = Optional[Callable[[_T], Any]]</code></p>