PR #736 Code style cleanup

to improve signal/noise ratio of PR

(using darker and ruff)

Author: soxofaan

openeo/extra/job_management/__init__.py

@@ -33,6 +33,7 @@
 from urllib3.util import Retry
 
 from openeo import BatchJob, Connection
+from openeo.extra.job_management.thread_worker import _JobManagerWorkerThreadPool
 from openeo.internal.processes.parse import (
     Parameter,
     Process,
@@ -41,8 +42,6 @@
 from openeo.rest import OpenEoApiError
 from openeo.rest.auth.auth import BearerAuth
 from openeo.util import LazyLoadCache, deep_get, repr_truncate, rfc3339
-from openeo.extra.job_management.thread_worker import _JobManagerWorkerThreadPool
-
 
 _log = logging.getLogger(__name__)
 
@@ -369,7 +368,15 @@ def run_loop():
             # TODO: support user-provided `stats`
             stats = collections.defaultdict(int)
 
-            while sum(job_db.count_by_status(statuses=["not_started", "created", "queued", "queued_for_start", "running"]).values()) > 0  and not self._stop_thread:
+            while (
+                sum(
+                    job_db.count_by_status(
+                        statuses=["not_started", "created", "queued", "queued_for_start", "running"]
+                    ).values()
+                )
+                > 0
+                and not self._stop_thread
+            ):
                 self._job_update_loop(job_db=job_db, start_job=start_job, stats=stats)
                 stats["run_jobs loop"] += 1
 
@@ -380,10 +387,8 @@ def run_loop():
                     if self._stop_thread:
                         break
 
-            
         self._thread = Thread(target=run_loop)
         self._thread.start()
-        
 
     def stop_job_thread(self, timeout_seconds: Optional[float] = _UNSET):
         """
@@ -402,7 +407,6 @@ def stop_job_thread(self, timeout_seconds: Optional[float] = _UNSET):
             if timeout_seconds is _UNSET:
                 timeout_seconds = 2 * self.poll_sleep
             self._thread.join(timeout_seconds)
-            
             if self._thread.is_alive():
                 _log.warning("Job thread did not stop after timeout")
         else:
@@ -502,7 +506,14 @@ def run_jobs(
 
         self._worker_pool = _JobManagerWorkerThreadPool()
 
-        while sum(job_db.count_by_status(statuses=["not_started", "created", "queued", "queued_for_start", "running"]).values()) > 0:
+        while (
+            sum(
+                job_db.count_by_status(
+                    statuses=["not_started", "created", "queued", "queued_for_start", "running"]
+                ).values()
+            )
+            > 0
+        ):
             self._job_update_loop(job_db=job_db, start_job=start_job, stats=stats)
             stats["run_jobs loop"] += 1
 
@@ -564,7 +575,6 @@ def _job_update_loop(
 
         for job, row in jobs_cancel:
             self.on_job_cancel(job, row)
- 
 
     def _launch_job(self, start_job, df, i, backend_name, stats: Optional[dict] = None):
         """Helper method for launching jobs
@@ -623,14 +633,12 @@ def _launch_job(self, start_job, df, i, backend_name, stats: Optional[dict] = No
                             _log.info(f"Job : {job.job_id} created, submitting to thread pool")
                             job_con = job.connection
                             self._worker_pool.submit_work(
-                                
-                                    _JobManagerWorkerThreadPool.WORK_TYPE_START_JOB,
-                                    (
-                                        job_con.root_url,
-                                        job_con.auth.bearer if isinstance(job_con.auth, BearerAuth) else None,
-                                        job.job_id,
-                                    ),
-                                
+                                _JobManagerWorkerThreadPool.WORK_TYPE_START_JOB,
+                                (
+                                    job_con.root_url,
+                                    job_con.auth.bearer if isinstance(job_con.auth, BearerAuth) else None,
+                                    job.job_id,
+                                ),
                             )
                             stats["job queued for start"] += 1
                             df.loc[i, "status"] = "queued_for_start"
@@ -706,7 +714,6 @@ def _cancel_prolonged_job(self, job: BatchJob, row):
             elapsed = current_time - job_running_start_time
 
             if elapsed > self._cancel_running_job_after:
-
                 _log.info(
                     f"Cancelling long-running job {job.job_id} (after {elapsed}, running since {job_running_start_time})"
                 )

openeo/extra/job_management/thread_worker.py

@@ -1,11 +1,12 @@
-
-from  typing import Tuple, Any, List, Dict
+import concurrent.futures
 import logging
+from typing import Any, Dict, List, Tuple
+
 import openeo
-import concurrent.futures
 
 _log = logging.getLogger(__name__)
 
+
 class _JobManagerWorkerThreadPool:
     """
     A worker thread pool for processing job management tasks asynchronously.
@@ -53,7 +54,6 @@ def submit_work(self, work_type: str, work_args: Tuple[Any, ...]) -> None:
         future = self._executor.submit(self._process_work_item, work_type, work_args)
         self._futures.append(future)
 
-
     def _process_work_item(self, work_type: str, work_args: Tuple[Any, ...]) -> Tuple[str, bool, str]:
         """
         Process a work item and return its result.
@@ -75,10 +75,9 @@ def _process_work_item(self, work_type: str, work_args: Tuple[Any, ...]) -> Tupl
         """
 
         root_url, bearer, job_id = work_args
-        
+
         try:
             if work_type == self.WORK_TYPE_START_JOB:
-                
                 conn = openeo.connect(root_url)
                 if bearer:
                     conn.authenticate_bearer_token(bearer)
@@ -107,12 +106,13 @@ def process_futures(self, stats: Dict[str, int]) -> None:
         """
 
         if not self._futures:
-            return # no futures to process
-        
+            return  # no futures to process
+
         done, _ = concurrent.futures.wait(
-                    self._futures,
-                    timeout=0,
-                    return_when=concurrent.futures.FIRST_COMPLETED)
+            self._futures,
+            timeout=0,
+            return_when=concurrent.futures.FIRST_COMPLETED,
+        )
 
         completed = []
 
@@ -128,7 +128,7 @@ def process_futures(self, stats: Dict[str, int]) -> None:
             except Exception as e:
                 _log.exception(f"Unexpected error processing future: {e}")
             completed.append(future)
-        
+
         _log.info(f"Processed {len(completed)} jobs")
         # Remove processed futures
         for future in completed:
@@ -170,7 +170,7 @@ def _validate_work_args(self, work_type: str, work_args: Tuple[Any, ...]) -> Non
             error_msg = f"job_id must be a string, got {job_id}"
             _log.error(error_msg)
             raise TypeError(error_msg)
-        
+
     def shutdown(self):
         """
         Shut down the thread pool gracefully.
@@ -180,4 +180,4 @@ def shutdown(self):
             - It waits for all pending tasks to complete before shutting down.
         """
         _log.info("Shutting down worker thread pool")
-        self._executor.shutdown(wait=True)
+        self._executor.shutdown(wait=True)

tests/extra/job_management/test_job_management.py

@@ -6,8 +6,9 @@
 import threading
 from pathlib import Path
 from time import sleep
-from typing import  Union
+from typing import Union
 from unittest import mock
+
 import dirty_equals
 import geopandas
 
@@ -364,8 +365,6 @@ def start_worker_thread():
             for filename in ["job-results.json", f"job_{job_id}.json", "result.data"]
         }
 
-        
-
     def test_on_error_log(self, tmp_path, requests_mock):
         backend = "http://foo.test"
         requests_mock.get(backend, json={"api_version": "1.1.0"})
@@ -598,8 +597,7 @@ def get_status(job_id, current_status):
             job_manager.run_jobs(df=df, start_job=self._create_year_job, job_db=job_db_path)
 
         final_df = CsvJobDatabase(job_db_path).read()
-        assert dirty_equals.IsPartialDict(id="job-2024", status=expected_status
-                ) == final_df.iloc[0].to_dict()
+        assert dirty_equals.IsPartialDict(id="job-2024", status=expected_status) == final_df.iloc[0].to_dict()
 
         assert dummy_backend_foo.batch_jobs == {
             "job-2024": {

tests/extra/job_management/test_thread_worker.py

@@ -1,23 +1,27 @@
-import pytest
+import logging
 import time
 from collections import defaultdict
-import logging
+
+import pytest
+
 from openeo.extra.job_management.thread_worker import _JobManagerWorkerThreadPool
 
+
 # Fixture for creating and cleaning up the worker thread pool
 @pytest.fixture
 def worker_pool():
     worker_pool = _JobManagerWorkerThreadPool()
     yield worker_pool
     worker_pool.shutdown()
 
+
 # Fixture for the shared stats dictionary
 @pytest.fixture
 def stats():
     return defaultdict(int)
 
-class TestJobManagerWorkerThreadPool:
 
+class TestJobManagerWorkerThreadPool:
     def test_worker_thread_lifecycle(self, worker_pool, caplog):
         with caplog.at_level(logging.INFO):
             assert not worker_pool._executor._shutdown
@@ -44,17 +48,12 @@ def test_start_job_success(self, worker_pool, stats, requests_mock):
             f"{backend_url}/jobs",
             json={"job_id": job_id, "status": "created"},
             status_code=201,
-            headers={"openeo-identifier": job_id}
+            headers={"openeo-identifier": job_id},
         )
         requests_mock.post(
-            f"{backend_url}/jobs/{job_id}/results",
-            json={"job_id": job_id, "status": "finished"},
-            status_code=202
-        )
-        requests_mock.get(
-            f"{backend_url}/jobs/{job_id}",
-            json={"id": job_id, "status": "finished"}
+            f"{backend_url}/jobs/{job_id}/results", json={"job_id": job_id, "status": "finished"}, status_code=202
         )
+        requests_mock.get(f"{backend_url}/jobs/{job_id}", json={"id": job_id, "status": "finished"})
         # Submit several valid jobs
         for _ in range(3):
             worker_pool.submit_work(worker_pool.WORK_TYPE_START_JOB, (backend_url, "token", job_id))
@@ -82,11 +81,10 @@ def test_process_futures_mixed_success_and_failure(self, worker_pool, stats, req
         requests_mock.post(
             f"{backend_url_success}/jobs/{job_id_success}/results",
             json={"job_id": job_id_success, "status": "finished"},
-            status_code=202
+            status_code=202,
         )
         requests_mock.get(
-            f"{backend_url_success}/jobs/{job_id_success}",
-            json={"id": job_id_success, "status": "finished"}
+            f"{backend_url_success}/jobs/{job_id_success}", json={"id": job_id_success, "status": "finished"}
         )
         # Failed job
         backend_url_failure = "https://failure.test"
@@ -105,29 +103,27 @@ def test_invalid_work_type(self, worker_pool, stats, requests_mock, caplog):
         backend_url = "https://foo.test"
         job_id = "job-123"
         requests_mock.get(backend_url, json={"api_version": "1.1.0"})
-        
+
         # Test that invalid work type raises ValueError
         with pytest.raises(ValueError) as exc_info:
             worker_pool.submit_work("invalid_work_type", (backend_url, "token", job_id))
-        
+
         assert "Unsupported work type: invalid_work_type" in str(exc_info.value)
         assert len(worker_pool._futures) == 0
 
-        
     @pytest.mark.parametrize(
-    "work_args,expected_log",
-    [
-        (("https://foo.test", "token"), "Expected 3 arguments for work type start_job, got 2"),
-        ((None, "token", "job-123"), "root_url must be a string"),
-    ]
+        "work_args,expected_log",
+        [
+            (("https://foo.test", "token"), "Expected 3 arguments for work type start_job, got 2"),
+            ((None, "token", "job-123"), "root_url must be a string"),
+        ],
     )
     def test_invalid_work_args(self, worker_pool, stats, caplog, work_args, expected_log):
-        with caplog.at_level(logging.ERROR):  
+        with caplog.at_level(logging.ERROR):
             with pytest.raises(Exception):  # Expect an exception to be raised
                 worker_pool.submit_work(worker_pool.WORK_TYPE_START_JOB, work_args)
-        
+
         # Verify the expected log message
         assert expected_log in caplog.text
         assert len(worker_pool._futures) == 0
         assert stats["job start"] == 0
-