diadefslib.py 8.5 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230
  1. # Licensed under the GPL: https://www.gnu.org/licenses/old-licenses/gpl-2.0.html
  2. # For details: https://github.com/PyCQA/pylint/blob/main/LICENSE
  3. # Copyright (c) https://github.com/PyCQA/pylint/blob/main/CONTRIBUTORS.txt
  4. """Handle diagram generation options for class diagram or default diagrams."""
  5. from __future__ import annotations
  6. import argparse
  7. from collections.abc import Generator
  8. from typing import Any
  9. import astroid
  10. from astroid import nodes
  11. from pylint.pyreverse.diagrams import ClassDiagram, PackageDiagram
  12. from pylint.pyreverse.inspector import Linker, Project
  13. from pylint.pyreverse.utils import LocalsVisitor
  14. # diagram generators ##########################################################
  15. class DiaDefGenerator:
  16. """Handle diagram generation options."""
  17. def __init__(self, linker: Linker, handler: DiadefsHandler) -> None:
  18. """Common Diagram Handler initialization."""
  19. self.config = handler.config
  20. self.module_names: bool = False
  21. self._set_default_options()
  22. self.linker = linker
  23. self.classdiagram: ClassDiagram # defined by subclasses
  24. def get_title(self, node: nodes.ClassDef) -> str:
  25. """Get title for objects."""
  26. title = node.name
  27. if self.module_names:
  28. title = f"{node.root().name}.{title}"
  29. return title # type: ignore[no-any-return]
  30. def _set_option(self, option: bool | None) -> bool:
  31. """Activate some options if not explicitly deactivated."""
  32. # if we have a class diagram, we want more information by default;
  33. # so if the option is None, we return True
  34. if option is None:
  35. return bool(self.config.classes)
  36. return option
  37. def _set_default_options(self) -> None:
  38. """Set different default options with _default dictionary."""
  39. self.module_names = self._set_option(self.config.module_names)
  40. all_ancestors = self._set_option(self.config.all_ancestors)
  41. all_associated = self._set_option(self.config.all_associated)
  42. anc_level, association_level = (0, 0)
  43. if all_ancestors:
  44. anc_level = -1
  45. if all_associated:
  46. association_level = -1
  47. if self.config.show_ancestors is not None:
  48. anc_level = self.config.show_ancestors
  49. if self.config.show_associated is not None:
  50. association_level = self.config.show_associated
  51. self.anc_level, self.association_level = anc_level, association_level
  52. def _get_levels(self) -> tuple[int, int]:
  53. """Help function for search levels."""
  54. return self.anc_level, self.association_level
  55. def show_node(self, node: nodes.ClassDef) -> bool:
  56. """True if builtins and not show_builtins."""
  57. if self.config.show_builtin:
  58. return True
  59. return node.root().name != "builtins" # type: ignore[no-any-return]
  60. def add_class(self, node: nodes.ClassDef) -> None:
  61. """Visit one class and add it to diagram."""
  62. self.linker.visit(node)
  63. self.classdiagram.add_object(self.get_title(node), node)
  64. def get_ancestors(
  65. self, node: nodes.ClassDef, level: int
  66. ) -> Generator[nodes.ClassDef, None, None]:
  67. """Return ancestor nodes of a class node."""
  68. if level == 0:
  69. return
  70. for ancestor in node.ancestors(recurs=False):
  71. if not self.show_node(ancestor):
  72. continue
  73. yield ancestor
  74. def get_associated(
  75. self, klass_node: nodes.ClassDef, level: int
  76. ) -> Generator[nodes.ClassDef, None, None]:
  77. """Return associated nodes of a class node."""
  78. if level == 0:
  79. return
  80. for association_nodes in list(klass_node.instance_attrs_type.values()) + list(
  81. klass_node.locals_type.values()
  82. ):
  83. for node in association_nodes:
  84. if isinstance(node, astroid.Instance):
  85. node = node._proxied
  86. if not (isinstance(node, nodes.ClassDef) and self.show_node(node)):
  87. continue
  88. yield node
  89. def extract_classes(
  90. self, klass_node: nodes.ClassDef, anc_level: int, association_level: int
  91. ) -> None:
  92. """Extract recursively classes related to klass_node."""
  93. if self.classdiagram.has_node(klass_node) or not self.show_node(klass_node):
  94. return
  95. self.add_class(klass_node)
  96. for ancestor in self.get_ancestors(klass_node, anc_level):
  97. self.extract_classes(ancestor, anc_level - 1, association_level)
  98. for node in self.get_associated(klass_node, association_level):
  99. self.extract_classes(node, anc_level, association_level - 1)
  100. class DefaultDiadefGenerator(LocalsVisitor, DiaDefGenerator):
  101. """Generate minimum diagram definition for the project :
  102. * a package diagram including project's modules
  103. * a class diagram including project's classes
  104. """
  105. def __init__(self, linker: Linker, handler: DiadefsHandler) -> None:
  106. DiaDefGenerator.__init__(self, linker, handler)
  107. LocalsVisitor.__init__(self)
  108. def visit_project(self, node: Project) -> None:
  109. """Visit a pyreverse.utils.Project node.
  110. create a diagram definition for packages
  111. """
  112. mode = self.config.mode
  113. if len(node.modules) > 1:
  114. self.pkgdiagram: PackageDiagram | None = PackageDiagram(
  115. f"packages {node.name}", mode
  116. )
  117. else:
  118. self.pkgdiagram = None
  119. self.classdiagram = ClassDiagram(f"classes {node.name}", mode)
  120. def leave_project(self, _: Project) -> Any:
  121. """Leave the pyreverse.utils.Project node.
  122. return the generated diagram definition
  123. """
  124. if self.pkgdiagram:
  125. return self.pkgdiagram, self.classdiagram
  126. return (self.classdiagram,)
  127. def visit_module(self, node: nodes.Module) -> None:
  128. """Visit an astroid.Module node.
  129. add this class to the package diagram definition
  130. """
  131. if self.pkgdiagram:
  132. self.linker.visit(node)
  133. self.pkgdiagram.add_object(node.name, node)
  134. def visit_classdef(self, node: nodes.ClassDef) -> None:
  135. """Visit an astroid.Class node.
  136. add this class to the class diagram definition
  137. """
  138. anc_level, association_level = self._get_levels()
  139. self.extract_classes(node, anc_level, association_level)
  140. def visit_importfrom(self, node: nodes.ImportFrom) -> None:
  141. """Visit astroid.ImportFrom and catch modules for package diagram."""
  142. if self.pkgdiagram:
  143. self.pkgdiagram.add_from_depend(node, node.modname)
  144. class ClassDiadefGenerator(DiaDefGenerator):
  145. """Generate a class diagram definition including all classes related to a
  146. given class.
  147. """
  148. def class_diagram(self, project: Project, klass: nodes.ClassDef) -> ClassDiagram:
  149. """Return a class diagram definition for the class and related classes."""
  150. self.classdiagram = ClassDiagram(klass, self.config.mode)
  151. if len(project.modules) > 1:
  152. module, klass = klass.rsplit(".", 1)
  153. module = project.get_module(module)
  154. else:
  155. module = project.modules[0]
  156. klass = klass.split(".")[-1]
  157. klass = next(module.ilookup(klass))
  158. anc_level, association_level = self._get_levels()
  159. self.extract_classes(klass, anc_level, association_level)
  160. return self.classdiagram
  161. # diagram handler #############################################################
  162. class DiadefsHandler:
  163. """Get diagram definitions from user (i.e. xml files) or generate them."""
  164. def __init__(self, config: argparse.Namespace) -> None:
  165. self.config = config
  166. def get_diadefs(self, project: Project, linker: Linker) -> list[ClassDiagram]:
  167. """Get the diagram's configuration data.
  168. :param project:The pyreverse project
  169. :type project: pyreverse.utils.Project
  170. :param linker: The linker
  171. :type linker: pyreverse.inspector.Linker(IdGeneratorMixIn, LocalsVisitor)
  172. :returns: The list of diagram definitions
  173. :rtype: list(:class:`pylint.pyreverse.diagrams.ClassDiagram`)
  174. """
  175. # read and interpret diagram definitions (Diadefs)
  176. diagrams = []
  177. generator = ClassDiadefGenerator(linker, self)
  178. for klass in self.config.classes:
  179. diagrams.append(generator.class_diagram(project, klass))
  180. if not diagrams:
  181. diagrams = DefaultDiadefGenerator(linker, self).visit(project)
  182. for diagram in diagrams:
  183. diagram.extract_relationships()
  184. return diagrams