UML Diagram maker overhaul

pyproject.toml

@@ -9,6 +9,8 @@ dependencies = [
   "pytz",
   "rich",
   "sh",
+  "graphviz",
+  "pylint",
 ]
 description = "Python utils to simplify DAQ development."
 name = "daqpytools"
@@ -41,7 +43,7 @@ docs = [
   "mkdocs-section-index",
   "mkdocs-exclude",
 ]
-test = ["pytest", "pytest-mypy", "pytest-cov", "types-pytz"]
+test = ["pytest", "pylint", "graphviz", "pytest-mypy", "pytest-cov", "types-pytz"]
 
 [project.scripts]
 daqpytools-generate-uml = "daqpytools.apps.generate_uml:main"

src/daqpytools/apps/generate_uml.py

@@ -20,11 +20,23 @@
 
 import click
 
+from daqpytools.uml.directory_coloring import (
+    assign_directory_colors,
+    build_directory_color_legend_dot,
+    inject_directory_colors,
+)
 from daqpytools.uml.dot_parsing import patch_dot
+from daqpytools.uml.github_links import inject_node_links, resolve_git_ref
 from daqpytools.uml.render import render_dot
 from daqpytools.uml.split_diagram import split_dot_file
 from daqpytools.uml.style_pyreverse import run_pyreverse
-from daqpytools.uml.utils import CONTEXT_SETTINGS, load_style_config, vprint
+from daqpytools.uml.utils import (
+    CONTEXT_SETTINGS,
+    load_color_theme_config,
+    load_link_config,
+    load_style_config,
+    vprint,
+)
 
 
 def validate_output_directory(
@@ -60,15 +72,50 @@ def resolve_output_directory(
     return cwd, resolved_output
 
 
-def style_dot_file(dot_path: Path, style: dict[str, str], concise: bool) -> Path:
+def style_dot_file(
+    dot_path: Path,
+    style: dict[str, str],
+    concise: bool,
+    color_context: tuple[str, list[str]] | None = None,
+    link_context: tuple[Path, str, str, str, bool] | None = None,
+) -> Path:
     """Patch a raw dot file and write the styled version next to it."""
     original = dot_path.read_text(encoding="utf-8")
     patched = patch_dot(original, style=style, concise=concise)
+    if color_context is not None:
+        package_name, palette = color_context
+        directory_colors = assign_directory_colors(patched, package_name, palette)
+        patched = inject_directory_colors(patched, package_name, directory_colors)
+    if link_context is not None:
+        cwd, org, repo, ref, link_line_number = link_context
+        patched = inject_node_links(patched, cwd, org, repo, ref, link_line_number)
     patched_path = dot_path.with_name(f"{dot_path.stem}_styled.dot")
     patched_path.write_text(patched, encoding="utf-8")
     return patched_path
 
 
+def write_directory_color_legend(
+    styled_dot_files: list[Path],
+    output_dir: Path,
+    color_context: tuple[str, list[str]],
+) -> Path | None:
+    """Write a standalone directory color legend dot file."""
+    package_name, palette = color_context
+    directory_colors: dict[str, str] = {}
+    for styled_dot in styled_dot_files:
+        dot_src = styled_dot.read_text(encoding="utf-8")
+        directory_colors.update(assign_directory_colors(dot_src, package_name, palette))
+
+    if not directory_colors:
+        return None
+
+    legend_dot = output_dir / "directory_color_legend.dot"
+    legend_dot.write_text(
+        build_directory_color_legend_dot(directory_colors), encoding="utf-8"
+    )
+    return legend_dot
+
+
 @click.command(context_settings=CONTEXT_SETTINGS)
 @click.argument("targets", nargs=-1, required=False)
 @click.option(
@@ -136,6 +183,21 @@ def style_dot_file(dot_path: Path, style: dict[str, str], concise: bool) -> Path
     default=None,
     help="Path to YAML style config file for the UML renderer.",
 )
+@click.option(
+    "--generate-linkable",
+    is_flag=True,
+    help="Add GitHub source links to nodes and render a linked SVG.",
+)
+@click.option(
+    "--link-line-number",
+    is_flag=True,
+    help="With --generate-linkable, link directly to a class's definition line.",
+)
+@click.option(
+    "--color-by-directory",
+    is_flag=True,
+    help="Color class/package boxes by their top-level source directory.",
+)
 def main(
     targets: tuple[str, ...],
     directory: Path | None,
@@ -148,6 +210,9 @@ def main(
     classes: tuple[str, ...],
     verbose: bool,
     style_config: Path | None,
+    generate_linkable: bool,
+    link_line_number: bool,
+    color_by_directory: bool,
 ) -> None:
     """Generate styled UML class diagrams from Python code."""
     pyreverse_args = build_pyreverse_args(targets, package, classes)
@@ -158,6 +223,23 @@ def main(
     cwd, resolved_output_dir = resolve_output_directory(directory, output_directory)
     style = load_style_config(style_config)
     render_format = None if output_format.lower() == "none" else output_format.lower()
+    package_name = (targets[0] if targets else package[0]).split(".", maxsplit=1)[0]
+
+    link_context = None
+    if generate_linkable:
+        link_config = load_link_config()
+        repo = package_name
+        ref = resolve_git_ref(cwd, link_config["default_ref"])
+        link_context = (cwd, link_config["github_org"], repo, ref, link_line_number)
+        vprint(
+            verbose,
+            f"[generate_uml] Linking nodes to {link_config['github_org']}/{repo}@{ref}",
+        )
+
+    color_context = None
+    if color_by_directory:
+        color_context = (package_name, load_color_theme_config())
+        vprint(verbose, "[generate_uml] Coloring nodes by source directory")
 
     vprint(verbose, f"[generate_uml] Running pyreverse in {cwd}")
     dot_files = run_pyreverse(
@@ -167,7 +249,15 @@ def main(
     styled_dot_files: list[Path] = []
     for dot_path in dot_files:
         vprint(verbose, f"[generate_uml] Styling {dot_path.name}")
-        styled_dot_files.append(style_dot_file(dot_path, style=style, concise=concise))
+        styled_dot_files.append(
+            style_dot_file(
+                dot_path,
+                style=style,
+                concise=concise,
+                color_context=color_context,
+                link_context=link_context,
+            )
+        )
 
     split_dot_files: list[Path] = []
     if split:
@@ -187,6 +277,14 @@ def main(
                 verbose, f"[generate_uml] Split diagrams written to {split_output_dir}"
             )
 
+    legend_dot = None
+    if color_context is not None:
+        legend_dot = write_directory_color_legend(
+            styled_dot_files, resolved_output_dir, color_context
+        )
+        if legend_dot is not None:
+            vprint(verbose, f"[generate_uml] Written color legend: {legend_dot}")
+
     if render_format is not None:
         for split_dot in split_dot_files:
             img_path = render_dot(
@@ -199,9 +297,35 @@ def main(
                 styled_dot, resolved_output_dir, fmt=render_format, verbose=verbose
             )
             vprint(verbose, f"[generate_uml] Written: {img_path}")
+
+        if legend_dot is not None:
+            img_path = render_dot(
+                legend_dot, resolved_output_dir, fmt=render_format, verbose=verbose
+            )
+            vprint(verbose, f"[generate_uml] Written color legend: {img_path}")
     else:
         vprint(verbose, "[generate_uml] Skipping rendering (--format none)")
 
+    # Linkable output must be SVG; only render it if it wasn't already produced above.
+    if generate_linkable and render_format != "svg":
+        for split_dot in split_dot_files:
+            img_path = render_dot(
+                split_dot, split_dot.parent, fmt="svg", verbose=verbose
+            )
+            vprint(verbose, f"[generate_uml] Written linkable SVG: {img_path}")
+
+        for styled_dot in styled_dot_files:
+            img_path = render_dot(
+                styled_dot, resolved_output_dir, fmt="svg", verbose=verbose
+            )
+            vprint(verbose, f"[generate_uml] Written linkable SVG: {img_path}")
+
+        if legend_dot is not None:
+            img_path = render_dot(
+                legend_dot, resolved_output_dir, fmt="svg", verbose=verbose
+            )
+            vprint(verbose, f"[generate_uml] Written color legend SVG: {img_path}")
+
     vprint(verbose, "[generate_uml] Complete")
     vprint(verbose, f"[generate_uml] Output directory: {resolved_output_dir.resolve()}")
 

src/daqpytools/uml/directory_coloring.py

@@ -0,0 +1,83 @@
+"""Color UML nodes by their top-level source directory."""
+
+import re
+from collections.abc import Sequence
+
+from daqpytools.uml.dot_parsing import append_bracket_attrs
+from daqpytools.uml.github_links import NODE_LINE_PATTERN
+
+LEGEND_NODE_NAME = "directory_color_legend"
+
+
+def get_directory_name(node_id: str, package_name: str) -> str | None:
+    """Return the top-level package directory represented by ``node_id``."""
+    segments = node_id.split(".")
+    if segments and segments[0] == "src":
+        segments = segments[1:]
+    if len(segments) < 2 or segments[0] != package_name:
+        return None
+    return segments[1]
+
+
+def assign_directory_colors(
+    dot_src: str, package_name: str, palette: Sequence[str]
+) -> dict[str, str]:
+    """Assign palette colors to directories in first-seen order."""
+    colors = {}
+    if not palette:
+        return colors
+
+    for line in dot_src.splitlines():
+        match = NODE_LINE_PATTERN.match(line.strip())
+        if not match:
+            continue
+        directory = get_directory_name(match.group(1), package_name)
+        if directory is not None and directory not in colors:
+            colors[directory] = palette[len(colors) % len(palette)]
+    return colors
+
+
+def inject_directory_colors(
+    dot_src: str, package_name: str, directory_colors: dict[str, str]
+) -> str:
+    """Apply per-directory fill colors to UML node lines."""
+    out = []
+    for line in dot_src.splitlines():
+        match = NODE_LINE_PATTERN.match(line.strip())
+        if match:
+            directory = get_directory_name(match.group(1), package_name)
+            color = directory_colors.get(directory or "")
+            if color is not None:
+                line = _replace_or_append_fillcolor(line, color)
+        out.append(line)
+    return "\n".join(out)
+
+
+def build_directory_color_legend_dot(directory_colors: dict[str, str]) -> str:
+    """Build a standalone Graphviz dot file for directory colors."""
+    rows = [
+        f'<TR><TD BGCOLOR="{color}">{directory}</TD></TR>'
+        for directory, color in directory_colors.items()
+    ]
+    return "\n".join(
+        [
+            'digraph "directory_color_legend" {',
+            '  graph [bgcolor="white" rankdir="TB" margin="0.1"];',
+            f'  "{LEGEND_NODE_NAME}" [',
+        "  shape=plain,",
+        "  label=<",
+        '    <TABLE BORDER="0" CELLBORDER="1" CELLSPACING="0" CELLPADDING="6">',
+        '      <TR><TD BGCOLOR="white"><B>Directory</B></TD></TR>',
+        *(f"      {row}" for row in rows),
+        "    </TABLE>",
+        "  >",
+            "  ];",
+            "}",
+        ]
+    )
+
+
+def _replace_or_append_fillcolor(line: str, color: str) -> str:
+    if re.search(r'\bfillcolor\s*=\s*"[^"]*"', line):
+        return re.sub(r'\bfillcolor\s*=\s*"[^"]*"', f'fillcolor="{color}"', line)
+    return append_bracket_attrs(line, f'fillcolor="{color}"')

src/daqpytools/uml/dot_parsing.py

@@ -46,23 +46,38 @@ def patch_dot(dot_src: str, style: dict[str, str], concise: bool = False) -> str
     )
 
 
+def append_bracket_attrs(line: str, new_attrs: str) -> str:
+    """Append attrs into a line's outermost ``[...]`` attribute list.
+
+    Uses balanced-bracket matching rather than a naive first-match regex,
+    since HTML labels (e.g. ``set[HandlerType]``) can contain literal
+    bracket pairs before the real attribute list closes.
+    """
+    start = line.find("[")
+    if start == -1:
+        return re.sub(r";?\s*$", f" [{new_attrs}];", line.rstrip())
+
+    depth = 0
+    end = None
+    for i in range(start, len(line)):
+        if line[i] == "[":
+            depth += 1
+        elif line[i] == "]":
+            depth -= 1
+            if depth == 0:
+                end = i
+                break
+
+    if end is None:
+        return re.sub(r";?\s*$", f" [{new_attrs}];", line.rstrip())
+
+    inner = line[start + 1 : end].strip().rstrip(",")
+    prefix = f"{inner}, " if inner else ""
+    return line[:start] + "[" + prefix + new_attrs + "]" + line[end + 1 :]
+
+
 def _append_edge_attrs(line: str, new_attrs: str) -> str:
-    if "[" in line:
-        return re.sub(
-            r"\[([^\]]*)\]",
-            lambda match: (
-                "["
-                + (
-                    match.group(1).strip().rstrip(",") + ", "
-                    if match.group(1).strip()
-                    else ""
-                )
-                + new_attrs
-                + "]"
-            ),
-            line,
-        )
-    return re.sub(r";?\s*$", f" [{new_attrs}];", line.rstrip())
+    return append_bracket_attrs(line, new_attrs)
 
 
 def fix_edges(dot_src: str, style: dict[str, str]) -> str: