data_generator.py

import pandas as pd
from parser import parse_prompt
from generators import india, usa, uk, uae, south_africa

country_module_map = {
    "india": india,
    "usa": usa,
    "uk": uk,
    "uae": uae,
    "south africa": south_africa
}

def generate_test_data(prompt):
    parsed = parse_prompt(prompt)
    country = parsed['country']
    fields = parsed['fields']
    count = parsed['count']

    generator = country_module_map.get(country.lower())
    if not generator:
        raise ValueError(f"Unsupported country: {country}")

    return generator.generate(fields, count)