from io import StringIO, BytesIO
import pandas as pd


class DfToStream:
    def __init__(self, *args, index=False, **kwargs):
        self.dfs = args
        self.index = index
        self.writer = None
        self.output = None

    def __enter__(self):
        self.output = BytesIO()
        self.writer = pd.ExcelWriter(self.output, engine='xlsxwriter')
        return self

    def __exit__(self, exc_type, exc_val, exc_tb):
        self.writer.close()
        self.output.seek(0)

    def to_stream(self):
        for item in self.dfs:
            df = item[0]
            sheet_name = item[1]
            df.to_excel(self.writer, encoding='utf-8', sheet_name=sheet_name, index=self.index)
        return self.output