Static type checkers and Language Servers not recognizing attributes of objects that are subclasses

Question

I apologize if the title looks a bit convoluted, it is a bit difficult to phrase the problem.

Say I am building a graph with nodes and edges. It has a base class Node with common attributes to all nodes, and each specific node type can inherit from it extending with more attributes.

Take the following as example with two node types:

from typing import TypeVar


# nodes.py
class Node:
    def __init__(self, base_param: str) -> None:
        self.base_param = base_param

    # other base methods...
        

class NodeType1(Node):
    def __init__(self, base_param: str, node_param_type1: int) -> None:
        super().__init__(base_param)
        self.node_param_type1 = node_param_type1

    # other type1 methods...


class NodeType2(Node):
    def __init__(self, base_param: str, node_param_type2: int) -> None:
        super().__init__(base_param)
        self.node_param_type2 = node_param_type2

    # other type2 methods...


NodeSubClassType = TypeVar("NodeSubClassType", bound=Node)

It also has an Edge base class with a source and a target, which are both subclasses of Node. Take the following as example with an edge connecting a NodeType1 to a NodeType2.

# edges.py
class Edge:
    def __init__(self, source: NodeSubClassType, target: NodeSubClassType) -> None:
        self.source = source
        self.target = target

    # other base methods...


class EdgeType1ToType2(Edge):
    # Must contain `NodeType1` as `source` and `NodeType2` as `target` when instantiated
    pass
    
    # or add other `type1 to type2` methods

My problem is that, whenever I create an instance of EdgeType1ToType2 passing the correct objects, only the base attribute base_param of the source and target gets recognized by static checkers (I tried mypy) and language servers (I tried pylsp). For instance, in the following script:

if __name__ == "__main__":

    node_type1 = NodeType1("base_param", 1)
    node_type2 = NodeType2("base_param", 2)
    edge_12 = EdgeType1ToType2(source=node_type1, target=node_type2)

    node_type1.base_param  # this has no issue
    edge_12.source.base_param  # this has no issue
    node_type1.node_param_type1  # this has no issue
    edge_12.source.node_param_type1  # this is not recognized by `mypy` or the LSP

If I run mypy I get as message

error: "NodeSubClassType" has no attribute "node_param_type1"  [attr-defined]
Found 1 error in 1 file (checked 1 source file)

only for the edge_12.source.node_param_type1 line. Also, the language server does not recognize node_param_type1 as an attribute of edge_12.source. Similar issues with edge_12.target.

I tried:

• Adding a superfluous __init__ passing proper parameters to the subclass like

  class EdgeType1ToType2(Edge):
  
      def __init__(self, source: NodeType1, target: NodeType2) -> None:
          super().__init__(source, target)
  

  but it does not work.

• and replacing the custom type NodeSubClassType with the superclass Node directly, but I got the same message with Node instead.

I expect a solution that succeeds in static check analysis and possibly gets recognized as valid attribute by language servers.

Answer 1

As @STerliakov commented, a possible solution is to create two explicit node types, one to be used as source and another target of an edge:

_Source = TypeVar("_Source", bound=Node)
_Target = TypeVar("_Target", bound=Node)

and then create the base edge class like this:

from typing import Generic

class Edge(Generic[_Source, _Target]):
    def __init__(self, source: _Source, target: _Target) -> None:
        self.source = source
        self.target = target

    # other base methods...

Then, the concrete subclasses should be created by passing the required node types explicitly:

class EdgeType1ToType2(Edge[NodeType1, NodeType2]):
    # Must contain `NodeType1` as `source` and `NodeType2` as `target` when instantiated
    pass

And then, when using the original script

if __name__ == "__main__":
    node_type1 = NodeType1("base_param", 1)
    node_type2 = NodeType2("base_param", 2)
    edge_12 = EdgeType1ToType2(source=node_type1, target=node_type2)

    node_type1.base_param  # this has no issue
    edge_12.source.base_param  # this has no issue
    node_type1.node_param_type1  # this has no issue
    edge_12.source.node_param_type1  # this has no issue

mypy shows no issue anymore, and when typing edge_12.source. the editor recognizes node_param_type1 as a valid attribute.