From e0a5ed57f0569a41ee72ed81c1242703cc67706e Mon Sep 17 00:00:00 2001
From: "W. Trevor King" <wking@tremily.us>
Date: Sat, 21 Jul 2012 11:17:08 -0400
Subject: [PATCH] Add igor.packed.walk for traversing a packed experiment
 filesystem.

---
 igor/packed.py | 11 +++++++++++
 igor/util.py   | 16 ++++++++++++++++
 test/test.py   | 25 ++++++++++++++++++++++++-
 3 files changed, 51 insertions(+), 1 deletion(-)

diff --git a/igor/packed.py b/igor/packed.py
index da773bb..8b5537a 100644
--- a/igor/packed.py
+++ b/igor/packed.py
@@ -22,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
@@ -181,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])
diff --git a/igor/util.py b/igor/util.py
index b91bf2b..ecc783a 100644
--- a/igor/util.py
+++ b/igor/util.py
@@ -110,3 +110,19 @@ def checksum(buffer, byte_order, oldcksum, numbytes):
         if oldcksum > 2**31:
             oldcksum -= 2**31
     return oldcksum & 0xffff
+
+def _bytes(obj, encoding='utf-8'):
+    """Convert bytes or strings into bytes
+
+    >>> _bytes(b'123')
+    '123'
+    >>> _bytes('123')
+    '123'
+    """
+    if _sys.version_info >= (3,):
+        if isinstance(obj, bytes):
+            return obj
+        else:
+            return bytes(obj, encoding)
+    else:
+        return bytes(obj)
diff --git a/test/test.py b/test/test.py
index 4a76476..1934493 100644
--- a/test/test.py
+++ b/test/test.py
@@ -1415,6 +1415,21 @@ filesystem:
           'angleQ1': <WaveRecord ...>,
           'radiusData': <WaveRecord ...>,
           'radiusQ1': <WaveRecord ...>}}
+<BLANKLINE>
+walking filesystem:
+walk callback on ([], root, {'K0': 0.0,...})
+walk callback on (['root'], K0, 0.0)
+walk callback on (['root'], K1, 0.0)
+walk callback on (['root'], K10, 0.0)
+...
+walk callback on (['root'], K9, 0.0)
+walk callback on (['root'], Packages, {'PolarGraphs': ...})
+walk callback on (['root', 'Packages'], PolarGraphs, {...})
+walk callback on (['root', 'Packages', 'PolarGraphs'], V_bottom, 232.0)
+...
+walk callback on (['root', 'Packages'], WMDataBase, {...})
+...
+walk callback on (['root'], radiusQ1, <WaveRecord ...>)
 """
 
 import os.path
@@ -1423,6 +1438,7 @@ from pprint import pformat
 from igor import LOG
 from igor.binarywave import load as loadibw
 from igor.packed import load as loadpxp
+from igor.packed import walk as _walk
 from igor.record.base import TextRecord
 from igor.record.folder import FolderStartRecord, FolderEndRecord
 from igor.record.variables import VariablesRecord
@@ -1438,7 +1454,11 @@ def dumpibw(filename):
     data = loadibw(path)
     pprint(data)
 
-def dumppxp(filename):
+def walk_callback(dirpath, key, value):
+    print('walk callback on ({}, {}, {})'.format(
+            dirpath, key, pformat(value)))
+
+def dumppxp(filename, walk=True):
     LOG.info('Testing {}\n'.format(filename))
     path = os.path.join(_data_dir, filename)
     records,filesystem = loadpxp(path)
@@ -1456,6 +1476,9 @@ def dumppxp(filename):
             pprint(record)
     print('\nfilesystem:')
     pprint(filesystem)
+    if walk:
+        print('\nwalking filesystem:')
+        _walk(filesystem, walk_callback)
 
 def pprint(data):
     lines = pformat(data).splitlines()
-- 
2.26.2