Untitled

from graphviz import Digraph

def create_railway_erd():
    erd = Digraph('Infrastructure Railway ERD', filename='railway_erd', format='png')
    erd.attr(rankdir='LR', size='10')
    
    # Entities
    entities = {
        'Railway Station': ['Name', 'Location'],
        'Train': ['Train Number', 'Capacity'],
        'Track': ['Track ID', 'Length'],
        'Passenger': ['Passenger ID', 'Name', 'Age'],
        'Ticket': ['Ticket ID', 'Price'],
        'Operator': ['Operator ID', 'Name'],
        'Maintenance': ['Maintenance ID', 'Schedule']
    }
    
    for entity, attributes in entities.items():
        erd.node(entity, shape='box', style='filled', fillcolor='lightblue')
        for attr in attributes:
            erd.node(attr, shape='ellipse', style='filled', fillcolor='lightcoral')
            erd.edge(entity, attr)
    
    # Relationships
    relationships = {
        ('Passenger', 'Ticket'): 'Purchases',
        ('Passenger', 'Train'): 'Boards',
        ('Train', 'Railway Station'): 'Stops At',
        ('Train', 'Track'): 'Travels On',
        ('Operator', 'Train'): 'Operates',
        ('Railway Station', 'Track'): 'Connected To',
        ('Train', 'Maintenance'): 'Requires'
    }
    
    for (entity1, entity2), relation in relationships.items():
        erd.edge(entity1, entity2, label=relation, fontsize='10', color='darkblue')
    
    return erd

erd_diagram = create_railway_erd()
erd_diagram.view()