PDF processing API using 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

# Import Flask and other functions
from flask import Flask, request, send_file, jsonify
from werkzeug.utils import secure_filename
import tempfile
imports
import gc
from io import BytesIO

# Import CORS to receive requests from any origin
from flask_cors import CORS

# import PyPDF2
from PyPDF2 import PdfReader, PdfWriter

# Create the Flask instance
app = Flask(__name__)

# enable CORS in the Flask app
CORS(app)

# The function that processes the PDF
# This example only removes links from the pdf
def process_pdf_file(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,
            "Mistake": "",
        }
    except Exception as e:
        return {
            "Success": False,
            "return_path": "",
            "Error": str(e),
        }

# Define the endpoint of the api, the allowed methods and the function that is called
@app.route('/upload', methods=['POST'])
def upload_file():
    # We verify that a file has been uploaded
    if 'file' not in request.files:
        return {"Success": False, "Error": "No file part in the request."}, 400

    # We verify that the file has a name (it is not an empty field)
    file = request.files['file']
    if file.filename == '':
        return {"Success": False, "Error": "No file selected."}, 400

    # Check if the file is a PDF
    if not file.filename.lower().endswith('.pdf'):
        return {"Success": False, "Error": "Invalid file type. Only PDF files are allowed."}, 400
    
    # temporarily save the file with a random name
    with tempfile.NamedTemporaryFile(delete=False, suffix=".pdf") as temp_pdf:
        file.save(temp_pdf.name)

        # The processing function is called
        result = process_pdf_file(temp_pdf.name)

        temp_pdf.close()

        # delete the initial file
        os.remove(temp_pdf.name)

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

        # We save the data of the modified file
        with open(output_path, 'rb') as f:
            file_data = BytesIO(f.read())

        # And delete it from memory
        os.remove(output_path)

        # Call the garbage collector
        # (I've had memory leak failures)
        gc.collect()

        # return the modified file
        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

# To run locally
if __name__ == '__main__':
    app.run()