ruff based code/import cleanup #741/#815

Author: soxofaan

openeo/extra/job_management/__init__.py

@@ -1,7 +1,13 @@
 from openeo.extra.job_management._interface import JobDatabaseInterface
-from openeo.extra.job_management._job_db import FullDataFrameJobDatabase, CsvJobDatabase, ParquetJobDatabase, create_job_db, get_job_db
-from openeo.extra.job_management.process_based import ProcessBasedJobCreator
+from openeo.extra.job_management._job_db import (
+    CsvJobDatabase,
+    FullDataFrameJobDatabase,
+    ParquetJobDatabase,
+    create_job_db,
+    get_job_db,
+)
 from openeo.extra.job_management._manager import MultiBackendJobManager
+from openeo.extra.job_management.process_based import ProcessBasedJobCreator
 
 __all__ = [
     "JobDatabaseInterface",
@@ -11,7 +17,5 @@
     "ProcessBasedJobCreator",
     "create_job_db",
     "get_job_db",
-    "MultiBackendJobManager"
+    "MultiBackendJobManager",
 ]
-
-

openeo/extra/job_management/_df_schema.py

@@ -1,14 +1,17 @@
 import dataclasses
 from typing import Any, Mapping
+
 import pandas as pd
 
+
 @dataclasses.dataclass(frozen=True)
 class _ColumnProperties:
     """Expected/required properties of a column in the job manager related dataframes"""
 
     dtype: str = "object"
     default: Any = None
 
+
 # Expected columns in the job DB dataframes.
 # TODO: make this part of public API when settled?
 # TODO: move non official statuses to seperate column (not_started, queued_for_start)
@@ -28,15 +31,16 @@ class _ColumnProperties:
     "costs": _ColumnProperties(dtype="float64"),
 }
 
+
 def _normalize(df: pd.DataFrame) -> pd.DataFrame:
-        """
-        Normalize given pandas dataframe (creating a new one):
-        ensure we have the required columns.
+    """
+    Normalize given pandas dataframe (creating a new one):
+    ensure we have the required columns.
 
-        :param df: The dataframe to normalize.
-        :return: a new dataframe that is normalized.
-        """
-        new_columns = {col: req.default for (col, req) in _COLUMN_REQUIREMENTS.items() if col not in df.columns}
-        df = df.assign(**new_columns)
+    :param df: The dataframe to normalize.
+    :return: a new dataframe that is normalized.
+    """
+    new_columns = {col: req.default for (col, req) in _COLUMN_REQUIREMENTS.items() if col not in df.columns}
+    df = df.assign(**new_columns)
 
-        return df
+    return df

openeo/extra/job_management/_interface.py

@@ -1,14 +1,9 @@
 import abc
-from typing import (
-    Callable,
-    Iterable,
-    List,
-    NamedTuple,
-    Union,
-)
+from typing import Iterable, List, Union
 
 import pandas as pd
 
+
 class JobDatabaseInterface(metaclass=abc.ABCMeta):
     """
     Interface for a database of job metadata to use with the :py:class:`MultiBackendJobManager`,
@@ -68,8 +63,3 @@ def get_by_indices(self, indices: Iterable[Union[int, str]]) -> pd.DataFrame:
         :return: DataFrame with jobs filtered by indices.
         """
         ...
-
-
-
-
-

openeo/extra/job_management/_job_db.py

@@ -1,16 +1,18 @@
 import abc
 import logging
-import shapely.errors
-import shapely.wkt
 from pathlib import Path
 from typing import Iterable, Union
+
 import pandas as pd
+import shapely.errors
+import shapely.wkt
 
+from openeo.extra.job_management._df_schema import _COLUMN_REQUIREMENTS, _normalize
 from openeo.extra.job_management._interface import JobDatabaseInterface
-from openeo.extra.job_management._df_schema import _normalize, _COLUMN_REQUIREMENTS
 
 _log = logging.getLogger(__name__)
 
+
 class FullDataFrameJobDatabase(JobDatabaseInterface):
     def __init__(self):
         super().__init__()
@@ -241,4 +243,4 @@ def create_job_db(path: Union[str, Path], df: pd.DataFrame, *, on_exists: str =
         job_db.initialize_from_df(df=df, on_exists=on_exists)
     else:
         raise NotImplementedError(f"Initialization of {type(job_db)} is not supported.")
-    return job_db
+    return job_db

openeo/extra/job_management/_manager.py

@@ -1,6 +1,5 @@
 import collections
 import contextlib
-import dataclasses
 import datetime
 import json
 import logging
@@ -25,28 +24,31 @@
 from urllib3.util import Retry
 
 from openeo import BatchJob, Connection
+from openeo.extra.job_management._df_schema import _normalize
+from openeo.extra.job_management._interface import JobDatabaseInterface
+from openeo.extra.job_management._job_db import get_job_db
 from openeo.extra.job_management._thread_worker import (
     _JobManagerWorkerThreadPool,
     _JobStartTask,
 )
-from openeo.extra.job_management._interface import JobDatabaseInterface
-from openeo.extra.job_management._job_db import get_job_db 
-from openeo.extra.job_management._df_schema import _normalize
-
 from openeo.rest import OpenEoApiError
 from openeo.rest.auth.auth import BearerAuth
 from openeo.util import deep_get, rfc3339
 
-
 _log = logging.getLogger(__name__)
 
+
 MAX_RETRIES = 50
+
+
 # Sentinel value to indicate that a parameter was not set
 _UNSET = object()
 
+
 def _start_job_default(row: pd.Series, connection: Connection, *args, **kwargs):
     raise NotImplementedError("No 'start_job' callable provided")
 
+
 class _Backend(NamedTuple):
     """Container for backend info/settings"""
 
@@ -56,8 +58,6 @@ class _Backend(NamedTuple):
     parallel_jobs: int
 
 
-
-
 class MultiBackendJobManager:
     """
     Tracker for multiple jobs on multiple backends.
@@ -125,7 +125,6 @@ def start_job(
         Added ``cancel_running_job_after`` parameter.
     """
 
-
     def __init__(
         self,
         poll_sleep: int = 60,
@@ -409,7 +408,7 @@ def run_jobs(
         assert not kwargs, f"Unexpected keyword arguments: {kwargs!r}"
 
         if isinstance(job_db, (str, Path)):
-            job_db = get_job_db(path=job_db) #TODO circular import
+            job_db = get_job_db(path=job_db)  # TODO circular import
 
         if not isinstance(job_db, JobDatabaseInterface):
             raise ValueError(f"Unsupported job_db {job_db!r}")
@@ -806,7 +805,7 @@ def _track_statuses(self, job_db: JobDatabaseInterface, stats: Optional[dict] =
         job_db.persist(active)
 
         return jobs_done, jobs_error, jobs_cancel
-    
+
 
 def _format_usage_stat(job_metadata: dict, field: str) -> str:
     value = deep_get(job_metadata, "usage", field, "value", default=0)

tests/extra/job_management/test_job_db.py

@@ -1,5 +1,5 @@
-
 import re
+
 import geopandas
 
 # TODO: can we avoid using httpretty?
@@ -18,7 +18,6 @@
     create_job_db,
     get_job_db,
 )
-
 from openeo.utils.version import ComparableVersion
 
 JOB_DB_DF_BASICS = pd.DataFrame(
@@ -378,4 +377,4 @@ def test_create_job_db(tmp_path, filename, expected):
     path = tmp_path / filename
     db = create_job_db(path=path, df=df)
     assert isinstance(db, expected)
-    assert path.exists()
+    assert path.exists()

tests/extra/job_management/test_manager.py

@@ -26,16 +26,15 @@
 
 import openeo
 from openeo import BatchJob
-from openeo.extra.job_management._manager import (
-    MAX_RETRIES,
-    MultiBackendJobManager,
-)
-
 from openeo.extra.job_management._job_db import (
     CsvJobDatabase,
     ParquetJobDatabase,
     create_job_db,
 )
+from openeo.extra.job_management._manager import (
+    MAX_RETRIES,
+    MultiBackendJobManager,
+)
 from openeo.extra.job_management._thread_worker import (
     Task,
     _JobManagerWorkerThreadPool,
@@ -604,7 +603,9 @@ def get_status(job_id, current_status):
         job_db_path = tmp_path / "jobs.csv"
 
         # Mock sleep() to not actually sleep, but skip one hour at a time
-        with mock.patch.object(openeo.extra.job_management._manager.time, "sleep", new=lambda s: time_machine.shift(60 * 60)):
+        with mock.patch.object(
+            openeo.extra.job_management._manager.time, "sleep", new=lambda s: time_machine.shift(60 * 60)
+        ):
             job_manager.run_jobs(df=df, start_job=self._create_year_job, job_db=job_db_path)
 
         final_df = CsvJobDatabase(job_db_path).read()
@@ -723,7 +724,9 @@ def get_status(job_id, current_status):
         job_db_path = tmp_path / "jobs.csv"
 
         # Mock sleep() to skip one hour at a time instead of actually sleeping
-        with mock.patch.object(openeo.extra.job_management._manager.time, "sleep", new=lambda s: time_machine.shift(60 * 60)):
+        with mock.patch.object(
+            openeo.extra.job_management._manager.time, "sleep", new=lambda s: time_machine.shift(60 * 60)
+        ):
             job_manager.run_jobs(df=df, start_job=self._create_year_job, job_db=job_db_path)
 
         final_df = CsvJobDatabase(job_db_path).read()
@@ -865,9 +868,8 @@ def test_refresh_bearer_token_before_start(
         dummy_backend_bar,
         job_manager_root_dir,
         sleep_mock,
-        requests_mock
+        requests_mock,
     ):
-
         client_id = "client123"
         client_secret = "$3cr3t"
         oidc_issuer = "https://oidc.test/"
@@ -902,6 +904,3 @@ def test_refresh_bearer_token_before_start(
         # Because of proactive+throttled token refreshing,
         # we should have 2 additional token requests now
         assert len(oidc_mock.grant_request_history) == 4
-
-
-