Skip to content

Instantly share code, notes, and snippets.

@IamMiracleAlex

IamMiracleAlex/data_export_util.py

Created February 27, 2023 11:55
Show Gist options
  • Save IamMiracleAlex/72cc6268412c6f5dbfd4887330c95f60 to your computer and use it in GitHub Desktop.
Save IamMiracleAlex/72cc6268412c6f5dbfd4887330c95f60 to your computer and use it in GitHub Desktop.
Download ZIP
This class creates and export reports to csv, html, pdf and xlsx
Raw
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)
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment