05_uml_class_diagram.py

Download
python 296 lines 8.9 KB
  1"""
  2UML Class Diagram Generator (Text-Based)
  3
  4Generates a formatted text representation of a UML class diagram.
  5Define classes with attributes and methods (using +/-/# visibility),
  6then specify relationships (inheritance, composition, association).
  7
  8Example domain: e-commerce system with Product, Order, Customer, Payment.
  9
 10Run:
 11    python 05_uml_class_diagram.py
 12"""
 13
 14from dataclasses import dataclass, field
 15from typing import Literal
 16
 17# ---------------------------------------------------------------------------
 18# Data model
 19# ---------------------------------------------------------------------------
 20
 21Visibility = Literal["+", "-", "#"]
 22RelType = Literal["inheritance", "composition", "aggregation", "association", "dependency"]
 23
 24
 25@dataclass
 26class Attribute:
 27    name: str
 28    type_hint: str
 29    visibility: Visibility = "+"
 30
 31    def render(self) -> str:
 32        return f"  {self.visibility} {self.name}: {self.type_hint}"
 33
 34
 35@dataclass
 36class Method:
 37    name: str
 38    params: str = ""
 39    return_type: str = "None"
 40    visibility: Visibility = "+"
 41
 42    def render(self) -> str:
 43        return f"  {self.visibility} {self.name}({self.params}): {self.return_type}"
 44
 45
 46@dataclass
 47class UMLClass:
 48    name: str
 49    is_abstract: bool = False
 50    attributes: list[Attribute] = field(default_factory=list)
 51    methods: list[Method] = field(default_factory=list)
 52
 53    def render(self) -> list[str]:
 54        lines: list[str] = []
 55        width = max(
 56            len(self.name) + 4,
 57            max((len(a.render()) for a in self.attributes), default=0) + 2,
 58            max((len(m.render()) for m in self.methods), default=0) + 2,
 59            20,
 60        )
 61        bar = "+" + "-" * width + "+"
 62
 63        lines.append(bar)
 64        label = f"<<abstract>> {self.name}" if self.is_abstract else self.name
 65        lines.append(f"|{label.center(width)}|")
 66        lines.append(bar)
 67
 68        if self.attributes:
 69            for attr in self.attributes:
 70                row = attr.render()
 71                lines.append(f"|{row:<{width}}|")
 72        else:
 73            lines.append(f"|{'(no attributes)':<{width}}|")
 74
 75        lines.append(bar)
 76
 77        if self.methods:
 78            for method in self.methods:
 79                row = method.render()
 80                lines.append(f"|{row:<{width}}|")
 81        else:
 82            lines.append(f"|{'(no methods)':<{width}}|")
 83
 84        lines.append(bar)
 85        return lines
 86
 87
 88@dataclass
 89class Relationship:
 90    source: str
 91    target: str
 92    rel_type: RelType
 93    label: str = ""
 94    source_mult: str = ""
 95    target_mult: str = ""
 96
 97    _ARROWS: dict[str, str] = field(default_factory=lambda: {
 98        "inheritance":  "──────────▷",
 99        "composition":  "◆─────────",
100        "aggregation":  "◇─────────",
101        "association":  "──────────▶",
102        "dependency":   "- - - - - ▶",
103    })
104
105    def render(self) -> str:
106        arrow = self._ARROWS[self.rel_type]
107        parts = [self.source]
108        if self.source_mult:
109            parts.append(f"[{self.source_mult}]")
110        parts.append(arrow)
111        if self.target_mult:
112            parts.append(f"[{self.target_mult}]")
113        parts.append(self.target)
114        line = " ".join(parts)
115        if self.label:
116            line += f"  ({self.label})"
117        return line
118
119
120class UMLDiagram:
121    def __init__(self, title: str):
122        self.title = title
123        self.classes: list[UMLClass] = []
124        self.relationships: list[Relationship] = []
125
126    def add_class(self, cls: UMLClass) -> "UMLDiagram":
127        self.classes.append(cls)
128        return self
129
130    def add_relationship(self, rel: Relationship) -> "UMLDiagram":
131        self.relationships.append(rel)
132        return self
133
134    def render(self) -> str:
135        sections: list[str] = []
136
137        header = f" UML Class Diagram: {self.title} "
138        border = "=" * len(header)
139        sections.append(border)
140        sections.append(header)
141        sections.append(border)
142        sections.append("")
143
144        # Classes
145        sections.append("[ Classes ]")
146        sections.append("")
147        for cls in self.classes:
148            sections.extend(cls.render())
149            sections.append("")
150
151        # Relationships
152        if self.relationships:
153            sections.append("[ Relationships ]")
154            sections.append("")
155            legend = {
156                "inheritance":  "──────────▷  Inheritance (is-a)",
157                "composition":  "◆─────────   Composition (owns, lifecycle tied)",
158                "aggregation":  "◇─────────   Aggregation (has-a, independent lifecycle)",
159                "association":  "──────────▶  Association (uses)",
160                "dependency":   "- - - - - ▶  Dependency (depends on)",
161            }
162            seen_types: set[str] = set()
163            for rel in self.relationships:
164                sections.append(f"  {rel.render()}")
165                seen_types.add(rel.rel_type)
166            sections.append("")
167            sections.append("[ Legend ]")
168            for rtype, desc in legend.items():
169                if rtype in seen_types:
170                    sections.append(f"  {desc}")
171
172        return "\n".join(sections)
173
174
175# ---------------------------------------------------------------------------
176# E-commerce domain model
177# ---------------------------------------------------------------------------
178
179def build_ecommerce_diagram() -> UMLDiagram:
180    diagram = UMLDiagram("E-Commerce System")
181
182    # --- Product ---
183    product = UMLClass(
184        name="Product",
185        attributes=[
186            Attribute("id", "UUID", "-"),
187            Attribute("name", "str", "-"),
188            Attribute("price", "Decimal", "-"),
189            Attribute("stock_qty", "int", "-"),
190        ],
191        methods=[
192            Method("get_price", "", "Decimal"),
193            Method("is_in_stock", "", "bool"),
194            Method("apply_discount", "pct: float", "Decimal"),
195        ],
196    )
197
198    # --- Customer ---
199    customer = UMLClass(
200        name="Customer",
201        attributes=[
202            Attribute("id", "UUID", "-"),
203            Attribute("name", "str", "-"),
204            Attribute("email", "str", "-"),
205            Attribute("address", "Address", "#"),
206        ],
207        methods=[
208            Method("place_order", "items: list[OrderLine]", "Order"),
209            Method("get_order_history", "", "list[Order]"),
210        ],
211    )
212
213    # --- OrderLine ---
214    order_line = UMLClass(
215        name="OrderLine",
216        attributes=[
217            Attribute("product", "Product", "+"),
218            Attribute("quantity", "int", "+"),
219            Attribute("unit_price", "Decimal", "+"),
220        ],
221        methods=[
222            Method("subtotal", "", "Decimal"),
223        ],
224    )
225
226    # --- Order ---
227    order = UMLClass(
228        name="Order",
229        attributes=[
230            Attribute("id", "UUID", "-"),
231            Attribute("status", "OrderStatus", "-"),
232            Attribute("created_at", "datetime", "-"),
233            Attribute("lines", "list[OrderLine]", "-"),
234        ],
235        methods=[
236            Method("total", "", "Decimal"),
237            Method("add_line", "line: OrderLine", "None"),
238            Method("cancel", "", "None"),
239            Method("confirm", "", "None"),
240        ],
241    )
242
243    # --- Payment (abstract) ---
244    payment = UMLClass(
245        name="Payment",
246        is_abstract=True,
247        attributes=[
248            Attribute("amount", "Decimal", "#"),
249            Attribute("currency", "str", "#"),
250            Attribute("timestamp", "datetime", "#"),
251        ],
252        methods=[
253            Method("process", "", "bool"),
254            Method("refund", "", "bool"),
255        ],
256    )
257
258    # --- CreditCardPayment ---
259    cc_payment = UMLClass(
260        name="CreditCardPayment",
261        attributes=[
262            Attribute("card_token", "str", "-"),
263            Attribute("last_four", "str", "-"),
264        ],
265        methods=[
266            Method("process", "", "bool"),
267            Method("refund", "", "bool"),
268        ],
269    )
270
271    for cls in [product, customer, order_line, order, payment, cc_payment]:
272        diagram.add_class(cls)
273
274    # Relationships
275    diagram.add_relationship(Relationship("Customer", "Order",
276        "association", "places", "1", "0..*"))
277    diagram.add_relationship(Relationship("Order", "OrderLine",
278        "composition", "contains", "1", "1..*"))
279    diagram.add_relationship(Relationship("OrderLine", "Product",
280        "association", "references", "1..*", "1"))
281    diagram.add_relationship(Relationship("Order", "Payment",
282        "association", "paid by", "1", "0..1"))
283    diagram.add_relationship(Relationship("CreditCardPayment", "Payment",
284        "inheritance"))
285
286    return diagram
287
288
289# ---------------------------------------------------------------------------
290# Entry point
291# ---------------------------------------------------------------------------
292
293if __name__ == "__main__":
294    diagram = build_ecommerce_diagram()
295    print(diagram.render())