Refactor PDF split logic for improved clarity

Replaces getPagesToSplit with getCustomPagesToSplit and introduces shouldSplitPage to centralize split mode logic. This refactor improves code readability and maintainability by separating custom page selection from split mode handling.

app/core/src/main/java/stirling/software/SPDF/controller/api/SplitPdfBySectionsController.java

@@ -1,7 +1,7 @@
 package stirling.software.SPDF.controller.api;
 
+import java.io.ByteArrayOutputStream;
 import java.io.IOException;
-import java.io.OutputStream;
 import java.nio.file.Files;
 import java.util.*;
 import java.util.zip.ZipEntry;
@@ -22,7 +22,6 @@ import org.springframework.web.bind.annotation.PostMapping;
 import org.springframework.web.bind.annotation.RequestMapping;
 import org.springframework.web.bind.annotation.RestController;
 import org.springframework.web.multipart.MultipartFile;
-import org.springframework.web.servlet.mvc.method.annotation.StreamingResponseBody;
 
 import io.swagger.v3.oas.annotations.Operation;
 import io.swagger.v3.oas.annotations.tags.Tag;
@@ -34,7 +33,6 @@ import stirling.software.SPDF.model.api.SplitPdfBySectionsRequest;
 import stirling.software.common.service.CustomPDFDocumentFactory;
 import stirling.software.common.util.ExceptionUtils;
 import stirling.software.common.util.GeneralUtils;
-import stirling.software.common.util.PDFService;
 import stirling.software.common.util.TempFile;
 import stirling.software.common.util.TempFileManager;
 import stirling.software.common.util.WebResponseUtils;
@@ -47,7 +45,6 @@ public class SplitPdfBySectionsController {
 
     private final CustomPDFDocumentFactory pdfDocumentFactory;
     private final TempFileManager tempFileManager;
-    private final PDFService pdfService;
 
     @PostMapping(value = "/split-pdf-by-sections", consumes = MediaType.MULTIPART_FORM_DATA_VALUE)
     @Operation(
@@ -57,8 +54,8 @@ public class SplitPdfBySectionsController {
                             + " which page to split, and how to split"
                             + " ( halves, thirds, quarters, etc.), both vertically and horizontally."
                             + " Input:PDF Output:ZIP-PDF Type:SISO")
-    public ResponseEntity<StreamingResponseBody> splitPdf(
-            @ModelAttribute SplitPdfBySectionsRequest request) throws Exception {
+    public ResponseEntity<?> splitPdf(@ModelAttribute SplitPdfBySectionsRequest request)
+            throws Exception {
         MultipartFile file = request.getFileInput();
         String pageNumbers = request.getPageNumbers();
         SplitTypes splitMode =
@@ -66,9 +63,10 @@ public class SplitPdfBySectionsController {
                         .map(SplitTypes::valueOf)
                         .orElse(SplitTypes.SPLIT_ALL);
 
-        try (PDDocument sourceDocument = pdfDocumentFactory.load(file)) {
-            Set<Integer> pagesToSplit =
-                    getPagesToSplit(pageNumbers, splitMode, sourceDocument.getNumberOfPages());
+        try (PDDocument sourceDocument = pdfDocumentFactory.load(file, true)) {
+            int totalPages = sourceDocument.getNumberOfPages();
+            Set<Integer> customPagesToSplit =
+                    getCustomPagesToSplit(pageNumbers, splitMode, totalPages);
 
             // Process the PDF based on split parameters
             int horiz = request.getHorizontalDivisions() + 1;
@@ -77,14 +75,13 @@ public class SplitPdfBySectionsController {
             String filename = GeneralUtils.generateFilename(file.getOriginalFilename(), "_split");
 
             if (merge) {
-                TempFile tempFile = new TempFile(tempFileManager, ".pdf");
-                try (PDDocument mergedDoc = pdfDocumentFactory.createNewDocument();
-                        OutputStream out = Files.newOutputStream(tempFile.getPath())) {
+                try (PDDocument mergedDoc =
+                                pdfDocumentFactory.createNewDocumentBasedOnOldDocument(
+                                        sourceDocument);
+                        ByteArrayOutputStream baos = new ByteArrayOutputStream()) {
                     LayerUtility layerUtility = new LayerUtility(mergedDoc);
-                    for (int pageIndex = 0;
-                            pageIndex < sourceDocument.getNumberOfPages();
-                            pageIndex++) {
-                        if (pagesToSplit.contains(pageIndex)) {
+                    for (int pageIndex = 0; pageIndex < totalPages; pageIndex++) {
+                        if (shouldSplitPage(pageIndex, totalPages, splitMode, customPagesToSplit)) {
                             addSplitPageToTarget(
                                     sourceDocument,
                                     pageIndex,
@@ -96,18 +93,16 @@ public class SplitPdfBySectionsController {
                             addPageToTarget(sourceDocument, pageIndex, mergedDoc, layerUtility);
                         }
                     }
-                    mergedDoc.save(out);
+                    mergedDoc.save(baos);
+                    return WebResponseUtils.baosToWebResponse(baos, filename + ".pdf");
                 }
-                return WebResponseUtils.pdfFileToWebResponse(tempFile, filename + ".pdf");
             } else {
                 TempFile zipTempFile = new TempFile(tempFileManager, ".zip");
                 try (ZipOutputStream zipOut =
                         new ZipOutputStream(Files.newOutputStream(zipTempFile.getPath()))) {
-                    for (int pageIndex = 0;
-                            pageIndex < sourceDocument.getNumberOfPages();
-                            pageIndex++) {
+                    for (int pageIndex = 0; pageIndex < totalPages; pageIndex++) {
                         int pageNum = pageIndex + 1;
-                        if (pagesToSplit.contains(pageIndex)) {
+                        if (shouldSplitPage(pageIndex, totalPages, splitMode, customPagesToSplit)) {
                             for (int i = 0; i < horiz; i++) {
                                 for (int j = 0; j < verti; j++) {
                                     try (PDDocument subDoc =
@@ -245,53 +240,37 @@ public class SplitPdfBySectionsController {
     }
 
     // Based on the mode, get the pages that need to be split and return the pages set
-    private Set<Integer> getPagesToSplit(String pageNumbers, SplitTypes splitMode, int totalPages) {
-        Set<Integer> pagesToSplit = new HashSet<>();
-
-        switch (splitMode) {
-            case CUSTOM:
-                if (pageNumbers == null || pageNumbers.isBlank()) {
-                    throw ExceptionUtils.createIllegalArgumentException(
-                            "error.argumentRequired",
-                            "{0} is required for {1} mode",
-                            "page numbers",
-                            "custom");
-                }
-                String[] pageOrderArr = pageNumbers.split(",");
-                List<Integer> pageListToSplit =
-                        GeneralUtils.parsePageList(pageOrderArr, totalPages, false);
-                pagesToSplit.addAll(pageListToSplit);
-                break;
-
-            case SPLIT_ALL:
-                for (int i = 0; i < totalPages; i++) {
-                    pagesToSplit.add(i);
-                }
-                break;
-
-            case SPLIT_ALL_EXCEPT_FIRST:
-                for (int i = 1; i < totalPages; i++) {
-                    pagesToSplit.add(i);
-                }
-                break;
-
-            case SPLIT_ALL_EXCEPT_LAST:
-                for (int i = 0; i < totalPages - 1; i++) {
-                    pagesToSplit.add(i);
-                }
-                break;
-
-            case SPLIT_ALL_EXCEPT_FIRST_AND_LAST:
-                for (int i = 1; i < totalPages - 1; i++) {
-                    pagesToSplit.add(i);
-                }
-                break;
-
-            default:
-                throw ExceptionUtils.createIllegalArgumentException(
-                        "error.invalidFormat", "Invalid {0} format: {1}", "split mode", splitMode);
+    private Set<Integer> getCustomPagesToSplit(
+            String pageNumbers, SplitTypes splitMode, int totalPages) {
+        if (splitMode != SplitTypes.CUSTOM) {
+            return Collections.emptySet();
         }
+        if (pageNumbers == null || pageNumbers.isBlank()) {
+            throw ExceptionUtils.createIllegalArgumentException(
+                    "error.argumentRequired",
+                    "{0} is required for {1} mode",
+                    "page numbers",
+                    "custom");
+        }
+        String[] pageOrderArr = pageNumbers.split(",");
+        List<Integer> pageListToSplit = GeneralUtils.parsePageList(pageOrderArr, totalPages, false);
+        return new HashSet<>(pageListToSplit);
+    }
 
-        return pagesToSplit;
+    private boolean shouldSplitPage(
+            int pageIndex, int totalPages, SplitTypes splitMode, Set<Integer> customPagesToSplit) {
+        return switch (splitMode) {
+            case CUSTOM -> customPagesToSplit.contains(pageIndex);
+            case SPLIT_ALL -> true;
+            case SPLIT_ALL_EXCEPT_FIRST -> pageIndex > 0;
+            case SPLIT_ALL_EXCEPT_LAST -> pageIndex < totalPages - 1;
+            case SPLIT_ALL_EXCEPT_FIRST_AND_LAST -> pageIndex > 0 && pageIndex < totalPages - 1;
+            default ->
+                    throw ExceptionUtils.createIllegalArgumentException(
+                            "error.invalidFormat",
+                            "Invalid {0} format: {1}",
+                            "split mode",
+                            splitMode);
+        };
     }
 }

testing/cucumber/features/steps/step_definitions.py

@@ -15,6 +15,7 @@ import re
 from PIL import Image, ImageDraw
 
 API_HEADERS = {"X-API-KEY": "123456789"}
+TIMEOUT = 80  # seconds
 
 #########
 # GIVEN #
@@ -298,7 +299,11 @@ def save_generated_pdf(context, filename):
 def step_send_get_request(context, endpoint):
     base_url = "http://localhost:8080"
     full_url = f"{base_url}{endpoint}"
-    response = requests.get(full_url, headers=API_HEADERS)
+    response = requests.get(
+        full_url,
+        headers=API_HEADERS,
+        timeout=getattr(context, "request_timeout", TIMEOUT),
+    )
     context.response = response
 
 
@@ -307,7 +312,12 @@ def step_send_get_request_with_params(context, endpoint):
     base_url = "http://localhost:8080"
     params = {row["parameter"]: row["value"] for row in context.table}
     full_url = f"{base_url}{endpoint}"
-    response = requests.get(full_url, params=params, headers=API_HEADERS)
+    response = requests.get(
+        full_url,
+        params=params,
+        headers=API_HEADERS,
+        timeout=getattr(context, "request_timeout", TIMEOUT),
+    )
     context.response = response
 
 
@@ -337,7 +347,12 @@ def step_send_api_request(context, endpoint):
         print(f"form_data {file.name} with {mime_type}")
         form_data.append((key, (file.name, file, mime_type)))
 
-    response = requests.post(url, files=form_data, headers=API_HEADERS)
+    response = requests.post(
+        url,
+        files=form_data,
+        headers=API_HEADERS,
+        timeout=getattr(context, "request_timeout", TIMEOUT),
+    )
     context.response = response