Add wrapper to pass args and kwargs to the function and improve test coverage

Change-Id: I18b548b75b20eb8dc963d23d1dcdc575faeb903a

Author: adrien-berchet

bluepyparallel/__init__.py

@@ -1,3 +1,4 @@
 """BluePyParallel functions."""
-from .evaluator import evaluate  # noqa
-from .parallel import init_parallel_factory  # noqa
+from bluepyparallel.evaluator import evaluate  # noqa
+from bluepyparallel.parallel import init_parallel_factory  # noqa
+from bluepyparallel.version import VERSION as __version__  # noqa

bluepyparallel/database.py

@@ -1,4 +1,4 @@
-"""Module"""
+"""Module used to provide a simple API to a database in which the results are stored."""
 import re
 
 import pandas as pd
@@ -9,12 +9,28 @@
 from sqlalchemy import schema
 from sqlalchemy import select
 from sqlalchemy.engine.reflection import Inspector
+from sqlalchemy.exc import OperationalError
 from sqlalchemy_utils import create_database
 from sqlalchemy_utils import database_exists
 
+try:  # pragma: no cover
+    import psycopg2
+
+    with_psycopg2 = True
+except ImportError:
+    with_psycopg2 = False
+
 
 class DataBase:
-    """A database API using SQLAlchemy."""
+    """A simple API to manage the database in which the results are inserted using SQLAlchemy.
+
+    Args:
+        url (str): The URL of the database following the RFC-1738 format (
+            https://docs.sqlalchemy.org/en/14/core/engines.html#database-urls)
+        create (bool): If set to True, the database will be automatically created by the
+            constructor.
+        args and kwargs: They will be passed to the :func:`sqlalchemy.create_engine` function.
+    """
 
     index_col = "df_index"
     _url_pattern = r"[a-zA-Z0-9_\-\+]+://.*"
@@ -25,25 +41,36 @@ def __init__(self, url, *args, create=False, **kwargs):
 
         self.engine = create_engine(url, *args, **kwargs)
 
-        if create and not database_exists(self.engine.url):
+        if create and not self.db_exists():
             create_database(self.engine.url)
 
-        self.connection = self.engine.connect()
+        self._connection = None
         self.metadata = None
         self.table = None
 
     def __del__(self):
+        """Close the connection and the engine to the database."""
         self.connection.close()
+        self.engine.dispose()
+
+    @property
+    def connection(self):
+        """Get a connection to the database."""
+        if self._connection is None:
+            self._connection = self.engine.connect()
+        return self._connection
 
     def get_url(self):
+        """Get the URL of the database."""
         return self.engine.url
 
     def create(self, df, table_name=None, schema_name=None):
+        """Create a table in the database in which the results will be written."""
         if table_name is None:
             table_name = "df"
         if schema_name is not None and schema_name not in self.connection.dialect.get_schema_names(
             self.connection
-        ):
+        ):  # pragma: no cover
             self.connection.execute(schema.CreateSchema(schema_name))
         new_df = df.loc[[]]
         new_df.to_sql(
@@ -55,11 +82,25 @@ def create(self, df, table_name=None, schema_name=None):
         )
         self.reflect(table_name, schema_name)
 
+    def db_exists(self):
+        """Check that the server and the database exist."""
+        if with_psycopg2:  # pragma: no cover
+            exceptions = (OperationalError, psycopg2.OperationalError)
+        else:
+            exceptions = (OperationalError,)
+
+        try:
+            return database_exists(self.engine.url)
+        except exceptions:  # pragma: no cover
+            return False
+
     def exists(self, table_name, schema_name=None):
+        """Test that the table exists in the database."""
         inspector = Inspector.from_engine(self.engine)
         return table_name in inspector.get_table_names(schema=schema_name)
 
     def reflect(self, table_name, schema_name=None):
+        """Reflect the table from the database."""
         self.metadata = MetaData()
         self.table = Table(
             table_name,
@@ -70,10 +111,12 @@ def reflect(self, table_name, schema_name=None):
         )
 
     def load(self):
+        """Load the table data from the database."""
         query = select([self.table])
         return pd.read_sql(query, self.connection, index_col=self.index_col)
 
     def write(self, row_id, result=None, exception=None, **input_values):
+        """Write a result entry or an exception into the table."""
         if result is not None:
             vals = result
         elif exception is not None:

bluepyparallel/evaluator.py

@@ -84,7 +84,7 @@ def evaluate(
 
     # Set default new columns
     if new_columns is None:
-        new_columns = [["data", ""]]
+        new_columns = []
 
     # Setup internal and new columns
     to_evaluate["exception"] = None
@@ -141,10 +141,9 @@ def evaluate(
     # Split the data into rows
     arg_list = list(to_evaluate.loc[task_ids, df.columns].to_dict("index").items())
 
+    res = []
     try:
-        res = []
-
-        # Collect the results
+        # Compute and collect the results
         for task_id, result, exception in tqdm(mapper(eval_func, arg_list), total=len(task_ids)):
             res.append(dict({"df_index": task_id, "exception": exception}, **result))
 
@@ -154,13 +153,13 @@ def evaluate(
                     task_id, result, exception, **to_evaluate.loc[task_id, df.columns].to_dict()
                 )
 
-        # Gather the results to the output DataFrame
-        res_df = pd.DataFrame(res)
-        res_df.set_index("df_index", inplace=True)
-        to_evaluate.loc[res_df.index, res_df.columns] = res_df
-
     except (KeyboardInterrupt, SystemExit) as ex:
         # To save dataframe even if program is killed
         logger.warning("Stopping mapper loop. Reason: %r", ex)
 
+    # Gather the results to the output DataFrame
+    res_df = pd.DataFrame(res)
+    res_df.set_index("df_index", inplace=True)
+    to_evaluate.loc[res_df.index, res_df.columns] = res_df
+
     return to_evaluate

bluepyparallel/parallel.py

@@ -14,20 +14,25 @@
     import dask_mpi
 
     dask_available = True
-except ImportError:
+except ImportError:  # pragma: no cover
     dask_available = False
 
 try:
     import ipyparallel
 
     ipyparallel_available = True
-except ImportError:
+except ImportError:  # pragma: no cover
     ipyparallel_available = False
 
 
 L = logging.getLogger(__name__)
 
 
+def _func_wrapper(data, func, func_args, func_kwargs):
+    """Function wrapper used to pass args and kwargs."""
+    return func(data, *func_args, **func_kwargs)
+
+
 class ParallelFactory:
     """Abstract class that should be subclassed to provide parallel functions."""
 
@@ -56,6 +61,10 @@ def get_mapper(self, batch_size=None, chunk_size=None, **kwargs):
     def shutdown(self):
         """Can be used to cleanup."""
 
+    def mappable_func(self, func, *args, **kwargs):
+        """Can be used to add args and kwargs to a function before calling the mapper."""
+        return partial(_func_wrapper, func=func, func_args=args, func_kwargs=kwargs)
+
     def _with_batches(self, mapper, func, iterable, batch_size=None):
         """Wrapper on mapper function creating batches of iterable to give to mapper.
 
@@ -95,7 +104,7 @@ def __init__(self, group=None, target=None, name=None, args=(), kwargs={}):
 
     def _get_daemon(self):  # pylint: disable=no-self-use
         """Get daemon flag"""
-        return False
+        return False  # pragma: no cover
 
     def _set_daemon(self, value):
         """Set daemon flag"""
@@ -114,7 +123,12 @@ class SerialFactory(ParallelFactory):
 
     def get_mapper(self, batch_size=None, chunk_size=None, **kwargs):
         """Get a map."""
-        return map
+
+        def _mapper(func, iterable, *func_args, **func_kwargs):
+            func = self.mappable_func(func, *func_args, **func_kwargs)
+            return self._with_batches(map, func, iterable)
+
+        return _mapper
 
 
 class MultiprocessingFactory(ParallelFactory):
@@ -134,7 +148,8 @@ def get_mapper(self, batch_size=None, chunk_size=None, **kwargs):
         """Get a NestedPool."""
         self._chunksize_to_kwargs(chunk_size, kwargs, label="chunksize")
 
-        def _mapper(func, iterable):
+        def _mapper(func, iterable, *func_args, **func_kwargs):
+            func = self.mappable_func(func, *func_args, **func_kwargs)
             return self._with_batches(
                 partial(self.pool.imap_unordered, **kwargs),
                 func,
@@ -164,12 +179,13 @@ def __init__(self, batch_size=None, chunk_size=None, profile=None, **kwargs):
 
     def get_mapper(self, batch_size=None, chunk_size=None, **kwargs):
         """Get an ipyparallel mapper using the profile name provided."""
-        if "ordered" not in kwargs:
+        if "ordered" not in kwargs:  # pragma: no cover
             kwargs["ordered"] = False
 
         self._chunksize_to_kwargs(chunk_size, kwargs)
 
-        def _mapper(func, iterable):
+        def _mapper(func, iterable, *func_args, **func_kwargs):
+            func = self.mappable_func(func, *func_args, **func_kwargs)
             return self._with_batches(
                 partial(self.lview.imap, **kwargs), func, iterable, batch_size=batch_size
             )
@@ -178,7 +194,7 @@ def _mapper(func, iterable):
 
     def shutdown(self):
         """Remove zmq."""
-        if self.rc is not None:
+        if self.rc is not None:  # pragma: no cover
             self.rc.close()
 
 
@@ -193,11 +209,11 @@ def __init__(
         """Initialize the dask factory."""
         dask_scheduler_path = scheduler_file or os.getenv(self._SCHEDULER_PATH)
         self.interactive = True
-        if dask_scheduler_path:
+        if dask_scheduler_path:  # pragma: no cover
             L.info("Connecting dask_mpi with scheduler %s", dask_scheduler_path)
-        if address:
+        if address:  # pragma: no cover
             L.info("Connecting dask_mpi with address %s", address)
-        if not dask_scheduler_path and not address:
+        if not dask_scheduler_path and not address:  # pragma: no cover
             self.interactive = False
             dask_mpi.initialize()
             L.info("Starting dask_mpi...")
@@ -213,19 +229,20 @@ def shutdown(self):
         """Close the scheduler and the cluster if it was created by the factory."""
         cluster = self.client.cluster
         self.client.close()
-        if not self.interactive:
+        if not self.interactive:  # pragma: no cover
             cluster.close()
 
     def get_mapper(self, batch_size=None, chunk_size=None, **kwargs):
         """Get a Dask mapper."""
         self._chunksize_to_kwargs(chunk_size, kwargs, label="batch_size")
 
-        def _mapper(func, iterable):
-            def _dask_mapper(func, iterable):
+        def _mapper(func, iterable, *func_args, **func_kwargs):
+            def _dask_mapper(func, iterable, **kwargs):
                 futures = self.client.map(func, iterable, **kwargs)
                 for _future, result in dask.distributed.as_completed(futures, with_results=True):
                     yield result
 
+            func = self.mappable_func(func, *func_args, **func_kwargs)
             return self._with_batches(_dask_mapper, func, iterable, batch_size=batch_size)
 
         return _mapper
@@ -245,9 +262,9 @@ def init_parallel_factory(parallel_lib, *args, **kwargs):
         None: SerialFactory,
         "multiprocessing": MultiprocessingFactory,
     }
-    if dask_available:
+    if dask_available:  # pragma: no cover
         parallel_factories["dask"] = DaskFactory
-    if ipyparallel_available:
+    if ipyparallel_available:  # pragma: no cover
         parallel_factories["ipyparallel"] = IPyParallelFactory
 
     try:

bluepyparallel/version.py

@@ -1,4 +1,3 @@
-"""Package version """
-
-VERSION = "0.0.3.dev0"
-__version__ = VERSION
+"""Package version"""
+# pragma: no cover
+VERSION = "0.0.3.dev1"

doc/source/api.rst

@@ -7,3 +7,4 @@ API Documentation
     bluepyparallel
     bluepyparallel.parallel
     bluepyparallel.evaluator
+    bluepyparallel.database

setup.py

@@ -26,7 +26,7 @@
     "sphinx-bluebrain-theme",
 ]
 
-VERSION = imp.load_source("", "bluepyparallel/version.py").__version__
+VERSION = imp.load_source("", "bluepyparallel/version.py").VERSION
 
 setup(
     name="BluePyParallel",