Merge pull request #1 from explosion/pre-release-fixes

add unit tests in ci

.github/workflows/unit_tests.yml

@@ -0,0 +1,47 @@
+name: Unit Tests
+
+on:
+  pull_request:
+  push:
+    branches:
+      - main
+
+jobs:
+  setup:
+    runs-on: ubuntu-latest
+    steps:
+      - uses: actions/checkout@v3
+      - name: Set up Python 3.9
+        uses: actions/setup-python@v4
+        with:
+          python-version: 3.9
+          cache: "pip" # caching pip dependencies
+
+      - name: Check out Prodigy
+        uses: actions/checkout@v3
+        with:
+          repository: explosion/prodigy
+          ref: v1.14.0
+          path: ./prodigy
+          ssh-key: ${{ secrets.GHA_PRODIGY_READ }}
+
+      - name: Install prodigy
+        run: |
+          ls -la
+          pip install ./prodigy
+
+      - name: Install dependencies
+        run: |
+          pip install --upgrade pip
+          pip install -e .
+          pip install ruff pytest
+      
+      - name: Run ruff
+        if: always()
+        shell: bash
+        run: python -m ruff prodigy_pdf tests
+
+      - name: Run pytest
+        if: always()
+        shell: bash
+        run: python -m pytest tests

prodigy_pdf/__init__.py

@@ -1,4 +1,4 @@
-from typing import Dict, List
+from typing import List
 import base64
 from io import BytesIO
 from pathlib import Path
@@ -10,14 +10,16 @@
 from prodigy.util import msg
 
 def page_to_image(page: pdfium.PdfPage) -> str:
+    """Turns a PdfPage into a base64 image for Prodigy"""
     pil_image = page.render().to_pil()
-    buffered = BytesIO()
-    pil_image.save(buffered, format="JPEG")
-    img_str = base64.b64encode(buffered.getvalue())
+    with BytesIO() as buffered:
+        pil_image.save(buffered, format="JPEG")
+        img_str = base64.b64encode(buffered.getvalue())
     return f"data:image/png;base64,{img_str.decode('utf-8')}"
 
 
-def generate_pdf_pages(pdf_paths: List[Path]) -> Dict:
+def generate_pdf_pages(pdf_paths: List[Path]):
+    """Generate dictionaries that contain an image for each page in the PDF"""
     for pdf_path in pdf_paths:
         pdf = pdfium.PdfDocument(pdf_path)
         n_pages = len(pdf)
@@ -30,6 +32,7 @@ def generate_pdf_pages(pdf_paths: List[Path]) -> Dict:
                     "pdf": pdf_path.parts[-1],
                 }
             })
+        pdf.close()
 
 
 @recipe(
@@ -47,6 +50,7 @@ def pdf_image_manual(
     labels:str,
     remove_base64:bool=False
 ) -> ControllerComponentsDict:
+    """Turns pdfs into images in order to annotate them."""
     # Read in stream as a list for progress bar.
     if not pdf_folder.exists():
         msg.fail(f"Folder `{pdf_folder}` does not exist.", exits=True)

pyproject.toml

@@ -0,0 +1,2 @@
+[tool.ruff]
+line-length = 120

setup.cfg

@@ -10,6 +10,7 @@ zip_safe = true
 python_requires = >=3.8
 install_requires =
     pypdfium2==4.20.0
+    Pillow==9.4.0
 
 [options.entry_points]
 prodigy_recipes =