From 37bf6e5737164e84bec9a1c7dbc03f8dcddcb5ca Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Paul=20M=C3=BCller?= <paul.mueller@biotec.tu-dresden.de>
Date: Tue, 27 Jan 2015 08:19:57 +0100
Subject: [PATCH] Added `override_count=False` keyword to
 `jobmanager.decorators.decorate_module_ProgressBar` to enable switching to
 `jobmanager.decorators.ProgressBarOverrideCount`.

---
 jobmanager/decorators.py | 86 +++++++++++++++++++++++-----------------
 1 file changed, 50 insertions(+), 36 deletions(-)

diff --git a/jobmanager/decorators.py b/jobmanager/decorators.py
index ea4c4eb..31d2ce4 100644
--- a/jobmanager/decorators.py
+++ b/jobmanager/decorators.py
@@ -156,41 +156,6 @@ class ProgressBar(object):
             return self.func(**callargs)
 
 
-
-def decorate_module_ProgressBar(module, **kwargs):
-    """ Decorates all decoratable functions in a module with a
-    ProgressBar.
-    
-    You can prevent wrapping of a function by not specifying the keyword
-    arguments as defined in `jobmanager.jobmanager.validCountKwargs` or
-    by defining a function `_jm_decorate_{func}".
-    
-    **kwargs are keyword arguments for ProgressBar
-    
-    Note that decorating all functions in a module might lead to
-    strange behavior of the progress bar for nested functions.
-    """
-    vdict = module.__dict__
-    for key in list(vdict.keys()):
-        if hasattr(vdict[key], "__call__"):
-            if getCountKwargs(vdict[key]) is not None:
-                newid = "_jm_decorate_{}".format(key)
-                if hasattr(module, newid):
-                    warings.warn("Wrapping of {} prevented by module.".
-                                 format(key))
-                else:
-                    # copy old function
-                    setattr(module, newid, vdict[key])
-                    # create new function
-                    wrapper = ProgressBar(getattr(module, newid), **kwargs)
-                    # set new function
-                    setattr(module, key, wrapper)
-                    if (kwargs.has_key("verbose") and
-                        kwargs["verbose"] > 0):
-                        print("Jobmanager wrapped {}.{}".format(
-                                                  module.__name__, key))
-
-
 class ProgressBarOverrideCount(ProgressBar):
     def __call__(self, *args, **kwargs):
         """ Calls `func` - previously defined in `__init__`.
@@ -237,4 +202,53 @@ class ProgressBarOverrideCount(ProgressBar):
                                   **self.kwargs) as pb:
             pb.start()
             return self.func(**callargs)    
-    
\ No newline at end of file
+
+    
+    
+def decorate_module_ProgressBar(module, override_count=False, **kwargs):
+    """ Decorates all decoratable functions in a module with a
+    ProgressBar.
+    
+    You can prevent wrapping of a function by not specifying the keyword
+    arguments as defined in `jobmanager.jobmanager.validCountKwargs` or
+    by defining a function `_jm_decorate_{func}".
+
+    Parameters
+    ----------
+    module : Python module
+        The module whose functions should be decorated.
+    override_count : bool
+        Override the default "count values" of the functions. If this
+        is set to `True`, then `decorators.ProgressBarOverrideCount` 
+        will be used instead of `decorators.ProgressBar`.
+    **kwargs : dict
+        Keyword arguments to the ProgressBar.
+    
+    Notes
+    -----
+    Decorating nested functions in a module might lead to unexpected
+    behavior.
+    """
+    vdict = module.__dict__
+    for key in list(vdict.keys()):
+        if hasattr(vdict[key], "__call__"):
+            if getCountKwargs(vdict[key]) is not None:
+                newid = "_jm_decorate_{}".format(key)
+                if hasattr(module, newid):
+                    warings.warn("Wrapping of {} prevented by module.".
+                                 format(key))
+                else:
+                    # copy old function
+                    setattr(module, newid, vdict[key])
+                    # create new function
+                    if override_count:
+                        wrapper = ProgressBarOverrideCount(getattr(module, newid), **kwargs)
+                    else:    
+                        wrapper = ProgressBar(getattr(module, newid), **kwargs)
+                    # set new function
+                    setattr(module, key, wrapper)
+                    if (kwargs.has_key("verbose") and
+                        kwargs["verbose"] > 0):
+                        print("Jobmanager wrapped {}.{}".format(
+                                                  module.__name__, key))
+