tasks/code_generation/python_algorithms/loader.py

"""
Python Algorithms Dataset Loader
Creates test data for Python algorithm evaluation.
"""

import os
import json
from typing import List, Dict, Any


def create_test_data(data_path: str = "python_algorithms_test_data.json"):
    """Create test data for Python algorithm evaluation."""
    
    test_data = {
        "sort_list": {
            "input": [64, 34, 25, 12, 22, 11, 90],
            "expected_output": [11, 12, 22, 25, 34, 64, 90]
        },
        "binary_search": {
            "input": {"arr": [1, 3, 5, 7, 9, 11, 13, 15], "target": 7},
            "expected_output": 3
        },
        "fibonacci": {
            "input": 10,
            "expected_output": 55
        },
        "two_sum": {
            "input": {"nums": [2, 7, 11, 15], "target": 9},
            "expected_output": [0, 1]
        },
        "merge_sort": {
            "input": [38, 27, 43, 3, 9, 82, 10],
            "expected_output": [3, 9, 10, 27, 38, 43, 82]
        },
        "bank_account": {
            "input": {"operations": ["deposit", "withdraw", "deposit"], "amounts": [100, 50, 25]},
            "expected_output": 75
        }
    }
    
    with open(data_path, 'w') as f:
        json.dump(test_data, f, indent=2)
    
    print(f"Created test data: {data_path}")
    return data_path


def load_test_data(data_path: str = "python_algorithms_test_data.json") -> Dict[str, Any]:
    """Load test data for evaluation."""
    if not os.path.exists(data_path):
        create_test_data(data_path)
    
    with open(data_path, 'r') as f:
        return json.load(f)


if __name__ == "__main__":
    create_test_data()