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 Node (list):
  24     def __init__(self, root, type, tags, *args, **kwargs):
  25         list.__init__(self, *args, **kwargs)
  26         self.root = root
  27         self.type = type
  28         self.tags = tags
  29
  30     def pre_extend(self, a):
  31         """
  32         >>> n = Node(None, None, None, ['a', 'b', 'c'])
  33         >>> n.pre_extend([1, 2, 3])
  34         >>> print n
  35         [1, 2, 3, 'a', 'b', 'c']
  36         """
  37         for i,element in enumerate(a):
  38             self.insert(i, element)
  39
  40
  41 class Tree (list):
  42     """
  43     >>> t = Tree('a')
  44     >>> t.append(Tree('1'))
  45     >>> t.append(Tree('2'))
  46     >>> t[0].append(Tree('A'))
  47     >>> print '\\n'.join(['|'.join(x) for x in t.traverse()])
  48     a
  49     a|1
  50     a|1|A
  51     a|2
  52     >>> print '\\n'.join(['|'.join(x) for x in t.traverse(prefix=['z'])])
  53     z|a
  54     z|a|1
  55     z|a|1|A
  56     z|a|2
  57     """
  58     def __init__(self, data=None, root=None, *args, **kwargs):
  59         self.data = data
  60         self.root = root
  61         list.__init__(self, *args, **kwargs)
  62
  63     def traverse(self, prefix=[], depth=0, type='both', dirtag=None):
  64         p = prefix + [self.data]
  65         if depth < len(p) and type in ['both', 'dirs']:
  66             yield Node(self.root, 'dir', [], p[depth:])
  67         for child in self:
  68             if hasattr(child, 'traverse'):
  69                 for grandchild in child.traverse(
  70                     p, depth=depth, type=type, dirtag=dirtag):
  71                     yield grandchild
  72             elif depth <= len(p) and type in ['both', 'files']:
  73                 n = Node(self.root, 'file', [], (p + [child])[depth:])
  74                 if dirtag != None:
  75                     n.tags = dirtag.node_tags(n)
  76                 yield n
  77
  78
  79 def dir_tree(root_dir):
  80     """Generate a directory tree.
  81
  82     >>> t = dir_tree('test/tag')
  83     >>> print '\\n'.join(['|'.join(x) for x in t.traverse(depth=1)])
  84     1
  85     1|a1.dat
  86     1|b1.dat
  87     x
  88     x|y
  89     x|y|b1.dat
  90     x|a2.dat
  91     x|b3.dat
  92     >>> print '\\n'.join(['|'.join(x) for x in t.traverse(depth=1, type='files')])
  93     1|a1.dat
  94     1|b1.dat
  95     x|y|b1.dat
  96     x|a2.dat
  97     x|b3.dat
  98     >>> print '\\n'.join(['|'.join(x) for x in t.traverse(depth=1, type='dirs')])
  99     1
 100     x
 101     x|y
 102     """
 103     if not os.path.exists(root_dir):
 104         raise ValueError(root_dir)
 105     trees = {}
 106     for dirpath,dirnames,filenames in os.walk(root_dir, followlinks=True):
 107         t = trees.get(
 108             dirpath,
 109             Tree(os.path.basename(dirpath), root=root_dir))
 110         if dirpath == root_dir:
 111             root = t
 112         for d in sorted(dirnames):
 113             t2 = Tree(d, root=root_dir)
 114             trees[os.path.join(dirpath, d)] = t2
 115             t.append(t2)
 116         t.extend(sorted(filenames))
 117     return root
 118
 119
 120 class Dirtag (object):
 121     """
 122     >>> d = Dirtag('test/raw', 'test/tag')
 123     >>> print '\\n'.join(['\t'.join([x.type, x.root, '|'.join(x),
 124     ...                              ','.join(['/'.join(t) for t in x.tags])])
 125     ...                   for x in d.elements()]
 126     ...                 )  # doctest: +NORMALIZE_WHITESPACE
 127     dir       test/raw      a
 128     file      test/raw      a|a1.dat      1
 129     file      test/raw      a|a2.dat      x
 130     file      test/raw      a|a3.dat
 131     dir       test/raw      b
 132     file      test/raw      b|b1.dat      1,x/y
 133     file      test/raw      b|b2.dat
 134     file      test/raw      b|b3.dat      x
 135     >>> print '\\n'.join(['|'.join(x) for x in d.elements(['x'])])
 136     x|y
 137     x|y|b1.dat
 138     x|a2.dat
 139     x|b3.dat
 140     >>> print '\\n'.join(['|'.join(x) for x in d.tags(['b', 'b2.dat'])])
 141     <BLANKLINE>
 142     >>> print '\\n'.join(['|'.join(x) for x in d.tags(['b', 'b1.dat'])])
 143     1
 144     x|y
 145
 146     >>> print d.tag_path(['a', 'a3.dat'], ['x', 'y'])
 147     test/tag/x/y/a3.dat
 148     >>> os.listdir('test/tag/x/y')
 149     ['b1.dat']
 150     >>> d.add_tag(['a', 'a3.dat'], ['x', 'y'])
 151     >>> os.listdir('test/tag/x/y')
 152     ['a3.dat', 'b1.dat']
 153     >>> print 'Z'+os.path.realpath('test/tag/x/y/a3.dat') # doctest: +ELLIPSIS
 154     Z.../test/raw/a/a3.dat
 155     >>> d.remove_tag(['a', 'a3.dat'], ['x', 'y'])
 156     >>> os.listdir('test/tag/x/y')
 157     ['b1.dat']
 158     """
 159     def __init__(self, raw_dir, tag_dir):
 160         self.raw_dir = raw_dir
 161         self.tag_dir = tag_dir
 162
 163     def elements(self, tag=None):
 164         if tag == None:
 165             tag = []
 166             nodes = dir_tree(self.raw_dir).traverse(depth=1, dirtag=self)
 167         else:
 168             p = [self.tag_dir]+tag
 169             nodes = dir_tree(os.path.join(*p)).traverse(
 170                 prefix=p[:-1], depth=len(p), dirtag=self)
 171         for node in nodes:
 172             node.pre_extend(tag)
 173             yield node
 174
 175     def tags(self, target):
 176         for node in dir_tree(self.tag_dir).traverse(depth=1):
 177             p = os.path.join(*([self.tag_dir]+node))
 178             t = os.path.abspath(os.path.join(*([self.raw_dir] + target)))
 179             if (os.path.islink(p) and os.path.realpath(p) == t):
 180                 yield os.path.join(node[:-1])
 181
 182     def tag_path(self, target, tag):
 183         # TODO: assumes unique basenames.  Check?
 184         return os.path.join(*([self.tag_dir] + tag + target[-1:]))
 185
 186     def add_tag(self, target, tag):
 187         tag_path = self.tag_path(target, tag)
 188         target_path = os.path.abspath(os.path.join(*([self.raw_dir]+target)))
 189         cwd = os.getcwd()
 190         try:
 191             os.chdir(os.path.dirname(tag_path))
 192             os.symlink(os.path.relpath(target_path),
 193                        os.path.basename(tag_path))
 194         finally:
 195             os.chdir(cwd)
 196
 197     def remove_tag(self, target, tag):
 198         tag_path = self.tag_path(target, tag)
 199         cwd = os.getcwd()
 200         try:
 201             os.chdir(os.path.dirname(tag_path))
 202             os.remove(os.path.basename(tag_path))
 203         finally:
 204             os.chdir(cwd)
 205
 206     def raw_node(self, node):
 207         """Return a node for the real file linked to by node.
 208         """
 209         if node.root == self.raw_dir:
 210             return node
 211         path = os.path.realpath(os.path.join(node.root, *node))
 212         raw_path = os.path.relpath(path, self.raw_dir)
 213         return raw_path.split(os.path.sep)
 214
 215     def node_tags(self, node):
 216         """Return a list of tags for a particular node.
 217         """
 218         return self.tags(self.raw_node(node))
 219
 220
 221 if __name__ == '__main__':
 222     import doctest
 223     doctest.testmod()