IamMiracleAlex/data_export_util.py

## data_export_util.py
from typing import Union
import csv

from django.http import HttpResponse
from django.db.models import QuerySet

from rest_framework.response import Response
from xhtml2pdf import pisa
import pandas as pd


class ExportItem:
    def __init__(
        self,
        data: Union[QuerySet, list],
        caption: str = "Reports data",
        fields: list = None,
        headers: list = None,
    ) -> None:
        """
        This class creates and export reports to csv, html, pdf and xlsx
        Args:
            data: a list of dictionaries, or a django queryset
            caption: a string to caption report
            fields: the fields of the model
            headers: the headers for the export
        """
        self.caption = caption

        if not isinstance(data, (list, QuerySet)):
            raise ValueError("Data must be a list or queryset")

        self.data = data

        if len(data) == 0:
            raise ValueError("Data must not be empty")

        if isinstance(data, list):
            self.fields = fields if fields else self.data[0].keys()
        else:
            self.fields = fields if fields else data.values()[0].keys()

        self.headers = headers if headers else self.fields

    def export_to_html(self):
        """
        Writes the html string to a html file
        """

        # Save the HTML code
        path = "report.html"
        content = self.create_html_data()

        file_obj = open(path, "w")
        file_obj.write(content)
        file_obj.close()

        download_file = open(path)

        response = HttpResponse(content=download_file, content_type="text/html")
        response["Content-Disposition"] = f'attachment; filename="{path}"'
        return response

    def format_title(self, value) -> str:
        """
        Format the title in a good form,
                Converts first_name -> First Name
        """

        new_val = value.split("_")
        return " ".join([x.capitalize() for x in new_val])

    def create_html_data(self):
        """
        Creates html data for reports
        """

        # Start the page
        content = (
            """
        <html>
            <head>
            <title>"""
            + self.caption
            + """</title>
            </head>
            <body>
                <center>
            \n
        """
        )

        # Add content to the body
        content += "<table style='border:1px solid black; '>\n"
        content += (
            "<caption style='font-weight: bold; font-size: 10px; padding-top: 3px;' >"
            + self.caption
            + "</caption>\n"
        )
        content += "<tr>\n"

        for k in self.headers:
            content += (
                "<th style='padding-top: 3px;' >" + self.format_title(k) + "</th>"
            )
        content += "</tr>\n"

        content += "  <tr>\n"

        if isinstance(self.data, list):
            for row in self.data:
                for k in self.fields:
                    content += (
                        "<td style='padding-top: 3px;' >"
                        + str(row.get(k, ""))
                        + "</td>\n"
                    )
                content += "</tr>\n"
        else:
            for row in self.data:
                for k in self.fields:
                    content += (
                        "<td style='padding-top: 3px;' >"
                        + str(getattr(row, str(k), ""))
                        + "</td>\n"
                    )
                content += "</tr>\n"

        content += "\t</table>\n"

        # Close the body and end the file
        content += """
            </center>
            </body>
        </html>
        """
        return content

    def export_to_pdf(self, html_string=""):
        """
        exports html string to pdf
        args:
        :html_string: Optional. This can be any html string or one generated from django's `render_to_string()` function
        """

        content = html_string if html_string else self.create_html_data()
        response = HttpResponse(content_type="application/pdf")
        response["Content-Disposition"] = 'attachment; filename="report.pdf"'
        pisa_status = pisa.CreatePDF(content, dest=response)

        if pisa_status.err:
            return Response(dict(detail=f"Error: {pisa_status.err}"), status=400)

        return response

    def export_to_csv(self):
        """
        export items to csv
        """

        response = HttpResponse(content_type="text/csv")
        response["Content-Disposition"] = 'attachment; filename="report.csv"'

        writer = csv.writer(response)
        writer.writerow(self.headers)

        if isinstance(self.data, list):
            [
                writer.writerow([obj.get(field, "") for field in self.fields])
                for obj in self.data
            ]
        else:
            [
                writer.writerow([getattr(obj, str(field), "") for field in self.fields])
                for obj in self.data
            ]

        return response

    def export_to_xlsx(self):
        """
        export items to xlsx
        """

        path = "report.xlsx"
        if isinstance(self.data, list):
            df = pd.DataFrame(self.data, columns=self.fields)
        else:
            df = pd.DataFrame(self.data.values(), columns=self.fields)

        df.to_excel(path)
        file_obj = open(path, "rb").read()

        response = HttpResponse(content=file_obj, content_type="mimetype/submimetype")
        response["Content-Disposition"] = "attachment; filename=report.xlsx"

        return response

    def export(self, file_type):
        """
        export items to csv, html, xlsx or pdf
        """

        if not file_type in ["csv", "xlsx", "pdf", "html"]:
            raise TypeError("The acceptable file types are: csv, xlsx, pdf, html")

        exports = {
            "csv": self.export_to_csv,
            "xlsx": self.export_to_xlsx,
            "html": self.export_to_html,
            "pdf": self.export_to_pdf,
        }
        return exports.get(file_type)()


################## SAMPLE  USAGE  #####################
# Example 1
# Return a pdf response for an api view, using a queryset

class ExportView(APIView):
    permission_classes = [IsAuthenticated]
    queryset = User.objects.all()

    def post(self, request):
        export = ExportItem(self.queryset)
        return export.export_to_pdf()


# Example 2
# Get the csv content, using a list of dicts

data = [
    {"id": 1, "name": "Local", "plan": "premium"},
    {"id": 2, "name": "secong user", "plan": "premium"},
]

export = ExportItem(data)
response = export.export("csv")
print(response.content)
	from typing import Union
	import csv

	from django.http import HttpResponse
	from django.db.models import QuerySet

	from rest_framework.response import Response
	from xhtml2pdf import pisa
	import pandas as pd


	class ExportItem:
	def __init__(
	self,
	data: Union[QuerySet, list],
	caption: str = "Reports data",
	fields: list = None,
	headers: list = None,
	) -> None:
	"""
	This class creates and export reports to csv, html, pdf and xlsx
	Args:
	data: a list of dictionaries, or a django queryset
	caption: a string to caption report
	fields: the fields of the model
	headers: the headers for the export
	"""
	self.caption = caption

	if not isinstance(data, (list, QuerySet)):
	raise ValueError("Data must be a list or queryset")

	self.data = data

	if len(data) == 0:
	raise ValueError("Data must not be empty")

	if isinstance(data, list):
	self.fields = fields if fields else self.data[0].keys()
	else:
	self.fields = fields if fields else data.values()[0].keys()

	self.headers = headers if headers else self.fields

	def export_to_html(self):
	"""
	Writes the html string to a html file
	"""

	# Save the HTML code
	path = "report.html"
	content = self.create_html_data()

	file_obj = open(path, "w")
	file_obj.write(content)
	file_obj.close()

	download_file = open(path)

	response = HttpResponse(content=download_file, content_type="text/html")
	response["Content-Disposition"] = f'attachment; filename="{path}"'
	return response

	def format_title(self, value) -> str:
	"""
	Format the title in a good form,
	Converts first_name -> First Name
	"""

	new_val = value.split("_")
	return " ".join([x.capitalize() for x in new_val])

	def create_html_data(self):
	"""
	Creates html data for reports
	"""

	# Start the page
	content = (
	"""
	<html>
	<head>
	<title>"""
	+ self.caption
	+ """</title>
	</head>
	<body>
	<center>
	\n
	"""
	)

	# Add content to the body
	content += "<table style='border:1px solid black; '>\n"
	content += (
	"<caption style='font-weight: bold; font-size: 10px; padding-top: 3px;' >"
	+ self.caption
	+ "</caption>\n"
	)
	content += "<tr>\n"

	for k in self.headers:
	content += (
	"<th style='padding-top: 3px;' >" + self.format_title(k) + "</th>"
	)
	content += "</tr>\n"

	content += " <tr>\n"

	if isinstance(self.data, list):
	for row in self.data:
	for k in self.fields:
	content += (
	"<td style='padding-top: 3px;' >"
	+ str(row.get(k, ""))
	+ "</td>\n"
	)
	content += "</tr>\n"
	else:
	for row in self.data:
	for k in self.fields:
	content += (
	"<td style='padding-top: 3px;' >"
	+ str(getattr(row, str(k), ""))
	+ "</td>\n"
	)
	content += "</tr>\n"

	content += "\t</table>\n"

	# Close the body and end the file
	content += """
	</center>
	</body>
	</html>
	"""
	return content

	def export_to_pdf(self, html_string=""):
	"""
	exports html string to pdf
	args:
	:html_string: Optional. This can be any html string or one generated from django's `render_to_string()` function
	"""

	content = html_string if html_string else self.create_html_data()
	response = HttpResponse(content_type="application/pdf")
	response["Content-Disposition"] = 'attachment; filename="report.pdf"'
	pisa_status = pisa.CreatePDF(content, dest=response)

	if pisa_status.err:
	return Response(dict(detail=f"Error: {pisa_status.err}"), status=400)

	return response

	def export_to_csv(self):
	"""
	export items to csv
	"""

	response = HttpResponse(content_type="text/csv")
	response["Content-Disposition"] = 'attachment; filename="report.csv"'

	writer = csv.writer(response)
	writer.writerow(self.headers)

	if isinstance(self.data, list):
	[
	writer.writerow([obj.get(field, "") for field in self.fields])
	for obj in self.data
	]
	else:
	[
	writer.writerow([getattr(obj, str(field), "") for field in self.fields])
	for obj in self.data
	]

	return response

	def export_to_xlsx(self):
	"""
	export items to xlsx
	"""

	path = "report.xlsx"
	if isinstance(self.data, list):
	df = pd.DataFrame(self.data, columns=self.fields)
	else:
	df = pd.DataFrame(self.data.values(), columns=self.fields)

	df.to_excel(path)
	file_obj = open(path, "rb").read()

	response = HttpResponse(content=file_obj, content_type="mimetype/submimetype")
	response["Content-Disposition"] = "attachment; filename=report.xlsx"

	return response

	def export(self, file_type):
	"""
	export items to csv, html, xlsx or pdf
	"""

	if not file_type in ["csv", "xlsx", "pdf", "html"]:
	raise TypeError("The acceptable file types are: csv, xlsx, pdf, html")

	exports = {
	"csv": self.export_to_csv,
	"xlsx": self.export_to_xlsx,
	"html": self.export_to_html,
	"pdf": self.export_to_pdf,
	}
	return exports.get(file_type)()







	################## SAMPLE USAGE #####################
	# Example 1
	# Return a pdf response for an api view, using a queryset

	class ExportView(APIView):
	permission_classes = [IsAuthenticated]
	queryset = User.objects.all()

	def post(self, request):
	export = ExportItem(self.queryset)
	return export.export_to_pdf()


	# Example 2
	# Get the csv content, using a list of dicts

	data = [
	{"id": 1, "name": "Local", "plan": "premium"},
	{"id": 2, "name": "secong user", "plan": "premium"},
	]

	export = ExportItem(data)
	response = export.export("csv")
	print(response.content)