Add normalize parameter to Audio.

matangover · matangover · commit 7044ff6830bc · 2019-03-15T11:42:28.000-04:00
diff --git a/IPython/lib/display.py b/IPython/lib/display.py
@@ -54,6 +54,12 @@ class Audio(DisplayObject):
     autoplay : bool
         Set to True if the audio should immediately start playing.
         Default is `False`.
+    normalize : bool
+        Whether audio should be normalized (rescaled) to the maximum possible
+        range. Default is `True`. When set to `False`, `data` must be between
+        -1 and 1 (inclusive), otherwise an error is raised.
+        Applies only when `data` is a list or array of samples; other types of
+        audio are never normalized.
 
     Examples
     --------
@@ -83,7 +89,7 @@ class Audio(DisplayObject):
     """
     _read_flags = 'rb'
 
-    def __init__(self, data=None, filename=None, url=None, embed=None, rate=None, autoplay=False):
+    def __init__(self, data=None, filename=None, url=None, embed=None, rate=None, autoplay=False, normalize=True):
         if filename is None and url is None and data is None:
             raise ValueError("No audio data found. Expecting filename, url, or data.")
         if embed is False and url is None:
@@ -99,7 +105,7 @@ def __init__(self, data=None, filename=None, url=None, embed=None, rate=None, au
         if self.data is not None and not isinstance(self.data, bytes):
             if rate is None:
                 raise ValueError("rate must be specified when data is a numpy array or list of audio samples.")
-            self.data = Audio._make_wav(data, rate)
+            self.data = Audio._make_wav(data, rate, normalize)
 
     def reload(self):
         """Reload the raw data from file or URL."""
@@ -115,16 +121,16 @@ def reload(self):
             self.mimetype = "audio/wav"
 
     @staticmethod
-    def _make_wav(data, rate):
+    def _make_wav(data, rate, normalize):
         """ Transform a numpy array to a PCM bytestring """
         import struct
         from io import BytesIO
         import wave
 
         try:
-            scaled, nchan = Audio._validate_and_normalize_with_numpy(data)
+            scaled, nchan = Audio._validate_and_normalize_with_numpy(data, normalize)
         except ImportError:
-            scaled, nchan = Audio._validate_and_normalize_without_numpy(data)
+            scaled, nchan = Audio._validate_and_normalize_without_numpy(data, normalize)
 
         fp = BytesIO()
         waveobj = wave.open(fp,mode='wb')
@@ -139,7 +145,7 @@ def _make_wav(data, rate):
         return val
 
     @staticmethod
-    def _validate_and_normalize_with_numpy(data):
+    def _validate_and_normalize_with_numpy(data, normalize):
         import numpy as np
 
         data = np.array(data, dtype=float)
@@ -154,21 +160,32 @@ def _validate_and_normalize_with_numpy(data):
             data = data.T.ravel()
         else:
             raise ValueError('Array audio input must be a 1D or 2D array')
-        scaled = np.int16(data/np.max(np.abs(data))*32767).tolist()
+        
+        max_abs_value = np.max(np.abs(data))
+        normalization_factor = Audio._get_normalization_factor(max_abs_value, normalize)
+        scaled = np.int16(data / normalization_factor * 32767).tolist()
         return scaled, nchan
 
+
     @staticmethod
-    def _validate_and_normalize_without_numpy(data):
+    def _validate_and_normalize_without_numpy(data, normalize):
         try:
-            maxabsvalue = float(max([abs(x) for x in data]))
+            max_abs_value = float(max([abs(x) for x in data]))
         except TypeError:
             raise TypeError('Only lists of mono audio are '
                 'supported if numpy is not installed')
 
-        scaled = [int(x/maxabsvalue*32767) for x in data]
+        normalization_factor = Audio._get_normalization_factor(max_abs_value, normalize)
+        scaled = [int(x / normalization_factor * 32767) for x in data]
         nchan = 1
         return scaled, nchan
 
+    @staticmethod
+    def _get_normalization_factor(max_abs_value, normalize):
+        if not normalize and max_abs_value > 1:
+            raise ValueError('Audio data must be between -1 and 1 when normalize=False.')
+        return max_abs_value if normalize else 1
+
     def _data_and_metadata(self):
         """shortcut for returning metadata with url information, if defined"""
         md = {}
diff --git a/IPython/lib/tests/test_display.py b/IPython/lib/tests/test_display.py
@@ -19,7 +19,10 @@
     import pathlib
 except ImportError:
     pass
-from unittest import mock
+from unittest import TestCase, mock
+import struct
+import wave
+from io import BytesIO
 
 # Third-party imports
 import nose.tools as nt
@@ -184,25 +187,66 @@ def test_audio_from_file():
     path = pjoin(dirname(__file__), 'test.wav')
     display.Audio(filename=path)
 
-def test_audio_from_numpy_array():
-    display.Audio(get_test_tone(), rate=44100)
-
-def test_audio_from_list_without_numpy():
-    # Simulate numpy not installed.
-    with mock.patch('numpy.array', side_effect=ImportError):
-        display.Audio(list(get_test_tone()), rate=44100)
-
-def test_audio_from_list_without_numpy_raises_for_nested_list():
-    # Simulate numpy not installed.
-    with mock.patch('numpy.array', side_effect=ImportError):
+class TestAudioDataWithNumpy(TestCase):
+    def test_audio_from_numpy_array(self):
+        test_tone = get_test_tone()
+        audio = display.Audio(test_tone, rate=44100)
+        nt.assert_equal(len(read_wav(audio.data)), len(test_tone))
+
+    def test_audio_from_list(self):
+        test_tone = get_test_tone()
+        audio = display.Audio(list(test_tone), rate=44100)
+        nt.assert_equal(len(read_wav(audio.data)), len(test_tone))
+
+    def test_audio_from_numpy_array_without_rate_raises(self):
+        nt.assert_raises(ValueError, display.Audio, get_test_tone())
+
+    def test_audio_data_normalization(self):
+        expected_max_value = numpy.iinfo(numpy.int16).max
+        for scale in [1, 0.5, 2]:
+            audio = display.Audio(get_test_tone(scale), rate=44100)
+            actual_max_value = numpy.max(numpy.abs(read_wav(audio.data)))
+            nt.assert_equal(actual_max_value, expected_max_value)
+
+    def test_audio_data_without_normalization(self):
+        max_int16 = numpy.iinfo(numpy.int16).max
+        for scale in [1, 0.5, 0.2]:
+            test_tone = get_test_tone(scale)
+            test_tone_max_abs = numpy.max(numpy.abs(test_tone))
+            expected_max_value = int(max_int16 * test_tone_max_abs)
+            audio = display.Audio(test_tone, rate=44100, normalize=False)
+            actual_max_value = numpy.max(numpy.abs(read_wav(audio.data)))
+            nt.assert_equal(actual_max_value, expected_max_value)
+
+    def test_audio_data_without_normalization_raises_for_invalid_data(self):
+        nt.assert_raises(
+            ValueError,
+            lambda: display.Audio([1.001], rate=44100, normalize=False))
+        nt.assert_raises(
+            ValueError,
+            lambda: display.Audio([-1.001], rate=44100, normalize=False))
+
+def simulate_numpy_not_installed():
+    return mock.patch('numpy.array', mock.MagicMock(side_effect=ImportError))
+
+@simulate_numpy_not_installed()
+class TestAudioDataWithoutNumpy(TestAudioDataWithNumpy):
+    # All tests from `TestAudioDataWithNumpy` are inherited.
+
+    def test_audio_raises_for_nested_list(self):
         stereo_signal = [list(get_test_tone())] * 2
-        nt.assert_raises(TypeError, lambda: display.Audio(stereo_signal, rate=44100))
-
-def test_audio_from_numpy_array_without_rate_raises():
-    nt.assert_raises(ValueError, display.Audio, get_test_tone())
-
-def get_test_tone():
-    return numpy.sin(2 * numpy.pi * 440 * numpy.linspace(0, 1, 44100))
+        nt.assert_raises(
+            TypeError,
+            lambda: display.Audio(stereo_signal, rate=44100))
+
+def get_test_tone(scale=1):
+    return numpy.sin(2 * numpy.pi * 440 * numpy.linspace(0, 1, 44100)) * scale
+
+def read_wav(data):
+    with wave.open(BytesIO(data)) as wave_file:
+        wave_data = wave_file.readframes(wave_file.getnframes())
+        num_samples = wave_file.getnframes() * wave_file.getnchannels()
+        return struct.unpack('<%sh' % num_samples, wave_data)
 
 def test_code_from_file():
     c = display.Code(filename=__file__)
