Untitled

import pandas as pd
import pytest as pytest


def get_data_frame():
    d = {'col1': [i for i in range(1, 1000)], 'col2': [i for i in range(1, 1000)], 'col3': [i for i in range(1, 1000)]}
    return pd.DataFrame(data=d)


def get_formula():
    return "=SUM(D10,D11)"


def v1():
    writer = pd.ExcelWriter("v1.xlsx", engine="openpyxl")
    df = get_data_frame()
    df.to_excel(
        writer,
        startrow=0,
        startcol=0,
        index=False,
    )
    for row in writer.book.worksheets[0].iter_rows(min_row=0, max_row=0, min_col=1):
        for cell in row:
            col = cell.column
            cell_value = cell.value
            if cell_value == "col1":
                for row in writer.book.worksheets[0].iter_rows(min_row=0,
                                                               max_row=writer.book.worksheets[0].max_row,
                                                               min_col=col, max_col=col):
                    for cell in row:
                        cell.value = get_formula()
            if cell_value == "col2":
                for row in writer.book.worksheets[0].iter_rows(min_row=0,
                                                               max_row=writer.book.worksheets[0].max_row,
                                                               min_col=col, max_col=col):
                    for cell in row:
                        cell.value = get_formula()
            if cell_value == "col3":
                for row in writer.book.worksheets[0].iter_rows(min_row=0,
                                                               max_row=writer.book.worksheets[0].max_row,
                                                               min_col=col, max_col=col):
                    for cell in row:
                        cell.value = get_formula()
    writer.close()


def v2():
    writer = pd.ExcelWriter("v2.xlsx", engine="openpyxl")
    df = get_data_frame()
    df["col1"] = get_formula()
    df["col2"] = get_formula()
    df["col3"] = get_formula()
    df.to_excel(
        writer,
        startrow=0,
        startcol=0,
        index=False,
    )
    writer.close()


@pytest.mark.benchmark(
    group="af-compare",
    min_rounds=20,
)
def test_generate_excel_af_v1(benchmark):
    benchmark(v1)


@pytest.mark.benchmark(
    group="af-compare",
    min_rounds=20,
)
def test_generate_excel_af_v2(benchmark):
    benchmark(v2)