[Scipy-svn] r3230 - trunk/Lib/io

[scipy-svn at scipy.org]

Author: oliphant
Date: 2007-08-14 12:01:55 -0500 (Tue, 14 Aug 2007)
New Revision: 3230

Added:
   trunk/Lib/io/wavfile.py
Log:
Add wavfile.py to read and write basic .wav files.

Added: trunk/Lib/io/wavfile.py
===================================================================
--- trunk/Lib/io/wavfile.py	2007-08-13 19:16:51 UTC (rev 3229)
+++ trunk/Lib/io/wavfile.py	2007-08-14 17:01:55 UTC (rev 3230)
@@ -0,0 +1,109 @@
+import numpy
+import struct
+
+# assumes file pointer is immediately
+#  after the 'fmt ' id
+def _read_fmt_chunk(fid):
+    res = struct.unpack('lhHLLHH',fid.read(20))
+    size, comp, noc, rate, sbytes, ba, bits = res
+    if (comp != 1 or size > 16):
+        print "Warning:  unfamiliar format bytes..."
+        if (size>16):
+            fid.read(size-16)
+    return size, comp, noc, rate, sbytes, ba, bits
+
+# assumes file pointer is immediately 
+#   after the 'data' id
+def _read_data_chunk(fid, noc, bits):
+    size = struct.unpack('l',fid.read(4))[0]
+    if bits == 8:
+        data = numpy.fromfile(fid, dtype=numpy.ubyte, count=size)
+        if noc > 1:
+            data = data.reshape(-1,noc)
+    else:
+        bytes = bits//8
+        dtype = 'i%d' % bytes
+        data = numpy.fromfile(fid, dtype=dtype, count=size//bytes)
+        if noc > 1:
+            data = data.reshape(-1,noc)
+    return data
+
+def _read_riff_chunk(fid):
+    str1 = fid.read(4)
+    fsize = struct.unpack('L', fid.read(4))[0] + 8
+    str2 = fid.read(4)
+    if (str1 != 'RIFF' or str2 != 'WAVE'):
+        raise ValueError, "Not a WAV file."
+    return fsize
+
+# open a wave-file
+def read(file):
+    """Return the sample rate (in samples/sec) and data from a WAV file
+
+    The file can be an open file or a filename.
+    The returned sample rate is a Python integer
+    The data is returned as a numpy array with a 
+        data-type determined from the file. 
+    """
+    if hasattr(file,'read'):
+        fid = file
+    else:
+        fid = open(file, 'rb')
+
+    fsize = _read_riff_chunk(fid)
+    noc = 1
+    bits = 8
+    while (fid.tell() < fsize):
+        # read the next chunk
+        chunk_id = fid.read(4)
+        if chunk_id == 'fmt ':
+            print "Reading fmt chunk"
+            size, comp, noc, rate, sbytes, ba, bits = _read_fmt_chunk(fid)
+        elif chunk_id == 'data':
+            print "Reading data chunk"
+            data = _read_data_chunk(fid, noc, bits) 
+        else:
+            print "Warning:  %s chunk not understood"
+            size = struct.unpack('L',fid.read(4))[0]
+            bytes = fid.read(size)
+    fid.close()
+    return rate, data
+
+# Write a wave-file
+# sample rate, data
+def write(filename, rate, data):
+    """Write a numpy array as a WAV file
+    
+    filename -- The name of the file to write (will be over-written)
+    rate -- The sample rate (in samples/sec).
+    data -- A 1-d or 2-d numpy array of integer data-type. 
+            The bits-per-sample will be determined by the data-type
+            To write multiple-channels, use a 2-d array of shape 
+            (Nsamples, Nchannels)
+
+    Writes a simple uncompressed WAV file. 
+    """
+    fid = open(filename, 'wb')
+    fid.write('RIFF')
+    fid.write('\x00\x00\x00\x00')
+    fid.write('WAVE')
+    # fmt chunk
+    fid.write('fmt ')
+    if data.ndim == 1:
+        noc = 1
+    else:
+        noc = data.shape[1]
+    bits = data.dtype.itemsize * 8
+    sbytes = rate*(bits / 8)*noc
+    ba = noc * (bits / 8)
+    fid.write(struct.pack('lhHLLHH', 16, 1, noc, rate, sbytes, ba, bits))
+    # data chunk
+    fid.write('data')
+    fid.write(struct.pack('l', data.nbytes))
+    data.tofile(fid)
+    # Determine file size and place it in correct
+    #  position at start of the file. 
+    size = fid.tell()
+    fid.seek(4)
+    fid.write(struct.pack('l', size-8))
+    fid.close()