File size: 7,517 Bytes
11561b8 a495f34 11561b8 a495f34 11561b8 a495f34 11561b8 a495f34 11561b8 a495f34 11561b8 a495f34 11561b8 a495f34 11561b8 a495f34 11561b8 a495f34 11561b8 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 |
import graphviz
import json
from tempfile import NamedTemporaryFile
import os
from graph_generator_utils import add_nodes_and_edges # Reusing common utility
def generate_wbs_diagram(json_input: str, base_color: str) -> str: # base_color is now correctly used
"""
Generates a Work Breakdown Structure (WBS) Diagram from JSON input.
Args:
json_input (str): A JSON string describing the WBS structure.
It must follow the Expected JSON Format Example below.
base_color (str): The hexadecimal color string (e.g., '#19191a') for the base
color of the nodes, from which a gradient will be generated.
Returns:
str: The filepath to the generated PNG image file.
Expected JSON Format Example:
{
"project_title": "Software Development Project",
"phases": [
{
"id": "phase_prep",
"label": "1. Preparation",
"tasks": [
{"id": "task_vision", "label": "1.1. Identify Vision"},
{"id": "task_design", "label": "1.2. Design & Staffing"}
]
},
{
"id": "phase_plan",
"label": "2. Planning",
"tasks": [
{"id": "task_cost", "label": "2.1. Cost Analysis"},
{"id": "task_benefit", "label": "2.2. Benefit Analysis"},
{"id": "task_risk", "label": "2.3. Risk Assessment"}
]
}
]
}
"""
try:
if not json_input.strip():
return "Error: Empty input"
data = json.loads(json_input)
if 'project_title' not in data or 'phases' not in data:
raise ValueError("Missing required fields: project_title or phases")
dot = graphviz.Digraph(
name='WBSDiagram',
format='png',
graph_attr={
'rankdir': 'TB', # Top-to-Bottom hierarchy
'splines': 'ortho', # Straight lines
'bgcolor': 'white', # White background
'pad': '0.5', # Padding
'ranksep': '0.8', # Adjust vertical separation between ranks
'nodesep': '0.5' # Adjust horizontal separation between nodes
}
)
# This line was REMOVED to ensure the passed base_color is used: base_color = '#19191a'
# Project Title node (main node)
dot.node(
'project_root',
data['project_title'],
shape='box',
style='filled,rounded',
fillcolor=base_color, # Use the selected base color
fontcolor='white',
fontsize='18'
)
# Logic for color gradient within WBS specific nodes (Phases and Tasks)
# This ensures the gradient works correctly with the passed base_color
def get_gradient_color(depth, base_hex_color, lightening_factor=0.12):
base_r = int(base_hex_color[1:3], 16)
base_g = int(base_hex_color[3:5], 16)
base_b = int(base_hex_color[5:7], 16)
current_r = base_r + int((255 - base_r) * depth * lightening_factor)
current_g = base_g + int((255 - base_g) * depth * lightening_factor)
current_b = base_b + int((255 - base_b) * depth * lightening_factor)
return f'#{min(255, current_r):02x}{min(255, current_g):02x}{min(255, current_b):02x}'
def get_font_color_for_background(depth, base_hex_color, lightening_factor=0.12):
# Calculate brightness/lightness of the node color at this depth
# and return black/white for text accordingly
base_r = int(base_hex_color[1:3], 16)
base_g = int(base_hex_color[3:5], 16)
base_b = int(base_hex_color[5:7], 16)
current_r = base_r + (255 - base_r) * depth * lightening_factor
current_g = base_g + (255 - base_g) * depth * lightening_factor
current_b = base_b + (255 - base_b) * depth * lightening_factor
# Simple luminance check (ITU-R BT.709 coefficients)
luminance = (0.2126 * current_r + 0.7152 * current_g + 0.0722 * current_b) / 255
return 'white' if luminance < 0.5 else 'black'
current_depth = 1 # Depth for phases
for phase in data['phases']:
phase_id = phase.get('id')
phase_label = phase.get('label')
tasks = phase.get('tasks', [])
if not all([phase_id, phase_label]):
raise ValueError(f"Invalid phase: {phase}")
phase_fill_color = get_gradient_color(current_depth, base_color)
phase_font_color = get_font_color_for_background(current_depth, base_color)
dot.node(
phase_id,
phase_label,
shape='box',
style='filled,rounded',
fillcolor=phase_fill_color,
fontcolor=phase_font_color,
fontsize='14'
)
dot.edge('project_root', phase_id, color='#4a4a4a', arrowhead='none') # Connect to root
task_depth = current_depth + 1 # Depth for tasks
task_font_size = max(9, 14 - (task_depth * 2))
task_nodes_in_phase = []
for task in tasks:
task_id = task.get('id')
task_label = task.get('label')
if not all([task_id, task_label]):
raise ValueError(f"Invalid task: {task}")
task_fill_color = get_gradient_color(task_depth, base_color)
task_font_color = get_font_color_for_background(task_depth, base_color)
dot.node(
task_id,
task_label,
shape='box',
style='filled,rounded',
fillcolor=task_fill_color,
fontcolor=task_font_color,
fontsize=str(task_font_size)
)
dot.edge(phase_id, task_id, color='#4a4a4a', arrowhead='none') # Connect task to phase
task_nodes_in_phase.append(task_id)
# Use subgraph to enforce vertical alignment for tasks within a phase
if task_nodes_in_phase: # Only create subgraph if there are tasks
dot.subgraph(name=f'cluster_{phase_id}')
# You would typically add nodes to the subgraph if you want
# specific layout within it. For WBS columnar, just using rank attribute might suffice
# or ensuring proper `ranksep` and `nodesep` at graph level.
# The crucial part for columnar WBS is often setting nodes to the same rank
# or controlling the order. This is implicitly handled by `rankdir=TB`.
# The subgraphs themselves mostly define visual grouping (border) in Graphviz.
# Let's try to ensure vertical flow is emphasized without complex subgraphs that
# might interfere with the main layout engine. The primary connection handles rank.
# Save to temporary file
with NamedTemporaryFile(delete=False, suffix='.png') as tmp:
dot.render(tmp.name, format='png', cleanup=True)
return tmp.name + '.png'
except json.JSONDecodeError:
return "Error: Invalid JSON format"
except Exception as e:
return f"Error: {str(e)}"
|