2 # Copyright (C) 2010 W. Trevor King <wking@drexel.edu>
4 # This file is part of Dirtag.
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.
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.
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/>.
24 def __init__(self, root, type, tags, *args, **kwargs):
25 list.__init__(self, *args, **kwargs)
30 def pre_extend(self, a):
32 >>> n = Node(None, None, None, ['a', 'b', 'c'])
33 >>> n.pre_extend([1, 2, 3])
35 [1, 2, 3, 'a', 'b', 'c']
37 for i,element in enumerate(a):
38 self.insert(i, element)
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()])
52 >>> print '\\n'.join(['|'.join(x) for x in t.traverse(prefix=['z'])])
58 def __init__(self, data=None, root=None, *args, **kwargs):
61 list.__init__(self, *args, **kwargs)
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:])
68 if hasattr(child, 'traverse'):
69 for grandchild in child.traverse(
70 p, depth=depth, type=type, dirtag=dirtag):
72 elif depth <= len(p) and type in ['both', 'files']:
73 n = Node(self.root, 'file', [], (p + [child])[depth:])
75 n.tags = dirtag.node_tags(n)
79 def dir_tree(root_dir):
80 """Generate a directory tree.
82 >>> t = dir_tree('test/tag')
83 >>> print '\\n'.join(['|'.join(x) for x in t.traverse(depth=1)])
92 >>> print '\\n'.join(['|'.join(x) for x in t.traverse(depth=1, type='files')])
98 >>> print '\\n'.join(['|'.join(x) for x in t.traverse(depth=1, type='dirs')])
103 if not os.path.exists(root_dir):
104 raise ValueError(root_dir)
106 for dirpath,dirnames,filenames in os.walk(root_dir, followlinks=True):
109 Tree(os.path.basename(dirpath), root=root_dir))
110 if dirpath == root_dir:
112 for d in sorted(dirnames):
113 t2 = Tree(d, root=root_dir)
114 trees[os.path.join(dirpath, d)] = t2
116 t.extend(sorted(filenames))
120 class Dirtag (object):
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
128 file test/raw a|a1.dat 1
129 file test/raw a|a2.dat x
130 file test/raw a|a3.dat
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'])])
140 >>> print '\\n'.join(['|'.join(x) for x in d.tags(['b', 'b2.dat'])])
142 >>> print '\\n'.join(['|'.join(x) for x in d.tags(['b', 'b1.dat'])])
146 >>> print d.tag_path(['a', 'a3.dat'], ['x', 'y'])
148 >>> os.listdir('test/tag/x/y')
150 >>> d.add_tag(['a', 'a3.dat'], ['x', 'y'])
151 >>> os.listdir('test/tag/x/y')
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')
159 def __init__(self, raw_dir, tag_dir):
160 self.raw_dir = raw_dir
161 self.tag_dir = tag_dir
163 def elements(self, tag=None):
166 nodes = dir_tree(self.raw_dir).traverse(depth=1, dirtag=self)
168 p = [self.tag_dir]+tag
169 nodes = dir_tree(os.path.join(*p)).traverse(
170 prefix=p[:-1], depth=len(p), dirtag=self)
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])
182 def tag_path(self, target, tag):
183 # TODO: assumes unique basenames. Check?
184 return os.path.join(*([self.tag_dir] + tag + target[-1:]))
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)))
191 os.chdir(os.path.dirname(tag_path))
192 os.symlink(os.path.relpath(target_path),
193 os.path.basename(tag_path))
197 def remove_tag(self, target, tag):
198 tag_path = self.tag_path(target, tag)
201 os.chdir(os.path.dirname(tag_path))
202 os.remove(os.path.basename(tag_path))
206 def raw_node(self, node):
207 """Return a node for the real file linked to by node.
209 if node.root == self.raw_dir:
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)
215 def node_tags(self, node):
216 """Return a list of tags for a particular node.
218 return self.tags(self.raw_node(node))
221 if __name__ == '__main__':