plumberserver/main.py

import fitz
from flask import Flask, request
from flask_cors import CORS
from werkzeug.utils import secure_filename
from waitress import serve
from os import getenv, path
from tempfile import gettempdir

app = Flask(__name__)
CORS(app)

bullets = ["▶"]

@app.post("/")
def process_file():
    file = request.files["file"]
    if file is None:
        return "file not found", 400

    assert file.filename is not None
    filename = path.join(gettempdir(), secure_filename(file.filename))
    file.save(filename)

    doc = fitz.open(filename)

    textpages = []
    for page in doc:
        assert isinstance(page, fitz.Page)
        textpages.append(page.get_textpage().extractDICT())
    doc.close()

    sizes: set[int] = set()
    for page in textpages:
        for block in page["blocks"]:
            for line in block["lines"]:
                for span in line["spans"]:
                    sizes.add(round(span["size"]))

    sorted_sizes = sorted(sizes)
    h1_size = sorted_sizes[-1]
    h2_size = sorted_sizes[-2]

    pass1 = []
    for p in textpages:
        pass1_blocks = []
        for block in p["blocks"]:
            pass1_lines = []
            for line in block["lines"]:
                is_bullet = False
                largest_size = 0
                smallest_left = 10000
                line_text = ""
                for span_index, span in enumerate(line["spans"]):
                    t = span["text"]
                    if t in bullets and span_index == 0:
                        is_bullet = True
                        continue

                    span_size = round(span["size"])
                    if span_size > largest_size:
                        largest_size = span_size
                    span_left = round(span["origin"][0])
                    if span_left < smallest_left:
                        smallest_left = span_left
                    line_text += t

                line_type = "p"
                if is_bullet:
                    line_type = "li"
                elif largest_size == h1_size:
                    line_type = "h1"
                elif largest_size == h2_size:
                    line_type = "h2"

                pass1_lines.append({
                    "type": line_type,
                    "value": line_text,
                    "left": smallest_left,
                })

            pass1_blocks.append(pass1_lines)
        pass1.append(pass1_blocks)

    return pass1

if __name__ == "__main__":
    
    port = getenv("PORT")
    portNum = 8080
    if port is not None and port.isnumeric():
        portNum = int(port)

    print("using port {}".format(portNum))
    serve(app, port=portNum)