Changing python type hints for subclasses-CodePudding

Imagine the following code to handle graphs using a base class and a DAG subclass:

class NodeBase:
    def some_node_func(self):
        pass

class GraphBase:
    def add(self, node: NodeBase):
        node.some_node_func()

class DirectedNode(NodeBase):
    def some_dag_func(self):
        pass

class DirectedGraph(GraphBase):
    def add(self, node: DirectedNode):
        node.some_node_func()
        node.some_dag_func()

When I try to use this code with mypy I get errors like:

error: Argument 1 of "add" is incompatible with supertype "GraphBase"; 
supertype defines the argument type as "NodeBase"  [override]

My question is superficially similar to Python: how to handle type hinting in subclasses? but I actually need different behavior from the DirectedGraph.add function that relies on the corresponding DirectedNode functionality.

I realize that this "violates the Liskov substitution principle" but I don't need to be able to add DirectedNode instances to non-DAG graphs.

How can I structure things like this so that mypy doesn't complain? I'd like to avoid just disabling the checks, if possible.

CodePudding user response：

Parameterize GraphBase on the type used to represent a node, instead of hard-coding references to NodeBase

from typing import Generic, TypeVar


N = TypeVar('N', bound=NodeBase)


class NodeBase:
    def some_node_func(self):
        pass

class GraphBase(Generic[N]):
    def add(self, node: N):
        node.some_node_func()

class DirectedNode(NodeBase):
    def some_dag_func(self):
        pass

class DirectedGraph(GraphBase[DirectedNode]):
    def add(self, node: DirectedNode):
        node.some_node_func()
        node.some_dag_func()