Simu: digitizer repacks partial callable

simulations/lib/sampling/digitizer.py

@@ -1,5 +1,5 @@
 import numpy as np
-from functools import wraps
+from functools import wraps, partial
 
 from . import sampling as smp
 from .sampler import Sampler
@@ -32,15 +32,29 @@ class Digitizer(Sampler):
         """
 
         if callable(signal):
+            # if signal is already a partial,
+            # try to rebuild it after setting the wrapper
+            if isinstance(signal, partial):
+                rebuild_partial = True
+                p_args = signal.args
+                p_kwargs = signal.keywords
+                signal = signal.func
+            else:
+                rebuild_partial = False
+
             @wraps(signal)
-            def func(*args, **kwargs):
+            def wrapper(*args, **kwargs):
                 return smp.quantise(
                         self.sample(signal(*args, **kwargs), signal_sample_frequency),
                         self.resolution,
                         self.bias
                         )
 
-            return func
+            # rebuild the partial if applicable
+            if rebuild_partial:
+                wrapper = partial(wrapper, *p_args, **p_kwargs)
+
+            return wrapper
 
         else:
             signal = np.asarray(signal)