Procesa PDFs en una API con Flask + Heroku + Gunicorn

  1
  2
  3
  4
  5
  6
  7
  8
  9
 10
 11
 12
 13
 14
 15
 16
 17
 18
 19
 20
 21
 22
 23
 24
 25
 26
 27
 28
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102

# Importar Flask y otras funciones
from flask import Flask, request, send_file, jsonify
from werkzeug.utils import secure_filename
import tempfile
import os
import gc
from io import BytesIO

# Importar CORS para recibir solicitudes de cualquier origen
from flask_cors import CORS

# Importar PyPDF2
from PyPDF2 import PdfReader, PdfWriter

# Crear la instancia de Flask
app = Flask(__name__)

# Habilitar CORS en la aplicación Flask
CORS(app)

# La función que procesa el PDF
# En este ejemplo lo único que se hace es eliminar los enlaces
def procesar_archivo_pdf(pdf_path):
    try:
        output = pdf_path[:-4] + "_out.pdf"
        reader = PdfReader(pdf_path)
        writer = PdfWriter()

        for i in range(len(reader.pages)):
            page = reader.pages[i]
            writer.add_page(page)
            
        writer.remove_links()

        with open(output, "wb") as fp:
            writer.write(fp)

        gc.collect()
        return {
            "Success": True,
            "return_path": output,
            "Error": "",
        }
    except Exception as e:
        return {
            "Success": False,
            "return_path": "",
            "Error": str(e),
        }

# Define el endpoint de la api, los métodos permitidos y la función que se llama
@app.route('/upload', methods=['POST'])
def upload_file():

    # Verificamos que se ha subido un archivo
    if 'file' not in request.files:
        return {"Success": False, "Error": "No file part in the request."}, 400

    # Verificamos que el archivo tiene nombre (no es un campo vacío)
    file = request.files['file']
    if file.filename == '':
        return {"Success": False, "Error": "No file selected."}, 400

    # Verificar si el archivo es un PDF
    if not file.filename.lower().endswith('.pdf'):
        return {"Success": False, "Error": "Invalid file type. Only PDF files are allowed."}, 400
    
    # guardamos temporalmente el archivo con un nombre aleatorio
    with tempfile.NamedTemporaryFile(delete=False, suffix=".pdf") as temp_pdf:
        file.save(temp_pdf.name)

        # Se llama a la función de procesamiento
        result = procesar_archivo_pdf(temp_pdf.name)

        temp_pdf.close()

        # Eliminar el archivo inicial
        os.remove(temp_pdf.name)

    if result["Success"]:
        output_path = result["return_path"]

        # Guardamos los datos del arhivo modificado
        with open(output_path, 'rb') as f:
            file_data = BytesIO(f.read())

        # Y lo borramos de memoria
        os.remove(output_path)

        # Llamada al colector de basura
        # (He tenido fallos de fugas de memoria)
        gc.collect()

        # Devolver el archivo modificado
        return send_file(file_data, as_attachment=True, download_name=secure_filename(os.path.basename(output_path)))
    else:
        gc.collect()
        return jsonify({"Success": False, "Error": result["Error"]}), 500

# Para ejecutar en local
if __name__ == '__main__':
    app.run()