Add typing to pylint.pyreverse.diagrams.py (#6547)

* Add typing to `pylint.pyreverse.diagrams.py` * Apply suggestions from code review Co-authored-by: Daniël van Noord <13665637+DanielNoord@users.noreply.github.com> * Remove unused method `nodes()` and use node classes through `nodes.<CLASS>` Co-authored-by: Daniël van Noord <13665637+DanielNoord@users.noreply.github.com>

Author: DudeNr33

pylint/pyreverse/diagrams.py

@@ -3,6 +3,12 @@
 # Copyright (c) https://github.com/PyCQA/pylint/blob/main/CONTRIBUTORS.txt
 
 """Diagram objects."""
+
+from __future__ import annotations
+
+from collections.abc import Iterable
+from typing import Any
+
 import astroid
 from astroid import nodes
 
@@ -13,11 +19,20 @@
 class Figure:
  """Base class for counter handling."""
 
+ def __init__(self) -> None:
+ self.fig_id: str = ""
+
 
 class Relationship(Figure):
  """A relationship from an object in the diagram to another."""
 
- def __init__(self, from_object, to_object, relation_type, name=None):
+ def __init__(
+ self,
+ from_object: DiagramEntity,
+ to_object: DiagramEntity,
+ relation_type: str,
+ name: str | None = None,
+ ):
  super().__init__()
  self.from_object = from_object
  self.to_object = to_object
@@ -28,59 +43,76 @@ def __init__(self, from_object, to_object, relation_type, name=None):
 class DiagramEntity(Figure):
  """A diagram object, i.e. a label associated to an astroid node."""
 
- def __init__(self, title="No name", node=None):
+ default_shape = ""
+
+ def __init__(
+ self, title: str = "No name", node: nodes.NodeNG | None = None
+ ) -> None:
  super().__init__()
  self.title = title
- self.node = node
+ self.node: nodes.NodeNG = node if node else nodes.NodeNG()
+ self.shape = self.default_shape
 
 
 class PackageEntity(DiagramEntity):
  """A diagram object representing a package."""
 
+ default_shape = "package"
+
 
 class ClassEntity(DiagramEntity):
  """A diagram object representing a class."""
 
- def __init__(self, title, node):
+ default_shape = "class"
+
+ def __init__(self, title: str, node: nodes.ClassDef) -> None:
  super().__init__(title=title, node=node)
- self.attrs = None
- self.methods = None
+ self.attrs: list[str] = []
+ self.methods: list[nodes.FunctionDef] = []
 
 
 class ClassDiagram(Figure, FilterMixIn):
  """Main class diagram handling."""
 
  TYPE = "class"
 
- def __init__(self, title, mode):
+ def __init__(self, title: str, mode: str) -> None:
  FilterMixIn.__init__(self, mode)
  Figure.__init__(self)
  self.title = title
- self.objects = []
- self.relationships = {}
- self._nodes = {}
- self.depends = []
+ # TODO: Specify 'Any' after refactor of `DiagramEntity`
+ self.objects: list[Any] = []
+ self.relationships: dict[str, list[Relationship]] = {}
+ self._nodes: dict[nodes.NodeNG, DiagramEntity] = {}
 
- def get_relationships(self, role):
+ def get_relationships(self, role: str) -> Iterable[Relationship]:
  # sorted to get predictable (hence testable) results
  return sorted(
  self.relationships.get(role, ()),
  key=lambda x: (x.from_object.fig_id, x.to_object.fig_id),
  )
 
- def add_relationship(self, from_object, to_object, relation_type, name=None):
+ def add_relationship(
+ self,
+ from_object: DiagramEntity,
+ to_object: DiagramEntity,
+ relation_type: str,
+ name: str | None = None,
+ ) -> None:
  """Create a relationship."""
  rel = Relationship(from_object, to_object, relation_type, name)
  self.relationships.setdefault(relation_type, []).append(rel)
 
- def get_relationship(self, from_object, relation_type):
+ def get_relationship(
+ self, from_object: DiagramEntity, relation_type: str
+ ) -> Relationship:
  """Return a relationship or None."""
  for rel in self.relationships.get(relation_type, ()):
  if rel.from_object is from_object:
  return rel
  raise KeyError(relation_type)
 
- def get_attrs(self, node):
+ def get_attrs(self, node: nodes.ClassDef) -> list[str]:
  """Return visible attributes, possibly with class name."""
  attrs = []
  properties = [
@@ -101,7 +133,7 @@ def get_attrs(self, node):
  attrs.append(node_name)
  return sorted(attrs)
 
- def get_methods(self, node):
+ def get_methods(self, node: nodes.ClassDef) -> list[nodes.FunctionDef]:
  """Return visible methods."""
  methods = [
  m
@@ -113,14 +145,14 @@ def get_methods(self, node):
  ]
  return sorted(methods, key=lambda n: n.name)
 
- def add_object(self, title, node):
+ def add_object(self, title: str, node: nodes.ClassDef) -> None:
  """Create a diagram object."""
  assert node not in self._nodes
- ent = DiagramEntity(title, node)
+ ent = ClassEntity(title, node)
  self._nodes[node] = ent
  self.objects.append(ent)
 
- def class_names(self, nodes_lst):
+ def class_names(self, nodes_lst: Iterable[nodes.NodeNG]) -> list[str]:
  """Return class names if needed in diagram."""
  names = []
  for node in nodes_lst:
@@ -136,30 +168,26 @@ def class_names(self, nodes_lst):
  names.append(node_name)
  return names
 
- def nodes(self):
- """Return the list of underlying nodes."""
- return self._nodes.keys()
-
- def has_node(self, node):
+ def has_node(self, node: nodes.NodeNG) -> bool:
  """Return true if the given node is included in the diagram."""
  return node in self._nodes
 
- def object_from_node(self, node):
+ def object_from_node(self, node: nodes.NodeNG) -> DiagramEntity:
  """Return the diagram object mapped to node."""
  return self._nodes[node]
 
- def classes(self):
+ def classes(self) -> list[ClassEntity]:
  """Return all class nodes in the diagram."""
- return [o for o in self.objects if isinstance(o.node, nodes.ClassDef)]
+ return [o for o in self.objects if isinstance(o, ClassEntity)]
 
- def classe(self, name):
+ def classe(self, name: str) -> ClassEntity:
  """Return a class by its name, raise KeyError if not found."""
  for klass in self.classes():
  if klass.node.name == name:
  return klass
  raise KeyError(name)
 
- def extract_relationships(self):
+ def extract_relationships(self) -> None:
  """Extract relationships between nodes in the diagram."""
  for obj in self.classes():
  node = obj.node
@@ -205,18 +233,25 @@ class PackageDiagram(ClassDiagram):
 
  TYPE = "package"
 
- def modules(self):
+ def modules(self) -> list[PackageEntity]:
  """Return all module nodes in the diagram."""
- return [o for o in self.objects if isinstance(o.node, nodes.Module)]
+ return [o for o in self.objects if isinstance(o, PackageEntity)]
 
- def module(self, name):
+ def module(self, name: str) -> PackageEntity:
  """Return a module by its name, raise KeyError if not found."""
  for mod in self.modules():
  if mod.node.name == name:
  return mod
  raise KeyError(name)
 
- def get_module(self, name, node):
+ def add_object(self, title: str, node: nodes.Module) -> None:
+ """Create a diagram object."""
+ assert node not in self._nodes
+ ent = PackageEntity(title, node)
+ self._nodes[node] = ent
+ self.objects.append(ent)
+
+ def get_module(self, name: str, node: nodes.Module) -> PackageEntity:
  """Return a module by its name, looking also for relative imports;
  raise KeyError if not found.
  """
@@ -232,29 +267,29 @@ def get_module(self, name, node):
  return mod
  raise KeyError(name)
 
- def add_from_depend(self, node, from_module):
+ def add_from_depend(self, node: nodes.ImportFrom, from_module: str) -> None:
  """Add dependencies created by from-imports."""
  mod_name = node.root().name
  obj = self.module(mod_name)
  if from_module not in obj.node.depends:
  obj.node.depends.append(from_module)
 
- def extract_relationships(self):
+ def extract_relationships(self) -> None:
  """Extract relationships between nodes in the diagram."""
  super().extract_relationships()
- for obj in self.classes():
+ for class_obj in self.classes():
  # ownership
  try:
- mod = self.object_from_node(obj.node.root())
- self.add_relationship(obj, mod, "ownership")
+ mod = self.object_from_node(class_obj.node.root())
+ self.add_relationship(class_obj, mod, "ownership")
  except KeyError:
  continue
- for obj in self.modules():
- obj.shape = "package"
+ for package_obj in self.modules():
+ package_obj.shape = "package"
  # dependencies
- for dep_name in obj.node.depends:
+ for dep_name in package_obj.node.depends:
  try:
- dep = self.get_module(dep_name, obj.node)
+ dep = self.get_module(dep_name, package_obj.node)
  except KeyError:
  continue
- self.add_relationship(obj, dep, "depends")
+ self.add_relationship(package_obj, dep, "depends")