README: fix reStructuredText for `pip install` example.

[igor.git] / igor / packed.py

diff --git a/igor/packed.py b/igor/packed.py
index 564725aabbb9df7177c02a31aaa6afea0474b50c..8b5537a8f83f095f6b0b8eb9325b690dbee7da03 100644 (file)
--- a/igor/packed.py
+++ b/igor/packed.py
@@ -1,4 +1,19 @@
-# Copyright
+# Copyright (C) 2012 W. Trevor King <wking@tremily.us>
+#
+# This file is part of igor.
+#
+# igor is free software: you can redistribute it and/or modify it under the
+# terms of the GNU Lesser General Public License as published by the Free
+# Software Foundation, either version 3 of the License, or (at your option) any
+# later version.
+#
+# igor is distributed in the hope that it will be useful, but WITHOUT ANY
+# WARRANTY; without even the implied warranty of MERCHANTABILITY or FITNESS FOR
+# A PARTICULAR PURPOSE.  See the GNU Lesser General Public License for more
+# details.
+#
+# You should have received a copy of the GNU Lesser General Public License
+# along with igor.  If not, see <http://www.gnu.org/licenses/>.
 
 "Read IGOR Packed Experiment files files into records."
 
@@ -7,6 +22,7 @@ from .struct import Structure as _Structure
 from .struct import Field as _Field
 from .util import byte_order as _byte_order
 from .util import need_to_reorder_bytes as _need_to_reorder_bytes
+from .util import _bytes
 from .record import RECORD_TYPE as _RECORD_TYPE
 from .record.base import UnknownRecord as _UnknownRecord
 from .record.base import UnusedRecord as _UnusedRecord
@@ -52,9 +68,13 @@ def load(filename, strict=True, ignore_unknown=True):
         while True:
             PackedFileRecordHeader.byte_order = initial_byte_order
             PackedFileRecordHeader.setup()
-            b = buffer(f.read(PackedFileRecordHeader.size))
+            b = bytes(f.read(PackedFileRecordHeader.size))
             if not b:
                 break
+            if len(b) < PackedFileRecordHeader.size:
+                raise ValueError(
+                    ('not enough data for the next record header ({} < {})'
+                     ).format(len(b), PackedFileRecordHeader.size))
             _LOG.debug('reading a new packed experiment file record')
             header = PackedFileRecordHeader.unpack_from(b)
             if header['version'] and not byte_order:
@@ -69,7 +89,11 @@ def load(filename, strict=True, ignore_unknown=True):
                     header = PackedFileRecordHeader.unpack_from(b)
                     _LOG.debug(
                         'reordered version: {}'.format(header['version']))
-            data = buffer(f.read(header['numDataBytes']))
+            data = bytes(f.read(header['numDataBytes']))
+            if len(data) < header['numDataBytes']:
+                raise ValueError(
+                    ('not enough data for the next record ({} < {})'
+                     ).format(len(b), header['numDataBytes']))
             record_type = _RECORD_TYPE.get(
                 header['recordType'] & PACKEDRECTYPE_MASK, _UnknownRecord)
             _LOG.debug('the new record has type {} ({}).'.format(
@@ -158,3 +182,13 @@ def _check_filename(dir_stack, filename):
     if filename in cwd:
         raise ValueError('collision on name {} in {}'.format(
                 filename, ':'.join(d for d,cwd in dir_stack)))
+
+def walk(filesystem, callback, dirpath=None):
+    """Walk a packed experiment filesystem, operating on each key,value pair.
+    """
+    if dirpath is None:
+        dirpath = []
+    for key,value in sorted((_bytes(k),v) for k,v in filesystem.items()):
+        callback(dirpath, key, value)
+        if isinstance(value, dict):
+            walk(filesystem=value, callback=callback, dirpath=dirpath+[key])