🔍 Code Extractor

A class that generates various types of diagrams including flowcharts, process flows, and network diagrams using matplotlib and networkx, returning base64-encoded PNG images.

Purpose

DiagramGenerator provides a unified interface for creating different diagram types from specifications. It takes a GraphicSpec object containing diagram parameters (type, steps, nodes, edges, direction) and generates visual diagrams as base64-encoded images. The class supports flowcharts (vertical/horizontal), process diagrams, and network diagrams, with automatic error handling and fallback to error diagrams when generation fails.

Source Code

class DiagramGenerator:
    """Generates diagrams like flowcharts, process flows, organizational charts"""
    
    def generate_diagram(self, spec: GraphicSpec) -> GraphicSpec:
        """Generate a diagram based on specification"""
        params = spec.parameters
        diagram_type = params.get('style', 'flowchart').lower()
        
        try:
            if diagram_type == 'flowchart':
                return self._create_flowchart(spec)
            elif diagram_type == 'process':
                return self._create_process_diagram(spec)
            elif diagram_type == 'network':
                return self._create_network_diagram(spec)
            else:
                return self._create_flowchart(spec)  # Default to flowchart
                
        except Exception as e:
            print(f"Error generating diagram: {e}")
            return self._create_error_diagram(spec, f"Diagram generation failed: {e}")
    
    def _create_flowchart(self, spec: GraphicSpec) -> GraphicSpec:
        """Create a flowchart diagram"""
        params = spec.parameters
        steps = params.get('steps', ['Start', 'Process', 'End'])
        direction = params.get('direction', 'vertical')
        
        fig, ax = plt.subplots(figsize=(10, 8))
        
        if direction == 'horizontal':
            self._draw_horizontal_flowchart(ax, steps)
        else:
            self._draw_vertical_flowchart(ax, steps)
        
        ax.set_xlim(-1, len(steps))
        ax.set_ylim(-1, len(steps))
        ax.axis('off')
        ax.set_title(spec.description, fontsize=14, fontweight='bold', pad=20)
        
        plt.tight_layout()
        spec.image_data = self._fig_to_base64(fig)
        spec.width = 1000
        spec.height = 800
        plt.close(fig)
        return spec
    
    def _draw_vertical_flowchart(self, ax, steps):
        """Draw a vertical flowchart"""
        y_positions = np.linspace(len(steps)-1, 0, len(steps))
        
        for i, (step, y) in enumerate(zip(steps, y_positions)):
            # Draw box
            box = FancyBboxPatch((0.1, y-0.2), 1.8, 0.4,
                               boxstyle="round,pad=0.05",
                               facecolor='white', edgecolor='black', linewidth=2)
            ax.add_patch(box)
            
            # Add text
            ax.text(1, y, step, ha='center', va='center', fontsize=11, fontweight='bold')
            
            # Draw arrow to next step
            if i < len(steps) - 1:
                ax.annotate('', xy=(1, y_positions[i+1]+0.2), xytext=(1, y-0.2),
                           arrowprops=dict(arrowstyle='->', lw=2, color='black'))
    
    def _draw_horizontal_flowchart(self, ax, steps):
        """Draw a horizontal flowchart"""
        x_positions = np.linspace(0, len(steps)-1, len(steps))
        
        for i, (step, x) in enumerate(zip(steps, x_positions)):
            # Draw box
            box = FancyBboxPatch((x-0.4, 0.3), 0.8, 0.4,
                               boxstyle="round,pad=0.05",
                               facecolor='white', edgecolor='black', linewidth=2)
            ax.add_patch(box)
            
            # Add text
            ax.text(x, 0.5, step, ha='center', va='center', fontsize=11, fontweight='bold')
            
            # Draw arrow to next step
            if i < len(steps) - 1:
                ax.annotate('', xy=(x_positions[i+1]-0.4, 0.5), xytext=(x+0.4, 0.5),
                           arrowprops=dict(arrowstyle='->', lw=2, color='black'))
    
    def _create_process_diagram(self, spec: GraphicSpec) -> GraphicSpec:
        """Create a process diagram"""
        # Similar to flowchart but with different styling
        return self._create_flowchart(spec)
    
    def _create_network_diagram(self, spec: GraphicSpec) -> GraphicSpec:
        """Create a network diagram using networkx"""
        params = spec.parameters
        nodes = params.get('nodes', ['A', 'B', 'C'])
        edges = params.get('edges', [('A', 'B'), ('B', 'C')])
        
        fig, ax = plt.subplots(figsize=(8, 8))
        
        # Create network graph
        G = nx.Graph()
        G.add_nodes_from(nodes)
        G.add_edges_from(edges)
        
        # Position nodes
        pos = nx.spring_layout(G, seed=42)
        
        # Draw network
        nx.draw(G, pos, ax=ax, with_labels=True, node_color='white',
               node_size=1500, font_size=12, font_weight='bold',
               edge_color='black', linewidths=2)
        
        ax.set_title(spec.description, fontsize=14, fontweight='bold')
        plt.tight_layout()
        
        spec.image_data = self._fig_to_base64(fig)
        spec.width = 800
        spec.height = 800
        plt.close(fig)
        return spec
    
    def _fig_to_base64(self, fig) -> str:
        """Convert matplotlib figure to base64 encoded image"""
        buffer = io.BytesIO()
        fig.savefig(buffer, format='png', dpi=150, bbox_inches='tight',
                   facecolor='white', edgecolor='none')
        buffer.seek(0)
        image_data = buffer.getvalue()
        buffer.close()
        
        return base64.b64encode(image_data).decode('utf-8')
    
    def _create_error_diagram(self, spec: GraphicSpec, error_msg: str) -> GraphicSpec:
        """Create an error placeholder diagram"""
        fig, ax = plt.subplots(figsize=(6, 4))
        ax.text(0.5, 0.5, f"Error: {error_msg}", ha='center', va='center',
               transform=ax.transAxes, fontsize=12,
               bbox=dict(boxstyle="round,pad=0.3", facecolor="lightgray"))
        ax.set_xlim(0, 1)
        ax.set_ylim(0, 1)
        ax.axis('off')
        
        spec.image_data = self._fig_to_base64(fig)
        spec.width = 600
        spec.height = 400
        plt.close(fig)
        return spec

Parameters

Parameter Details

No constructor parameters: The class has no __init__ method defined, so it uses the default constructor with no parameters. Instantiation is straightforward with no required arguments.

Return Value

The main method generate_diagram() returns a GraphicSpec object with populated image_data (base64-encoded PNG), width, and height attributes. All internal methods also return modified GraphicSpec objects. The image_data field contains a base64-encoded string that can be decoded to display or save the generated diagram.

Class Interface

Methods

Dependencies

• matplotlib
• numpy
• networkx
• io
• base64

Required Imports

import matplotlib.pyplot as plt
from matplotlib.patches import FancyBboxPatch
import numpy as np
import networkx as nx
import io
import base64

Usage Example

from dataclasses import dataclass
from typing import Dict, Any

@dataclass
class GraphicSpec:
    description: str
    parameters: Dict[str, Any]
    image_data: str = ''
    width: int = 0
    height: int = 0

# Instantiate the generator
generator = DiagramGenerator()

# Create a flowchart specification
flowchart_spec = GraphicSpec(
    description='User Login Process',
    parameters={
        'style': 'flowchart',
        'steps': ['Start', 'Enter Credentials', 'Validate', 'Success'],
        'direction': 'vertical'
    }
)

# Generate the diagram
result = generator.generate_diagram(flowchart_spec)
print(f'Generated image size: {result.width}x{result.height}')
print(f'Image data length: {len(result.image_data)}')

# Create a network diagram
network_spec = GraphicSpec(
    description='Network Topology',
    parameters={
        'style': 'network',
        'nodes': ['Server', 'Router', 'Client1', 'Client2'],
        'edges': [('Server', 'Router'), ('Router', 'Client1'), ('Router', 'Client2')]
    }
)

network_result = generator.generate_diagram(network_spec)

# Decode and save the image
import base64
with open('diagram.png', 'wb') as f:
    f.write(base64.b64decode(network_result.image_data))

Best Practices

• Always ensure the GraphicSpec object has a valid 'parameters' dictionary before calling generate_diagram()
• The class automatically closes matplotlib figures to prevent memory leaks, but be aware of memory usage when generating many diagrams
• Default values are provided for missing parameters: style defaults to 'flowchart', steps default to ['Start', 'Process', 'End'], direction defaults to 'vertical'
• Error handling is built-in: if diagram generation fails, an error diagram is returned instead of raising an exception
• The class is stateless and thread-safe - each method call is independent and doesn't modify class state
• For network diagrams, ensure edges reference nodes that exist in the nodes list to avoid networkx errors
• The generated base64 image data can be directly embedded in HTML img tags using data:image/png;base64, prefix
• DPI is set to 150 for good quality; modify _fig_to_base64 if different resolution is needed
• The class modifies the input GraphicSpec object in-place, setting image_data, width, and height attributes

Tags

Similar Components

AI-powered semantic similarity - components with related functionality:

• class GraphicsGenerator 71.2% similar

  GraphicsGenerator is a coordinator class that orchestrates the generation of different types of graphics (charts, diagrams, illustrations, and sketches) by delegating to specialized generator classes.

  From: /tf/active/vicechatdev/e-ink-llm/graphics_generator.py

• class IllustrationGenerator 69.8% similar

  A class that generates educational illustrations and technical drawings for mathematical and scientific concepts using matplotlib and programmatic rendering.

  From: /tf/active/vicechatdev/e-ink-llm/graphics_generator.py

• class ChartGenerator 64.8% similar

  A class that generates various types of charts (bar, line, pie, scatter) optimized for e-ink displays with high contrast and clear visibility.

  From: /tf/active/vicechatdev/e-ink-llm/graphics_generator.py

• function demo_graphics_generation 56.7% similar

  Demonstrates the generation of three types of graphics (bar chart, process diagram, and mathematical illustration) using the GraphicsGenerator class with e-ink optimized styling.

  From: /tf/active/vicechatdev/e-ink-llm/demo_hybrid_mode.py

• function generate_diagram_data 54.5% similar

  Transforms Neo4j schema information into a structured format suitable for graph visualization, creating separate node and edge data structures.

  From: /tf/active/vicechatdev/neo4j_schema_report.py