printer.py 3.8 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132
  1. # Copyright (c) 2021 Pierre Sassoulas <pierre.sassoulas@gmail.com>
  2. # Copyright (c) 2021 Daniël van Noord <13665637+DanielNoord@users.noreply.github.com>
  3. # Copyright (c) 2021 Ashley Whetter <ashley@awhetter.co.uk>
  4. # Copyright (c) 2021 Marc Mueller <30130371+cdce8p@users.noreply.github.com>
  5. # Copyright (c) 2021 Andreas Finkler <andi.finkler@gmail.com>
  6. # Licensed under the GPL: https://www.gnu.org/licenses/old-licenses/gpl-2.0.html
  7. # For details: https://github.com/PyCQA/pylint/blob/main/LICENSE
  8. """
  9. Base class defining the interface for a printer.
  10. """
  11. from abc import ABC, abstractmethod
  12. from enum import Enum
  13. from typing import List, NamedTuple, Optional
  14. from astroid import nodes
  15. from pylint.pyreverse.utils import get_annotation_label
  16. class NodeType(Enum):
  17. CLASS = "class"
  18. INTERFACE = "interface"
  19. PACKAGE = "package"
  20. class EdgeType(Enum):
  21. INHERITS = "inherits"
  22. IMPLEMENTS = "implements"
  23. ASSOCIATION = "association"
  24. USES = "uses"
  25. class Layout(Enum):
  26. LEFT_TO_RIGHT = "LR"
  27. RIGHT_TO_LEFT = "RL"
  28. TOP_TO_BOTTOM = "TB"
  29. BOTTOM_TO_TOP = "BT"
  30. class NodeProperties(NamedTuple):
  31. label: str
  32. attrs: Optional[List[str]] = None
  33. methods: Optional[List[nodes.FunctionDef]] = None
  34. color: Optional[str] = None
  35. fontcolor: Optional[str] = None
  36. class Printer(ABC):
  37. """Base class defining the interface for a printer"""
  38. def __init__(
  39. self,
  40. title: str,
  41. layout: Optional[Layout] = None,
  42. use_automatic_namespace: Optional[bool] = None,
  43. ) -> None:
  44. self.title: str = title
  45. self.layout = layout
  46. self.use_automatic_namespace = use_automatic_namespace
  47. self.lines: List[str] = []
  48. self._indent = ""
  49. self._open_graph()
  50. def _inc_indent(self) -> None:
  51. """increment indentation"""
  52. self._indent += " "
  53. def _dec_indent(self) -> None:
  54. """decrement indentation"""
  55. self._indent = self._indent[:-2]
  56. @abstractmethod
  57. def _open_graph(self) -> None:
  58. """Emit the header lines, i.e. all boilerplate code that defines things like layout etc."""
  59. def emit(self, line: str, force_newline: Optional[bool] = True) -> None:
  60. if force_newline and not line.endswith("\n"):
  61. line += "\n"
  62. self.lines.append(self._indent + line)
  63. @abstractmethod
  64. def emit_node(
  65. self,
  66. name: str,
  67. type_: NodeType,
  68. properties: Optional[NodeProperties] = None,
  69. ) -> None:
  70. """Create a new node. Nodes can be classes, packages, participants etc."""
  71. @abstractmethod
  72. def emit_edge(
  73. self,
  74. from_node: str,
  75. to_node: str,
  76. type_: EdgeType,
  77. label: Optional[str] = None,
  78. ) -> None:
  79. """Create an edge from one node to another to display relationships."""
  80. @staticmethod
  81. def _get_method_arguments(method: nodes.FunctionDef) -> List[str]:
  82. if method.args.args:
  83. arguments: List[nodes.AssignName] = [
  84. arg for arg in method.args.args if arg.name != "self"
  85. ]
  86. else:
  87. arguments = []
  88. annotations = dict(zip(arguments, method.args.annotations[1:]))
  89. for arg in arguments:
  90. annotation_label = ""
  91. ann = annotations.get(arg)
  92. if ann:
  93. annotation_label = get_annotation_label(ann)
  94. annotations[arg] = annotation_label
  95. return [
  96. f"{arg.name}: {ann}" if ann else f"{arg.name}"
  97. for arg, ann in annotations.items()
  98. ]
  99. def generate(self, outputfile: str) -> None:
  100. """Generate and save the final outputfile."""
  101. self._close_graph()
  102. with open(outputfile, "w", encoding="utf-8") as outfile:
  103. outfile.writelines(self.lines)
  104. @abstractmethod
  105. def _close_graph(self) -> None:
  106. """Emit the lines needed to properly close the graph."""