cleanup test and fix dig-sim

Author: engineerjoe440

electricpy/sim.py

@@ -79,7 +79,8 @@ def digifiltersim(fin, filter, freqs, NN=1000, dt=0.01, title="",
     figsize:    tuple, optional
                 The figure dimensions for each subplot, default=None
     """
-    if (figsize != None): _plt.figure(figsize=figsize)
+    if figsize is not None:
+        _plt.figure(figsize=figsize)
     flen = len(freqs)
     for i in range(flen):
         # Gather frequency
@@ -94,7 +95,7 @@ def digifiltersim(fin, filter, freqs, NN=1000, dt=0.01, title="",
             x[k] = fin(k * dt, freq)
 
         # Identify how many rows were provided
-        sz = filter.size
+        sz = len(filter) if isinstance(filter, (tuple, list)) else filter.size
         if (sz < 5):
             raise ValueError("ERROR: Too few filter arguments provided. " +
                              "Refer to documentation for proper format.")
@@ -131,10 +132,11 @@ def digifiltersim(fin, filter, freqs, NN=1000, dt=0.01, title="",
         _plt.plot(ytime, 'k', label="Output")
         _plt.title(title)
         _plt.grid(which='both')
-        if legend: _plt.legend(title="Frequency = " + str(freq) + "Hz")
-        if xlim != False:
+        if legend:
+            _plt.legend(title="Frequency = " + str(freq) + "Hz")
+        if not xlim:
             _plt.xlim(xlim)
-        elif xmxscale != None:
+        elif xmxscale is not None:
             _plt.xlim((0, xmxscale / (freq * dt)))
 
     _plt.tight_layout()

test/test_sim.py

@@ -7,10 +7,14 @@
 
 
 def test_digifiltersim_basic():
-    fin = lambda t, f: np.sin(2 * np.pi * f * t)
-    filter_row = [0.0, 0.0, 1.0, 0.0, 0.0]
-
-    plot_module = sim.digifiltersim(fin, filter_row, freqs=[1], NN=10, dt=0.1, legend=False)
+    plot_module = sim.digifiltersim(
+        lambda t, f: np.sin(2 * np.pi * f * t),
+        [0.0, 0.0, 1.0, 0.0, 0.0],
+        freqs=[1],
+        NN=10,
+        dt=0.1,
+        legend=False
+    )
 
     assert plot_module is not None