dirtag/__init__.py

   1 #!/usr/bin/python
   2 # Copyright (C) 2010 W. Trevor King <wking@drexel.edu>
   3 #
   4 # This file is part of Dirtag.
   5 #
   6 # Dirtag is free software: you can redistribute it and/or modify it
   7 # under the terms of the GNU General Public License as published by
   8 # the Free Software Foundation, either version 3 of the License, or
   9 # (at your option) any later version.
  10 #
  11 # Dirtag is distributed in the hope that it will be useful, but
  12 # WITHOUT ANY WARRANTY; without even the implied warranty of
  13 # MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
  14 # General Public License for more details.
  15 #
  16 # You should have received a copy of the GNU General Public License
  17 # along with Dirtag.  If not, see <http://www.gnu.org/licenses/>.
  18
  19 import os
  20 import os.path
  21
  22
  23 class Tree (list):
  24     """
  25     >>> t = Tree('a')
  26     >>> t.append(Tree('1'))
  27     >>> t.append(Tree('2'))
  28     >>> t[0].append(Tree('A'))
  29     >>> print '\\n'.join(['|'.join(x) for x in t.traverse()])
  30     a
  31     a|1
  32     a|1|A
  33     a|2
  34     >>> print '\\n'.join(['|'.join(x) for x in t.traverse(prefix=['z'])])
  35     z|a
  36     z|a|1
  37     z|a|1|A
  38     z|a|2
  39     >>> print '\\n'.join(['|'.join(x) for x in t.traverse(depth=1)])
  40     1
  41     1|A
  42     2
  43     """
  44     def __init__(self, data=None, *args, **kwargs):
  45         self.data = data
  46         list.__init__(self, *args, **kwargs)
  47
  48     def traverse(self, prefix=[], depth=0):
  49         p = prefix + [self.data]
  50         if depth < len(p):
  51             yield p[depth:]
  52         for child in self:
  53             if hasattr(child, 'traverse'):
  54                 for grandchild in child.traverse(p, depth=depth):
  55                     yield grandchild
  56             elif depth <= len(p):
  57                 yield (p + [child])[depth:]
  58
  59
  60 def dir_tree(root_dir):
  61     """Generate a directory tree.
  62
  63     >>> t = dir_tree('test/tag')
  64     >>> print '\\n'.join(['|'.join(x) for x in t.traverse(depth=1)])
  65     1
  66     1|a1.dat
  67     1|b1.dat
  68     x
  69     x|y
  70     x|y|b1.dat
  71     x|a2.dat
  72     x|b3.dat
  73     """
  74     if not os.path.exists(root_dir):
  75         raise ValueError(root_dir)
  76     trees = {}
  77     for dirpath,dirnames,filenames in os.walk(root_dir, followlinks=True):
  78         t = trees.get(
  79             dirpath,
  80             Tree(os.path.basename(dirpath)))
  81         if dirpath == root_dir:
  82             root = t
  83         for d in sorted(dirnames):
  84             t2 = Tree(d)
  85             trees[os.path.join(dirpath, d)] = t2
  86             t.append(t2)
  87         t.extend(sorted(filenames))
  88     return root
  89
  90
  91 class Dirtag (object):
  92     """
  93     >>> d = Dirtag('test/raw', 'test/tag')
  94     >>> print '\\n'.join(['|'.join(x) for x in d.elements()])
  95     a
  96     a|a1.dat
  97     a|a2.dat
  98     a|a3.dat
  99     b
 100     b|b1.dat
 101     b|b2.dat
 102     b|b3.dat
 103     >>> print '\\n'.join(['|'.join(x) for x in d.elements(['x'])])
 104     x|y
 105     x|y|b1.dat
 106     x|a2.dat
 107     x|b3.dat
 108     >>> print '\\n'.join(['|'.join(x) for x in d.tags(['b', 'b2.dat'])])
 109     <BLANKLINE>
 110     >>> print '\\n'.join(['|'.join(x) for x in d.tags(['b', 'b1.dat'])])
 111     1
 112     x|y
 113
 114     >>> print d.tag_path(['a', 'a3.dat'], ['x', 'y'])
 115     test/tag/x/y/a3.dat
 116     >>> os.listdir('test/tag/x/y')
 117     ['b1.dat']
 118     >>> d.add_tag(['a', 'a3.dat'], ['x', 'y'])
 119     >>> os.listdir('test/tag/x/y')
 120     ['a3.dat', 'b1.dat']
 121     >>> print 'Z'+os.path.realpath('test/tag/x/y/a3.dat') # doctest: +ELLIPSIS
 122     Z.../test/raw/a/a3.dat
 123     >>> d.remove_tag(['a', 'a3.dat'], ['x', 'y'])
 124     >>> os.listdir('test/tag/x/y')
 125     ['b1.dat']
 126     """
 127     def __init__(self, raw_dir, tag_dir):
 128         self.raw_dir = raw_dir
 129         self.tag_dir = tag_dir
 130
 131     def elements(self, tag=None):
 132         if tag == None:
 133             tag = []
 134             nodes = dir_tree(self.raw_dir).traverse(depth=1)
 135         else:
 136             p = [self.tag_dir]+tag
 137             nodes = dir_tree(os.path.join(*p)).traverse(
 138                 prefix=p[:-1], depth=len(p))
 139         for node in nodes:
 140             yield tag+node
 141
 142     def tags(self, target):
 143         for node in dir_tree(self.tag_dir).traverse(depth=1):
 144             p = os.path.join(*([self.tag_dir]+node))
 145             t = os.path.abspath(os.path.join(*([self.raw_dir] + target)))
 146             if (os.path.islink(p) and os.path.realpath(p) == t):
 147                 yield os.path.join(node[:-1])
 148
 149     def tag_path(self, target, tag):
 150         # TODO: assumes unique basenames.  Check?
 151         return os.path.join(*([self.tag_dir] + tag + target[-1:]))
 152
 153     def add_tag(self, target, tag):
 154         tag_path = self.tag_path(target, tag)
 155         target_path = os.path.abspath(os.path.join(*([self.raw_dir]+target)))
 156         cwd = os.getcwd()
 157         try:
 158             os.chdir(os.path.dirname(tag_path))
 159             os.symlink(os.path.relpath(target_path),
 160                        os.path.basename(tag_path))
 161         finally:
 162             os.chdir(cwd)
 163
 164     def remove_tag(self, target, tag):
 165         tag_path = self.tag_path(target, tag)
 166         cwd = os.getcwd()
 167         try:
 168             os.chdir(os.path.dirname(tag_path))
 169             os.remove(os.path.basename(tag_path))
 170         finally:
 171             os.chdir(cwd)
 172
 173 if __name__ == '__main__':
 174     import doctest
 175     doctest.testmod()