✨ allow setting base URL from env

Author: ianardee

mindee/endpoints.py

@@ -7,7 +7,8 @@
 from mindee.logger import logger
 from mindee.versions import __version__, get_platform, python_version
 
-MINDEE_API_URL = "https://api.mindee.net/v1"
+MINDEE_BASE_URL = "https://api.mindee.net/v1"
+MINDEE_BASE_URL_NAME = "MINDEE_BASE_URL"
 MINDEE_API_KEY_NAME = "MINDEE_API_KEY"
 
 PLATFORM = get_platform()
@@ -26,6 +27,7 @@ class Endpoint:
     version: str
     api_key: str = ""
     timeout: int = DEFAULT_TIMEOUT
+    _mindee_url: str = MINDEE_BASE_URL
     _url_root: str
 
     def __init__(
@@ -49,9 +51,10 @@ def __init__(
             self.api_key = api_key
         else:
             self.set_api_key_from_env()
+        self.set_base_url_from_env()
 
         self._url_root = (
-            f"{MINDEE_API_URL}/products/{self.owner}/{self.url_name}/v{self.version}"
+            f"{self._mindee_url}/products/{self.owner}/{self.url_name}/v{self.version}"
         )
 
     @property
@@ -62,12 +65,19 @@ def base_headers(self) -> Dict[str, str]:
             "User-Agent": USER_AGENT,
         }
 
+    def set_base_url_from_env(self) -> None:
+        """Set the base URL from an environment variable, if present."""
+        env_val = os.getenv(MINDEE_BASE_URL_NAME, "")
+        if env_val:
+            self._mindee_url = env_val
+            logger.debug("Base URL set from environment")
+
     def set_api_key_from_env(self) -> None:
         """Set the endpoint's API key from an environment variable, if present."""
-        env_key = os.getenv(MINDEE_API_KEY_NAME, "")
-        if env_key:
-            self.api_key = env_key
-            logger.debug("Set API key from environment")
+        env_val = os.getenv(MINDEE_API_KEY_NAME, "")
+        if env_val:
+            self.api_key = env_val
+            logger.debug("API key set from environment")
 
     def predict_req_post(
         self,

mindee/input/sources.py

@@ -87,10 +87,15 @@ def process_pdf(
         pages_count = self.count_doc_pages()
         if on_min_pages > pages_count:
             return
+        all_pages = list(range(pages_count))
         if behavior == KEEP_ONLY:
-            pages_to_keep = set(page_indexes)
+            pages_to_keep = set()
+            for page_n in page_indexes:
+                try:
+                    pages_to_keep.add(all_pages[page_n])
+                except IndexError:
+                    logger.warning("Page index not in source document: %s", page_n)
         elif behavior == REMOVE:
-            all_pages = list(range(pages_count))
             pages_to_remove = set()
             for page_n in page_indexes:
                 try:

tests/data

@@ -74,8 +74,8 @@ def test_pdf_keep_5_first_pages():
 def test_pdf_keep_invalid_pages():
     input_obj = PathInput(f"{PDF_DATA_DIR}/multipage.pdf")
     assert input_obj.is_pdf() is True
-    input_obj.process_pdf(behavior=KEEP_ONLY, on_min_pages=2, page_indexes=[16, 17])
-    assert input_obj.count_doc_pages() == 12
+    input_obj.process_pdf(behavior=KEEP_ONLY, on_min_pages=2, page_indexes=[0, 1, 17])
+    assert input_obj.count_doc_pages() == 2
 
 
 def test_pdf_remove_5_last_pages():
@@ -104,8 +104,12 @@ def test_pdf_remove_invalid_pages():
 def test_pdf_keep_no_pages():
     input_obj = PathInput(f"{PDF_DATA_DIR}/multipage.pdf")
     assert input_obj.is_pdf() is True
+    # empty page indexes
     with pytest.raises(RuntimeError):
         input_obj.process_pdf(behavior=KEEP_ONLY, on_min_pages=2, page_indexes=[])
+    # all invalid pages
+    with pytest.raises(RuntimeError):
+        input_obj.process_pdf(behavior=KEEP_ONLY, on_min_pages=2, page_indexes=[16, 17])
 
 
 def test_pdf_remove_all_pages():