More levels

DataProcessingTools/levels.py

@@ -1,5 +1,8 @@
 import os
-levels = ['day', 'session', 'array', 'channel']
+import glob
+
+levels = ['subjects', 'subject', 'day', 'session', 'array', 'channel','cell']
+level_patterns_s = ["*", "*", "[0-9]*", "session[0-9]*", "array[0-9]*", "channel[0-9]*", "cell[0-9]*"]
 
 
 def level(cwd=None):
@@ -33,3 +36,24 @@ def resolve_level(target_level, cwd=None):
         pl.append("..")
     print(pl)
     return os.path.join(*pl)
+
+
+def get_level_dirs(target_level, cwd=None):
+    """
+    Get the directories representing `target_level` under the directory
+    pointed to by `cwd`.
+    """
+    if cwd is None:
+        cwd = os.getcwd()
+    this_level = level(cwd)
+    this_idx = levels.index(this_level)
+    target_idx = levels.index(target_level)
+    if target_idx <= this_idx:
+        rel_path = resolve_level(target_level, cwd)
+        pattern = level_patterns_s[target_idx]
+        gpattern = os.path.join(cwd, rel_path, pattern)
+        dirs = glob.glob(gpattern)
+    else:
+        patterns = level_patterns_s[this_idx+1:target_idx+1]
+        dirs = glob.glob(os.path.join(cwd, *patterns))
+    return dirs

README.md

@@ -18,7 +18,10 @@ cwd = "Pancake/20130923/session01/array01"
 ll = DPT.levels.level(cwd)
 
 # Resolve the relative path from `cwd` to the `session` directory.
-rr = DPHT.resolve_level("session", cwd)
+rr = DPT.levels.resolve_level("session", cwd)
+
+# Find all e.g. cell directories under the current directory
+celldirs = DPT.levels.get_level_dirs("cell", cwd)
 ```
 
 ### Objects

setup.py

@@ -1,7 +1,7 @@
 from setuptools import setup
 
 setup(name="DataProcessingTools",
-      version="0.3.0",
+      version="0.4.0",
       description="""Tools for processing data with hierarchical organization""",
       url="https://github.com/grero/DataProcessingTools.git",
       author="Roger Herikstad",

tests/test_levels.py

@@ -1,4 +1,6 @@
 import DataProcessingTools as DPT
+import tempfile
+import os
 
 
 def test_level():
@@ -8,3 +10,23 @@ def test_level():
 
     rr = DPT.levels.resolve_level("session", cwd)
     assert rr == "./../../.."
+
+    tdir = tempfile.gettempdir()
+    cwd = os.getcwd()
+    os.chdir(tdir)
+    dir1 = "Pancake/20130923/session01/array02/channel33"
+    dir2 = "Pancake/20130923/session01/array02/channel34"
+    for d in [dir1,dir2]:
+        if not os.path.isdir(d):
+            os.makedirs(d)
+    dirs = DPT.levels.get_level_dirs("channel", "Pancake/20130923/session01/array02")
+    assert dirs[0] == "Pancake/20130923/session01/array02/channel33"
+    assert dirs[1] == "Pancake/20130923/session01/array02/channel34"
+
+    dirs = DPT.levels.get_level_dirs("session", "Pancake/20130923/session01/array02")
+    assert dirs[0] == "Pancake/20130923/session01/array02/./../../session01"
+
+    dirs = DPT.levels.get_level_dirs("array", "Pancake/20130923/session01/array02")
+    assert dirs[0] == "Pancake/20130923/session01/array02/./../array02"
+
+    os.chdir(cwd)