Repository: RhodiumGroup/rhg_compute_tools

Issue:
Add client.map helper functions
## Goals

Provide functions for more semantic use of client.map and other dask workflows
* [x] expand/collapse for multi-arg function handling
* [x] provide [jrnr](https://github.com/ClimateImpactLab/jrnr)-like parameter iteration tools

## Prototype

```python
import functools, itertools

def expand(func):
    '''
    Decorator to expand an (args, kwargs) tuple in function calls
    
    Intended for use with the :py:func:`collapse` function
    
    Parameters
    ----------
    func : function
        Function to have arguments expanded. Func can have any
        number of positional and keyword arguments.
    
    Returns
    -------
    wrapped : function
        Wrapped version of ``func`` which accepts a single 
        ``(args, kwargs)`` tuple.
    
    Examples
    --------
    
    .. code-block:: python

        In [1]: @expand
           ...: def my_func(a, b, exp=1):
           ...:     return (a * b)**exp
           ...:

        In [2]: my_func((2, 3))
        6
        
        In [3]: my_func((2, 3, 2))
        36
        
        In [4]: my_func(tuple([]), {'b': 4, 'c': 2, 'a': 1})
        16
        
    This function can be used in combination with the ``collapse`` helper function,
    which allows more natural parameter calls

    .. code-block:: python
        
        In [5]: my_func(collapse(2, 3, exp=2))
        36

    These can then be paired to enable many parameterized function calls:

    .. code-block:: python

        In [6]: func_calls = [collapse(a, a+1, exp=a) for a in range(5)]

        In [7]: map(my_func, func_calls)
        [1, 2, 36, 1728, 160000]

    '''

    @functools.wraps(func)
    def inner(ak):
        return func(*ak[0], **ak[1])
    return inner


def collapse(*args, **kwargs):
    '''
    Collapse positional and keyword arguments into an (args, kwargs) tuple
    
    Intended for use with the :py:func:`expand` decorator
    
    Parameters
    ----------
    *args
        Variable length argument list.
    **kwargs
        Arbitrary keyword arguments.
        
    Returns
    -------
    args : tuple
        Positional arguments tuple
    kwargs : dict
        Keyword argument dictionary
    '''
    return (args, kwargs)


def collapse_product(*args, **kwargs):
    '''
    
    Parameters
    ----------
    
    *args
        Variable length list of iterables
    **kwargs
        Keyword arguments, whose values must be iterables
        
    Returns
    -------
    iterator
        Generator with collapsed arguments
        
    See Also
    --------
    
    Function :py:func:`collapse`
    
    Examples
    --------
    
    .. code-block:: python

        In [1]: @expand
           ...: def my_func(a, b, exp=1):
           ...:     return (a * b)**exp
           ...:
    
        In [3]: product_args = collapse_product(
           ...:     [0, 1, 2],
           ...:     [0.5, 2],
           ...:     exp=[0, 1])
        
        In [4]: list(product_args)  # doctest: NORMALIZE_WHITESPACE
           [
               ((0, 0.5), {'exp': 0}),
               ((0, 0.5), {'exp': 1}),
               ((0, 2), {'exp': 0}),
               ((0, 2), {'exp': 1}),
               ((1, 0.5), {'exp': 0}),
               ((1, 0.5), {'exp': 1}),
               ((1, 2), {'exp': 0}),
               ((1, 2), {'exp': 1}),
               ((2, 0.5), {'exp': 0}),
               ((2, 0.5), {'exp': 1}),
               ((2, 2), {'exp': 0}),
               ((2, 2), {'exp': 1})]
               
        In [5]: list(map(my_func, product_args))
        [1.0, 0.0, 1, 0, 1.0, 0.5, 1, 2, 1.0, 1.0, 1, 4]
    '''
    num_args = len(args)
    kwarg_keys = list(kwargs.keys())
    kwarg_vals = [kwargs[k] for k in kwarg_keys]
    
    format_iterations = lambda x: (tuple(x[:num_args]), dict(zip(kwarg_keys, x[num_args:])))
    
    return map(format_iterations, itertools.product(*args, *kwarg_vals))
    
```

Comments:

Patch:
diff --git a/.travis.yml b/.travis.yml
index a67a703..f07100e 100644
--- a/.travis.yml
+++ b/.travis.yml
@@ -46,6 +46,12 @@ language: python
 python:
 - 3.6
 script:
-- if [[ "$TEST_ENV" == "conda" ]]; then pytest && sphinx-apidoc -o docs rhg_compute_tools
-  && sphinx-build -W -b html -d docs/_build/doctrees docs/. docs/_build/html; else
-  tox; fi;
+- if [[ "$TEST_ENV" == "conda" ]];
+  then
+    pytest &&
+    rm docs/rhg_compute_tools*.rst &&
+    sphinx-apidoc -o docs rhg_compute_tools &&
+    sphinx-build -W -b html -d docs/_build/doctrees docs/. docs/_build/html;
+  else
+    tox;
+  fi;
diff --git a/docs/.gitignore b/docs/.gitignore
index a11931c..e69de29 100644
--- a/docs/.gitignore
+++ b/docs/.gitignore
@@ -1,2 +0,0 @@
-/rhg_compute_tools.rst
-/rhg_compute_tools.*.rst
diff --git a/docs/conf.py b/docs/conf.py
index c70b86d..47013ef 100755
--- a/docs/conf.py
+++ b/docs/conf.py
@@ -46,6 +46,7 @@ extensions = [
     'sphinx.ext.viewcode',
     'sphinx.ext.extlinks',
     'sphinx.ext.napoleon',
+    'sphinx.ext.doctest',
     'sphinx.ext.intersphinx']
 
 extlinks = {
diff --git a/docs/rhg_compute_tools.rst b/docs/rhg_compute_tools.rst
index d01b7ea..a3efd4e 100644
--- a/docs/rhg_compute_tools.rst
+++ b/docs/rhg_compute_tools.rst
@@ -1,6 +1,9 @@
-API
+rhg\_compute\_tools package
 ===========================
 
+Submodules
+----------
+
 rhg\_compute\_tools.gcs module
 ------------------------------
 
@@ -17,6 +20,14 @@ rhg\_compute\_tools.kubernetes module
     :undoc-members:
     :show-inheritance:
 
+rhg\_compute\_tools.utils module
+--------------------------------
+
+.. automodule:: rhg_compute_tools.utils
+    :members:
+    :undoc-members:
+    :show-inheritance:
+
 
 Module contents
 ---------------
diff --git a/rhg_compute_tools/utils.py b/rhg_compute_tools/utils.py
new file mode 100644
index 0000000..eb66010
--- /dev/null
+++ b/rhg_compute_tools/utils.py
@@ -0,0 +1,151 @@
+import functools
+import itertools
+
+
+def expand(func):
+    '''
+    Decorator to expand an (args, kwargs) tuple in function calls
+
+    Intended for use with the :py:func:`collapse` function
+
+    Parameters
+    ----------
+    func : function
+        Function to have arguments expanded. Func can have any
+        number of positional and keyword arguments.
+
+    Returns
+    -------
+    wrapped : function
+        Wrapped version of ``func`` which accepts a single
+        ``(args, kwargs)`` tuple.
+
+    Examples
+    --------
+
+    .. code-block:: python
+
+        >>> @expand
+        ... def my_func(a, b, exp=1):
+        ...     return (a * b)**exp
+        ...
+
+        >>> my_func(((2, 3), {}))
+        6
+
+        >>> my_func(((2, 3, 2), {}))
+        36
+
+        >>> my_func((tuple([]), {'b': 4, 'exp': 2, 'a': 1}))
+        16
+
+    This function can be used in combination with the ``collapse`` helper
+    function, which allows more natural parameter calls
+
+    .. code-block:: python
+
+        >>> my_func(collapse(2, 3, exp=2))
+        36
+
+    These can then be paired to enable many parameterized function calls:
+
+    .. code-block:: python
+
+        >>> func_calls = [collapse(a, a+1, exp=a) for a in range(5)]
+
+        >>> list(map(my_func, func_calls))
+        [1, 2, 36, 1728, 160000]
+
+    '''
+
+    @functools.wraps(func)
+    def inner(ak):
+        return func(*ak[0], **ak[1])
+    return inner
+
+
+def collapse(*args, **kwargs):
+    '''
+    Collapse positional and keyword arguments into an (args, kwargs) tuple
+
+    Intended for use with the :py:func:`expand` decorator
+
+    Parameters
+    ----------
+    *args
+        Variable length argument list.
+    **kwargs
+        Arbitrary keyword arguments.
+
+    Returns
+    -------
+    args : tuple
+        Positional arguments tuple
+    kwargs : dict
+        Keyword argument dictionary
+    '''
+    return (args, kwargs)
+
+
+def collapse_product(*args, **kwargs):
+    '''
+
+    Parameters
+    ----------
+
+    *args
+        Variable length list of iterables
+    **kwargs
+        Keyword arguments, whose values must be iterables
+
+    Returns
+    -------
+    iterator
+        Generator with collapsed arguments
+
+    See Also
+    --------
+
+    Function :py:func:`collapse`
+
+    Examples
+    --------
+
+    .. code-block:: python
+
+        >>> @expand
+        ... def my_func(a, b, exp=1):
+        ...     return (a * b)**exp
+        ...
+
+        >>> product_args = list(collapse_product(
+        ...     [0, 1, 2],
+        ...     [0.5, 2],
+        ...     exp=[0, 1]))
+
+        >>> product_args  # doctest: +NORMALIZE_WHITESPACE
+        [((0, 0.5), {'exp': 0}),
+         ((0, 0.5), {'exp': 1}),
+         ((0, 2), {'exp': 0}),
+         ((0, 2), {'exp': 1}),
+         ((1, 0.5), {'exp': 0}),
+         ((1, 0.5), {'exp': 1}),
+         ((1, 2), {'exp': 0}),
+         ((1, 2), {'exp': 1}),
+         ((2, 0.5), {'exp': 0}),
+         ((2, 0.5), {'exp': 1}),
+         ((2, 2), {'exp': 0}),
+         ((2, 2), {'exp': 1})]
+
+        >>> list(map(my_func, product_args))
+        [1.0, 0.0, 1, 0, 1.0, 0.5, 1, 2, 1.0, 1.0, 1, 4]
+    '''
+    num_args = len(args)
+    kwarg_keys = list(kwargs.keys())
+    kwarg_vals = [kwargs[k] for k in kwarg_keys]
+
+    format_iterations = lambda x: (
+        tuple(x[:num_args]),
+        dict(zip(kwarg_keys, x[num_args:])))
+
+    return map(format_iterations, itertools.product(*args, *kwarg_vals))


Test patch:
diff --git a/tests/test_rhg_compute_tools.py b/tests/test_rhg_compute_tools.py
index daaaa95..2b98710 100644
--- a/tests/test_rhg_compute_tools.py
+++ b/tests/test_rhg_compute_tools.py
@@ -6,7 +6,7 @@
 import pytest
 import inspect
 
-from rhg_compute_tools import gcs, kubernetes
+from rhg_compute_tools import gcs, kubernetes, utils
 
 
 @pytest.fixture
@@ -28,7 +28,7 @@ def test_content(response):
 
 
 def test_docstrings():
-    for mod in [gcs, kubernetes]:
+    for mod in [gcs, kubernetes, utils]:
         for cname, component in mod.__dict__.items():
             if cname.startswith('_'):
                 continue