Iman_Asset_Lite/asset_lite/api/data_import_api.py

122 lines
3.2 KiB
Python

import csv
import os
from io import StringIO
import frappe
from frappe import _
from frappe.utils.csvutils import read_csv_content
from frappe.utils.file_manager import save_file
from frappe.utils.xlsxutils import (
make_xlsx,
read_xls_file_from_attached_file,
read_xlsx_file_from_attached_file,
)
def _trim_cell(value):
if isinstance(value, str):
trimmed = value.strip()
return trimmed if trimmed else None
return value
def _trim_rows(rows):
trimmed_cells = 0
out = []
for row in rows:
new_row = []
for cell in row:
if isinstance(cell, str):
original = cell
trimmed = _trim_cell(cell)
if original != (trimmed or ""):
trimmed_cells += 1
new_row.append(trimmed)
else:
new_row.append(cell)
out.append(new_row)
return out, trimmed_cells
def _read_rows_from_content(content, extension):
extension = (extension or "").lower()
if extension == "csv":
return read_csv_content(content)
if extension == "xlsx":
return read_xlsx_file_from_attached_file(fcontent=content)
if extension == "xls":
return read_xls_file_from_attached_file(content)
frappe.throw(_("Import file must be .csv, .xlsx, or .xls"))
def _write_trimmed_content(rows, extension):
extension = (extension or "").lower()
if extension == "csv":
buffer = StringIO()
writer = csv.writer(buffer)
for row in rows:
writer.writerow(["" if cell is None else cell for cell in row])
return buffer.getvalue().encode("utf-8"), "csv"
if extension == "xls":
extension = "xlsx"
if extension == "xlsx":
return make_xlsx(rows, "Import").getvalue(), "xlsx"
frappe.throw(_("Import file must be .csv, .xlsx, or .xls"))
def _trimmed_filename(original_name, output_extension):
base, _ = os.path.splitext(original_name or "import")
return f"{base}_trimmed.{output_extension}"
@frappe.whitelist()
def trim_and_upload_import_file(data_import):
doc = frappe.get_doc("Data Import", data_import)
doc.check_permission("write")
if not doc.import_file:
frappe.throw(_("Attach an import file first."))
if doc.google_sheets_url:
frappe.throw(_("Trim and Upload is only available for attached files, not Google Sheets URLs."))
file_name = frappe.db.get_value("File", {"file_url": doc.import_file}, "name")
if not file_name:
frappe.throw(_("Import file not found in File attachments."))
source_file = frappe.get_doc("File", file_name)
content = source_file.get_content()
extension = os.path.splitext(source_file.file_name or doc.import_file)[1][1:].lower()
rows = _read_rows_from_content(content, extension)
if not rows:
frappe.throw(_("Import file is empty."))
trimmed_rows, trimmed_cells = _trim_rows(rows)
output_content, output_extension = _write_trimmed_content(trimmed_rows, extension)
output_name = _trimmed_filename(source_file.file_name or os.path.basename(doc.import_file), output_extension)
new_file = save_file(
output_name,
output_content,
dt=doc.doctype,
dn=doc.name,
is_private=source_file.is_private,
df="import_file",
)
doc.import_file = new_file.file_url
doc.template_options = ""
doc.template_warnings = ""
doc.save()
return {
"import_file": new_file.file_url,
"file_name": new_file.file_name,
"trimmed_cells": trimmed_cells,
"summary": _("Trimmed {0} cell(s) and uploaded {1}.").format(trimmed_cells, new_file.file_name),
}