app.py

#!/usr/bin/python
#
# Copyright 2025 Google LLC
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#      http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""Qt application to view PDFs and extract regions to SVG.

Sample applications include extracting charts, graphs, model architectures,
algorithms, equations, etc. from PDFs to share them via blog posts or social
media such that they remain clear and scalable vector graphics, as opposed to
pixelated raster images.

Before using, please see these instructions to install necessary dependencies:
https://github.com/mbrukman/pdf-extract-svg?tab=readme-ov-file#installation
"""

import os
import re
import subprocess
import sys

from PySide6.QtWidgets import (
    QApplication, QMainWindow, QFileDialog, QLabel, QWidget,
    QVBoxLayout, QHBoxLayout, QPushButton, QScrollArea, QLineEdit,
)
from PySide6.QtGui import (
    QPixmap, QPainter, QPen, QColor, QFont, QMouseEvent, QPaintEvent,
    QCloseEvent, QCursor,
)
from PySide6.QtCore import (
    Qt, QRect, QPoint, QSize, QPointF,
)


# DPI for the preview image. Higher values are clearer but use more memory.
PREVIEW_DPI = 150


class PDFViewerLabel(QLabel):
    """
    A custom QLabel to display the PDF page and handle mouse events for selection.
    """
    def __init__(self, parent=None):
        super().__init__(parent)
        self.setMouseTracking(True)
        self.selection_rect: QRect = QRect()
        self.start_point: QPoint = QPoint()
        self.end_point: QPoint = QPoint()
        self.is_selecting: bool = False
        self.is_selection_moving: bool = False
        self.drag_offset: QPoint = QPoint()

    def __is_position_inside_selection(self, position: QPointF) -> bool:
        """Check if the given position is inside the selection rectangle."""
        return (not self.selection_rect.isNull() and
                self.selection_rect.contains(position.toPoint(), proper=True))

    def mousePressEvent(self, event: QMouseEvent) -> None: # pylint: disable=invalid-name
        """Handle mouse click to start the selection or move existing selection."""
        if event.button() == Qt.MouseButton.LeftButton:
            # Check if clicking inside existing selection to move it
            if self.__is_position_inside_selection(event.position()):
                self.is_selection_moving = True
                self.drag_offset = event.position().toPoint() - self.selection_rect.topLeft()
            else:
                # Start new selection
                self.start_point = event.position().toPoint()
                self.selection_rect = QRect(self.start_point, QSize())
                self.is_selecting = True
            self.update()

    def mouseMoveEvent(self, event: QMouseEvent) -> None: # pylint: disable=invalid-name
        """Handle mouse drag while selection is in progress or being moved."""
        # Update cursor based on position
        # TODO(mbrukman): this is currently not working; the cursor shape does not update.
        if self.__is_position_inside_selection(event.position()):
            self.setCursor(QCursor(Qt.CursorShape.SizeAllCursor))
        else:
            self.setCursor(QCursor(Qt.CursorShape.ArrowCursor))
        self.update()

        if self.is_selecting:
            self.end_point = event.position().toPoint()
            self.selection_rect = QRect(self.start_point, self.end_point).normalized()
            self.update() # Trigger a repaint
        elif self.is_selection_moving:
            new_top_left = event.position().toPoint() - self.drag_offset

            # Constrain to widget bounds
            rect_width = self.selection_rect.width()
            rect_height = self.selection_rect.height()

            # Ensure we don't move outside the image area
            # The pixmap might be smaller than the widget if not resized, but usually it fits.
            # We constrain to the widget's rect for simplicity, or the pixmap rect if available.
            max_w = self.width()
            max_h = self.height()
            if self.pixmap():
                max_w = self.pixmap().width()
                max_h = self.pixmap().height()

            x = max(0, min(new_top_left.x(), max_w - rect_width))
            y = max(0, min(new_top_left.y(), max_h - rect_height))

            self.selection_rect.moveTo(x, y)
            self.update()

    def mouseReleaseEvent(self, event: QMouseEvent) -> None: # pylint: disable=invalid-name
        """Handle mouse button release while making a selection."""
        if event.button() == Qt.MouseButton.LeftButton:
            self.is_selecting = False
            self.is_selection_moving = False
            self.update()

    def paintEvent(self, event: QPaintEvent) -> None: # pylint: disable=invalid-name
        """Draw the selected region on top of the rendered PDF page."""
        super().paintEvent(event) # Draw the pixmap first

        # Draw the selection box if available (both during dragging and after).
        if not self.selection_rect.isNull():
            painter = QPainter(self)
            pen = QPen(QColor(0, 120, 215, 200), 2, Qt.PenStyle.SolidLine)
            painter.setPen(pen)

            fill_color = QColor(0, 120, 215, 50)
            painter.fillRect(self.selection_rect, fill_color)
            painter.drawRect(self.selection_rect)


class MainWindow(QMainWindow):
    """Handles the UI of the application."""

    def __init__(self):
        """Initialize the UI elements and connect them to handlers."""

        super().__init__()
        self.setWindowTitle("PDF Vector Extractor")
        self.setGeometry(100, 100, 1000, 800)

        # Member variables
        self.pdf_path: str = ''
        self.temp_image_base: str = "pdf_extractor_temp_page"
        self.current_page: int = 0
        self.total_pages: int = 0
        self.page_size_points: tuple[float, float] = (0, 0) # Store page size in points (w, h)

        # --- UI Setup ---
        main_layout: QVBoxLayout = QVBoxLayout()
        control_layout: QHBoxLayout = QHBoxLayout()

        self.btn_open: QPushButton = QPushButton("Open PDF")
        self.btn_prev: QPushButton = QPushButton("<< Prev")
        self.btn_next: QPushButton = QPushButton("Next >>")
        self.lbl_page_prefix: QLabel = QLabel("Page: ")
        self.txt_page_num: QLineEdit = QLineEdit()
        self.txt_page_num.setFixedWidth(50)
        self.txt_page_num.setAlignment(Qt.AlignmentFlag.AlignCenter)
        self.lbl_total_pages: QLabel = QLabel(" / N/A")
        self.btn_save: QPushButton = QPushButton("Save Region as SVG")

        self.btn_prev.setEnabled(False)
        self.btn_next.setEnabled(False)
        self.btn_save.setEnabled(False)
        self.txt_page_num.setEnabled(False)

        control_layout.addWidget(self.btn_open)
        control_layout.addStretch()
        control_layout.addWidget(self.btn_prev)
        control_layout.addWidget(self.lbl_page_prefix)
        control_layout.addWidget(self.txt_page_num)
        control_layout.addWidget(self.lbl_total_pages)
        control_layout.addWidget(self.btn_next)
        control_layout.addStretch()
        control_layout.addWidget(self.btn_save)

        # Scrollable area for the PDF viewer
        self.scroll_area: QScrollArea = QScrollArea(parent=self)
        # Allow the widget to be larger than the viewport, enabling scrolling.
        self.scroll_area.setWidgetResizable(False)
        # A dark background makes the page stand out, using a stylesheet.
        self.scroll_area.setStyleSheet("background-color: #3c3c3c;")
        self.scroll_area.setMouseTracking(True)
        self.scroll_area.viewport().setMouseTracking(True)

        # Custom label for viewing and selection
        self.viewer: PDFViewerLabel = PDFViewerLabel(parent=self.scroll_area)
        self.viewer.setAlignment(Qt.AlignmentFlag.AlignCenter)
        self.scroll_area.setWidget(self.viewer)

        main_layout.addLayout(control_layout)
        main_layout.addWidget(self.scroll_area)

        container: QWidget = QWidget()
        container.setLayout(main_layout)
        self.setCentralWidget(container)

        # --- Connections ---
        self.btn_open.clicked.connect(self.open_pdf)
        self.btn_prev.clicked.connect(self.prev_page)
        self.btn_next.clicked.connect(self.next_page)
        self.btn_save.clicked.connect(self.save_svg)
        self.txt_page_num.returnPressed.connect(self.jump_to_page)

    def display_error(self, message: str) -> None:
        """Helper to display error messages in the viewer."""
        self.viewer.clear()
        self.viewer.setText(message)
        self.viewer.setFont(QFont("Arial", 12))
        self.viewer.setStyleSheet("color: red;")

    def open_pdf(self) -> None:
        """Tries to open a PDF and, if successful, renders the first page."""
        path, _ = QFileDialog.getOpenFileName(self, "Open PDF", "", "PDF Files (*.pdf)")
        if path:
            self.pdf_path = path
            self.current_page = 0
            if self.get_pdf_info():
                self.render_page()

    def get_pdf_info(self) -> bool:
        """Gets total pages using pdfinfo and returns status."""
        try:
            command: list[str] = [
                "pdfinfo",
                self.pdf_path,
            ]
            result = subprocess.run(
                command, capture_output=True, text=True, check=True, encoding='utf-8',
                errors='ignore',
            )
            for line in result.stdout.splitlines():
                if "Pages:" in line:
                    self.total_pages = int(line.split(":")[1].strip())
            return True
        except FileNotFoundError:
            self.display_error("Error: `pdfinfo` command not found.")
            return False
        except subprocess.CalledProcessError as e:
            self.display_error(f"Failed to get PDF info.\n"
                               f"PDF may be corrupted or encrypted.\n\n"
                               f"Poppler Error:\n{e.stderr}")
            return False

    def update_page_size(self) -> None:
        """Updates the size in points for the current page."""
        try:
            command: list[str] = [
                "pdfinfo",
                "-f", str(self.current_page + 1),
                "-l", str(self.current_page + 1),
                self.pdf_path
            ]
            result = subprocess.run(
                command, capture_output=True, text=True, check=True, encoding='utf-8',
                errors='ignore',
            )
            if result.stderr:
                print(f'pdfinfo warnings: {result.stderr}', file=sys.stderr)
            for line in result.stdout.splitlines():
                # Line format: "Page (page number) size: (width) x (height) pts"
                match = re.match(r'Page\s+\d+\s+size:\s+([\d\.]+)\s+x\s+([\d\.]+)\s+', line)
                if match:
                    width = float(match.group(1))
                    height = float(match.group(2))
                    self.page_size_points = (width, height)
                    return
            self.page_size_points = (0, 0)
        except (subprocess.CalledProcessError, FileNotFoundError, IndexError) as e:
            print(f"Could not get page size for page {self.current_page + 1}: {e}", file=sys.stderr)
            self.page_size_points = (0, 0)

    def render_page(self) -> None:
        """Render a single page from a PDF to a file and display it."""
        if not self.pdf_path:
            return

        self.update_page_size()
        temp_image_path = f"{self.temp_image_base}-{self.current_page + 1}"

        # Clean up old temp files before creating a new one
        self.cleanup_temp_files()

        try:
            # Use pdftoppm to render the current page to a PNG
            command: list[str] = [
                "pdftoppm",
                "-f", str(self.current_page + 1),
                "-l", str(self.current_page + 1),
                "-png",
                "-r", str(PREVIEW_DPI),
                "-singlefile",
                self.pdf_path,
                temp_image_path, # pdftoppm appends ".png"
            ]
            result = subprocess.run(command, check=True, capture_output=True)

            if result.stderr:
                print("pdftoppm warnings:\n"
                      f"{result.stderr.decode('utf-8', 'ignore')}", file=sys.stderr)

            generated_file = temp_image_path + ".png"
            if os.path.exists(generated_file):
                pixmap = QPixmap(generated_file)
                self.viewer.setStyleSheet("") # Reset stylesheet
                self.viewer.setPixmap(pixmap)
                self.viewer.adjustSize() # Resize the label to fit the pixmap
            else:
                raise FileNotFoundError(f"pdftoppm did not create expected file: {generated_file}")

            self.update_ui_state()

        except subprocess.CalledProcessError as e:
            self.display_error(f"Failed to render page {self.current_page + 1}.\n\n"
                               f"Poppler Error:\n{e.stderr.decode('utf-8', 'ignore')}")
        except (FileNotFoundError) as e:
            self.display_error(f"Failed to render page {self.current_page + 1}.\nError: {e}")
        finally:
            # Clean up the generated PNG immediately after loading it
            self.cleanup_temp_files()

    def update_ui_state(self) -> None:
        """Enable/disable buttons based on current state."""
        self.txt_page_num.setText(str(self.current_page + 1))
        self.lbl_total_pages.setText(f" / {self.total_pages}")
        self.btn_prev.setEnabled(self.current_page > 0)
        self.btn_next.setEnabled(self.current_page < self.total_pages - 1)
        self.btn_save.setEnabled(True)
        self.txt_page_num.setEnabled(True)

    def jump_to_page(self) -> None:
        """Jump to the page specified in the text box."""
        try:
            page_num = int(self.txt_page_num.text())
            if 1 <= page_num <= self.total_pages:
                self.current_page = page_num - 1
                self.render_page()
            else:
                # Reset to current page if out of range
                self.txt_page_num.setText(str(self.current_page + 1))
        except ValueError:
            # Reset to current page if invalid input
            self.txt_page_num.setText(str(self.current_page + 1))

    def prev_page(self) -> None:
        """Move to the previous page in this file, if we're not at the beginning."""
        if self.current_page > 0:
            self.current_page -= 1
            self.render_page()

    def next_page(self) -> None:
        """Move to the next page in this file, if we're not at the end."""
        if self.current_page < self.total_pages - 1:
            self.current_page += 1
            self.render_page()

    def save_svg(self) -> None:
        """Write selected region to an SVG."""
        selection = self.viewer.selection_rect
        if selection.isNull() or selection.width() < 2 or selection.height() < 2:
            self.statusBar().showMessage("Please select a valid region first.", 3000)
            return

        output_path, _ = QFileDialog.getSaveFileName(self, "Save SVG", "", "SVG Files (*.svg)")
        if not output_path:
            return

        # --- Coordinate Conversion ---
        pixmap = self.viewer.pixmap()
        if not pixmap or pixmap.isNull():
            self.statusBar().showMessage("Cannot save: No page rendered.", 4000)
            return

        pixmap_size = pixmap.size()
        pdf_w_pt, pdf_h_pt = self.page_size_points

        if pixmap_size.width() == 0 or pixmap_size.height() == 0:
            self.statusBar().showMessage("Cannot save: Invalid page dimensions.", 4000)
            return

        # Calculate the scale factor
        scale_x = pdf_w_pt / pixmap_size.width()
        scale_y = pdf_h_pt / pixmap_size.height()

        # Apply the scale to the selection rectangle
        x_pt = selection.x() * scale_x
        y_pt = selection.y() * scale_y
        w_pt = selection.width() * scale_x
        h_pt = selection.height() * scale_y

        # --- Run pdftocairo to extract the SVG ---
        try:
            command: list[str] = [
                "pdftocairo",
                "-svg",
                "-f", str(self.current_page + 1),
                "-l", str(self.current_page + 1),
                "-x", f'{int(x_pt)}', # Needs to be integer valued!
                "-y", f'{int(y_pt)}', # Needs to be integer valued!
                "-W", f'{int(w_pt)}', # Needs to be integer valued!
                "-H", f'{int(h_pt)}', # Needs to be integer valued!
                "-paperw", f'{int(w_pt)}', # Set paper width to crop width; needs to be int
                "-paperh", f'{int(h_pt)}', # Set paper height to crop height; needs to be int
                "-nocenter",
                self.pdf_path,
                output_path,
            ]
            result = subprocess.run(
                command, check=True, capture_output=True, text=True, encoding='utf-8',
                errors='ignore',
            )

            # Check stderr for warnings even if the command succeeds
            if result.stderr:
                print(f"pdftocairo warnings:\n{result.stderr}", file=sys.stderr)
                self.statusBar().showMessage("Saved with warnings. See console for details.", 5000)
            else:
                self.statusBar().showMessage(f"Successfully saved to {output_path}", 5000)

        except subprocess.CalledProcessError as e:
            # The error you are seeing will be caught here
            print(f"Error saving SVG: {e.stderr.strip()}", file=sys.stderr)
            self.statusBar().showMessage(f"Error saving SVG: {e.stderr.strip()}", 10000)
        except FileNotFoundError:
            self.statusBar().showMessage("Error: 'pdftocairo' command not found.", 5000)

    def cleanup_temp_files(self) -> None:
        """Clean up any temporary image files."""
        for f in os.listdir('.'):
            if f.startswith(self.temp_image_base) and f.endswith(".png"):
                try:
                    os.remove(f)
                except OSError as e:
                    print(f"Error removing temp file {f}: {e}", file=sys.stderr)

    def closeEvent(self, event: QCloseEvent) -> None: # pylint: disable=invalid-name
        """Clean up temporary files on exit."""
        self.cleanup_temp_files()
        super().closeEvent(event)


def main(argv):
    all_tools = ["pdftocairo", "pdftoppm", "pdfinfo"]
    missing_tools = []

    for tool in all_tools:
        try:
            # Use --version to check for existence without processing a file
            subprocess.run([tool, "-v"], check=True, capture_output=True)
        except (subprocess.CalledProcessError, FileNotFoundError):
            missing_tools.append(tool)

    if missing_tools:
        print(f"Error: Poppler tools not found: {', '.join(missing_tools)}.\n"
              "Please install 'poppler-utils' and ensure it's in your PATH.\n"
              "More info: "
              "https://github.com/mbrukman/pdf-extract-svg?tab=readme-ov-file#installation",
              file=sys.stderr)
        sys.exit(1)

    app = QApplication(argv)
    window = MainWindow()
    window.show()
    sys.exit(app.exec())


if __name__ == "__main__":
    main(sys.argv)