FFT_tools: try to import pyplot up front

[FFT-tools.git] / FFT_tools.py

diff --git a/FFT_tools.py b/FFT_tools.py
index 9a6041ae70c9776331f03508beaee9b11a881cdd..95b998162df757f6009613fe838e280dc1a80f23 100644 (file)
--- a/FFT_tools.py
+++ b/FFT_tools.py
@@ -30,6 +30,11 @@ import logging as _logging
 import unittest as _unittest
 
 import numpy as _numpy
+try:
+    import matplotlib.pyplot as _pyplot
+except ImportError as e:
+    _pyplot = None
+    _pyplot_import_error = e
 
 
 __version__ = '0.4'
@@ -475,6 +480,8 @@ class TestUnitaryRFFT (_unittest.TestCase):
             expected[i] = 1.0 / _numpy.abs(a) * _numpy.sinc(f / a)
 
         if TEST_PLOTS:
+            if _pyplot is None:
+                raise _pyplot_import_error
             figure = _pyplot.figure()
             time_axes = figure.add_subplot(2, 1, 1)
             time_axes.plot(_numpy.arange(0, dt * samples, dt), x)
@@ -539,6 +546,8 @@ class TestUnitaryRFFT (_unittest.TestCase):
                 1.0 / a, _numpy.pi * f)
 
         if TEST_PLOTS:
+            if _pyplot is None:
+                raise _pyplot_import_error
             figure = _pyplot.figure()
             time_axes = figure.add_subplot(2, 1, 1)
             time_axes.plot(_numpy.arange(0, dt * samples, dt), x)
@@ -604,6 +613,8 @@ class TestUnitaryPowerSpectrum (_unittest.TestCase):
             'The total power should be {} ({})'.format(Pexp, P))
 
         if TEST_PLOTS:
+            if _pyplot is None:
+                raise _pyplot_import_error
             figure = _pyplot.figure()
             time_axes = figure.add_subplot(2, 1, 1)
             time_axes.plot(
@@ -657,6 +668,8 @@ class TestUnitaryPowerSpectrum (_unittest.TestCase):
                 expected_amp, power[0]))
 
         if TEST_PLOTS:
+            if _pyplot is None:
+                raise _pyplot_import_error
             figure = _pyplot.figure()
             time_axes = figure.add_subplot(2, 1, 1)
             time_axes.plot(
@@ -726,6 +739,8 @@ class TestUnitaryPowerSpectrum (_unittest.TestCase):
                 expected[0], power[0]))
 
         if TEST_PLOTS:
+            if _pyplot is None:
+                raise _pyplot_import_error
             figure = _pyplot.figure()
             time_axes = figure.add_subplot(2, 1, 1)
             time_axes.plot(
@@ -779,6 +794,8 @@ class TestUnitaryAvgPowerSpectrum (_unittest.TestCase):
             'The total power should be {} ({})'.format(Pexp, P))
 
         if TEST_PLOTS:
+            if _pyplot is None:
+                raise _pyplot_import_error
             figure = _pyplot.figure()
             time_axes = figure.add_subplot(2, 1, 1)
             time_axes.plot(