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     trees = {}
  75     for dirpath,dirnames,filenames in os.walk(root_dir, followlinks=True):
  76         t = trees.get(
  77             dirpath,
  78             Tree(os.path.basename(dirpath)))
  79         if dirpath == root_dir:
  80             root = t
  81         for d in sorted(dirnames):
  82             t2 = Tree(d)
  83             trees[os.path.join(dirpath, d)] = t2
  84             t.append(t2)
  85         t.extend(sorted(filenames))
  86     return root
  87
  88
  89 class Dirtag (object):
  90     """
  91     >>> d = Dirtag('test/raw', 'test/tag')
  92     >>> print '\\n'.join(['|'.join(x) for x in d.elements()])
  93     a
  94     a|a1.dat
  95     a|a2.dat
  96     a|a3.dat
  97     b
  98     b|b1.dat
  99     b|b2.dat
 100     b|b3.dat
 101     >>> print '\\n'.join(['|'.join(x) for x in d.elements(['x'])])
 102     x|y
 103     x|y|b1.dat
 104     x|a2.dat
 105     x|b3.dat
 106     >>> print '\\n'.join(['|'.join(x) for x in d.tags(['b', 'b2.dat'])])
 107     <BLANKLINE>
 108     >>> print '\\n'.join(['|'.join(x) for x in d.tags(['b', 'b1.dat'])])
 109     1
 110     x|y
 111
 112     >>> print d.tag_path(['a', 'a3.dat'], ['x', 'y'])
 113     test/tag/x/y/a3.dat
 114     >>> os.listdir('test/tag/x/y')
 115     ['b1.dat']
 116     >>> d.add_tag(['a', 'a3.dat'], ['x', 'y'])
 117     >>> os.listdir('test/tag/x/y')
 118     ['a3.dat', 'b1.dat']
 119     >>> print 'Z'+os.path.realpath('test/tag/x/y/a3.dat') # doctest: +ELLIPSIS
 120     Z.../test/raw/a/a3.dat
 121     >>> d.remove_tag(['a', 'a3.dat'], ['x', 'y'])
 122     >>> os.listdir('test/tag/x/y')
 123     ['b1.dat']
 124     """
 125     def __init__(self, raw_dir, tag_dir):
 126         self.raw_dir = raw_dir
 127         self.tag_dir = tag_dir
 128
 129     def elements(self, tag=None):
 130         if tag == None:
 131             tag = []
 132             nodes = dir_tree(self.raw_dir).traverse(depth=1)
 133         else:
 134             p = [self.tag_dir]+tag
 135             nodes = dir_tree(os.path.join(*p)).traverse(
 136                 prefix=p[:-1], depth=len(p))
 137         for node in nodes:
 138             yield tag+node
 139
 140     def tags(self, target):
 141         for node in dir_tree(self.tag_dir).traverse(depth=1):
 142             p = os.path.join(*([self.tag_dir]+node))
 143             t = os.path.abspath(os.path.join(*([self.raw_dir] + target)))
 144             if (os.path.islink(p) and os.path.realpath(p) == t):
 145                 yield os.path.join(node[:-1])
 146
 147     def tag_path(self, target, tag):
 148         # TODO: assumes unique basenames.  Check?
 149         return os.path.join(*([self.tag_dir] + tag + target[-1:]))
 150
 151     def add_tag(self, target, tag):
 152         tag_path = self.tag_path(target, tag)
 153         target_path = os.path.abspath(os.path.join(*([self.raw_dir]+target)))
 154         cwd = os.getcwd()
 155         try:
 156             os.chdir(os.path.dirname(tag_path))
 157             os.symlink(os.path.relpath(target_path),
 158                        os.path.basename(tag_path))
 159         finally:
 160             os.chdir(cwd)
 161
 162     def remove_tag(self, target, tag):
 163         tag_path = self.tag_path(target, tag)
 164         cwd = os.getcwd()
 165         try:
 166             os.chdir(os.path.dirname(tag_path))
 167             os.remove(os.path.basename(tag_path))
 168         finally:
 169             os.chdir(cwd)
 170
 171 if __name__ == '__main__':
 172     import doctest
 173     doctest.testmod()