41132a106c48bbc80774bb2d850cf8c742922108
[hooke.git] / hooke / playlist.py
1 # Copyright (C) 2010 W. Trevor King <wking@drexel.edu>
2 #
3 # This file is part of Hooke.
4 #
5 # Hooke is free software: you can redistribute it and/or modify it
6 # under the terms of the GNU Lesser General Public License as
7 # published by the Free Software Foundation, either version 3 of the
8 # License, or (at your option) any later version.
9 #
10 # Hooke is distributed in the hope that it will be useful, but WITHOUT
11 # ANY WARRANTY; without even the implied warranty of MERCHANTABILITY
12 # or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU Lesser General
13 # Public License for more details.
14 #
15 # You should have received a copy of the GNU Lesser General Public
16 # License along with Hooke.  If not, see
17 # <http://www.gnu.org/licenses/>.
18
19 """The `playlist` module provides a :class:`Playlist` and its subclass
20 :class:`FilePlaylist` for manipulating lists of
21 :class:`hooke.curve.Curve`\s.
22 """
23
24 import copy
25 import hashlib
26 import os
27 import os.path
28 import types
29
30 import yaml
31 from yaml.representer import RepresenterError
32
33 from .command_stack import CommandStack
34 from .curve import Curve
35 from .util.itertools import reverse_enumerate
36
37
38 class NoteIndexList (list):
39     """A list that keeps track of a "current" item and additional notes.
40
41     :attr:`index` (i.e. "bookmark") is the index of the currently
42     current curve.  Also keep a :class:`dict` of additional information
43     (:attr:`info`).
44     """
45     def __init__(self, name=None):
46         super(NoteIndexList, self).__init__()
47         self._set_default_attrs()
48         self.__setstate__({'name': name})
49
50     def __str__(self):
51         return str(self.__unicode__())
52
53     def __unicode__(self):
54         return u'<%s %s>' % (self.__class__.__name__, self.name)
55
56     def __repr__(self):
57         return self.__str__()
58
59     def _set_default_attrs(self):
60         self._default_attrs = {
61             'info': {},
62             'name': None,
63             '_index': 0,
64             }
65
66     def __getstate__(self):
67         return self.__dict__.copy()
68
69     def __setstate__(self, state):
70         self._set_default_attrs()
71         if state == True:
72             return
73         self.__dict__.update(self._default_attrs)
74         try:
75             self.__dict__.update(state)
76         except TypeError, e:
77             print state, type(state), e
78         if self.info in [None, {}]:
79             self.info = {}
80
81     def _setup_item(self, item):
82         """Perform any required initialization before returning an item.
83         """
84         pass
85
86     def index(self, value=None, *args, **kwargs):
87         """Extend `list.index`, returning the current index if `value`
88         is `None`.
89         """
90         if value == None:
91             return self._index
92         return super(NoteIndexList, self).index(value, *args, **kwargs)
93
94     def current(self, load=True):
95         if len(self) == 0:
96             return None
97         item = self[self._index]
98         if load == True:
99             self._setup_item(item)
100         return item
101
102     def jump(self, index):
103         if len(self) == 0:
104             self._index = 0
105         else:
106             self._index = index % len(self)
107
108     def next(self):
109         self.jump(self._index + 1)
110
111     def previous(self):
112         self.jump(self._index - 1)
113
114     def items(self, reverse=False):
115         """Iterate through `self` calling `_setup_item` on each item
116         before yielding.
117
118         Notes
119         -----
120         Updates :attr:`_index` during the iteration so
121         :func:`~hooke.plugin.curve.current_curve_callback` works as
122         expected in :class:`~hooke.command.Command`\s called from
123         :class:`~hooke.plugin.playlist.ApplyCommand`.  After the
124         iteration completes, :attr:`_index` is restored to its
125         original value.
126         """
127         index = self._index
128         items = self
129         if reverse == True:
130             items = reverse_enumerate(self)
131         else:
132             items = enumerate(self)
133         for i,item in items:
134             self._index = i
135             self._setup_item(item)
136             yield item
137         self._index = index
138
139     def filter(self, keeper_fn=lambda item:True, load_curves=True,
140                *args, **kwargs):
141         c = copy.deepcopy(self)
142         if load_curves == True:
143             items = c.items(reverse=True)
144         else:
145             items = reversed(c)
146         for item in items: 
147             if keeper_fn(item, *args, **kwargs) != True:
148                 c.remove(item)
149         try: # attempt to maintain the same current item
150             c._index = c.index(self.current())
151         except ValueError:
152             c._index = 0
153         return c
154
155
156 class Playlist (NoteIndexList):
157     """A :class:`NoteIndexList` of :class:`hooke.Curve`\s.
158
159     Keeps a list of :attr:`drivers` for loading curves.
160     """
161     def __init__(self, drivers, name=None):
162         super(Playlist, self).__init__(name=name)
163         self.drivers = drivers
164
165     def _set_default_attrs(self):
166         super(Playlist, self)._set_default_attrs()
167         self._default_attrs['drivers'] = []
168         # List of loaded curves, see :meth:`._setup_item`.
169         self._default_attrs['_loaded'] = []
170         self._default_attrs['_max_loaded'] = 100  # curves to hold in memory simultaneously.
171
172     def __setstate__(self, state):
173         super(Playlist, self).__setstate__(state)
174         if self.drivers in [None, {}]:
175             self.drivers = []
176         if self._loaded in [None, {}]:
177             self._loaded = []
178
179     def append_curve(self, curve):
180         self.append(curve)
181
182     def append_curve_by_path(self, path, info=None, identify=True, hooke=None):
183         path = os.path.normpath(path)
184         c = Curve(path, info=info)
185         c.set_hooke(hooke)
186         if identify == True:
187             c.identify(self.drivers)
188         self.append(c)
189         return c
190
191     def _setup_item(self, curve):
192         if curve != None and curve not in self._loaded:
193             if curve not in self:
194                 self.append(curve)
195             if curve.driver == None:
196                 c.identify(self.drivers)
197             if curve.data == None or max([d.size for d in curve.data]) == 0:
198                 curve.load()
199             self._loaded.append(curve)
200             if len(self._loaded) > self._max_loaded:
201                 oldest = self._loaded.pop(0)
202                 oldest.unload()
203
204     def unload(self, curve):
205         "Inverse of .`_setup_item`."
206         curve.unload()
207         try:
208             self._loaded.remove(curve)
209         except ValueError:
210             pass
211
212
213 def playlist_path(path):
214     """Normalize playlist path extensions.
215
216     Examples
217     --------
218     >>> print playlist_path('playlist')
219     playlist.hkp
220     >>> print playlist_path('playlist.hkp')
221     playlist.hkp
222     >>> print playlist_path(None)
223     None
224     """
225     if path == None:
226         return None
227     if not path.endswith('.hkp'):
228         path += '.hkp'
229     return path
230
231
232 class FilePlaylist (Playlist):
233     """A file-backed :class:`Playlist`.
234
235     Examples
236     --------
237
238     >>> p = FilePlaylist(drivers=['Driver A', 'Driver B'])
239     >>> p.append(Curve('dummy/path/A'))
240     >>> p.append(Curve('dummy/path/B'))
241
242     The data-type is pickleable, to ensure we can move it between
243     processes with :class:`multiprocessing.Queue`\s.
244
245     >>> import pickle
246     >>> s = pickle.dumps(p)
247     >>> z = pickle.loads(s)
248     >>> for curve in z:
249     ...     print curve
250     <Curve A>
251     <Curve B>
252     >>> print z.drivers
253     ['Driver A', 'Driver B']
254
255     The data-type is also YAMLable (see :mod:`hooke.util.yaml`).
256
257     >>> s = yaml.dump(p)
258     >>> z = yaml.load(s)
259     >>> for curve in z:
260     ...     print curve
261     <Curve A>
262     <Curve B>
263     >>> print z.drivers
264     ['Driver A', 'Driver B']
265     """
266     version = '0.2'
267
268     def __init__(self, drivers, name=None, path=None):
269         super(FilePlaylist, self).__init__(drivers, name)
270         self.path = self._base_path = None
271         self.set_path(path)
272         self.relative_curve_paths = True
273         self._relative_curve_paths = False
274
275     def _set_default_attrs(self):
276         super(FilePlaylist, self)._set_default_attrs()
277         self._default_attrs['relative_curve_paths'] = True
278         self._default_attrs['_relative_curve_paths'] = False
279         self._default_attrs['_digest'] = None
280
281     def __getstate__(self):
282         state = super(FilePlaylist, self).__getstate__()
283         assert 'version' not in state, state
284         state['version'] = self.version
285         return state
286
287     def __setstate__(self, state):
288         if 'version' in state:
289             version = state.pop('version')
290             assert version == FilePlaylist.version, (
291                 'invalid version %s (%s) != %s (%s)'
292                 % (version, type(version),
293                    FilePlaylist.version, type(FilePlaylist.version)))
294         super(FilePlaylist, self).__setstate__(state)
295
296     def set_path(self, path):
297         orig_base_path = getattr(self, '_base_path', None)
298         if path == None:
299             if self._base_path == None:
300                 self._base_path = os.getcwd()
301         else:
302             path = playlist_path(path)
303             self.path = path
304             self._base_path = os.path.dirname(os.path.abspath(
305                 os.path.expanduser(self.path)))
306             if self.name == None:
307                 self.name = os.path.basename(path)
308         if self._base_path != orig_base_path:
309             self.update_curve_paths()
310
311     def update_curve_paths(self):
312         for curve in self:
313             curve.set_path(self._curve_path(curve.path))
314
315     def _curve_path(self, path):
316         if self._base_path == None:
317             self._base_path = os.getcwd()
318         path = os.path.join(self._base_path, path)
319         if self._relative_curve_paths == True:
320             path = os.path.relpath(path, self._base_path)
321         return path
322
323     def append_curve(self, curve):
324         curve.set_path(self._curve_path(curve.path))
325         super(FilePlaylist, self).append_curve(curve)
326
327     def append_curve_by_path(self, path, *args, **kwargs):
328         path = self._curve_path(path)
329         super(FilePlaylist, self).append_curve_by_path(path, *args, **kwargs)
330
331     def is_saved(self):
332         return self.digest() == self._digest
333
334     def digest(self):
335         r"""Compute the sha1 digest of the flattened playlist
336         representation.
337
338         Examples
339         --------
340
341         >>> root_path = os.path.sep + 'path'
342         >>> p = FilePlaylist(drivers=[],
343         ...                  path=os.path.join(root_path, 'to','playlist'))
344         >>> p.info['note'] = 'An example playlist'
345         >>> c = Curve(os.path.join(root_path, 'to', 'curve', 'one'))
346         >>> c.info['note'] = 'The first curve'
347         >>> p.append_curve(c)
348         >>> c = Curve(os.path.join(root_path, 'to', 'curve', 'two'))
349         >>> c.info['note'] = 'The second curve'
350         >>> p.append_curve(c)
351         >>> p.digest()
352         'f\xe26i\xb98i\x1f\xb61J7:\xf2\x8e\x1d\xde\xc3}g'
353         """
354         string = self.flatten()
355         return hashlib.sha1(string).digest()
356
357     def flatten(self):
358         """Create a string representation of the playlist.
359
360         A playlist is a YAML document with the following minimal syntax::
361
362             !!python/object/new:hooke.playlist.FilePlaylist
363             state:
364               version: '0.2'
365             listitems:
366             - !!python/object:hooke.curve.Curve
367               path: /path/to/curve/one
368             - !!python/object:hooke.curve.Curve
369               path: /path/to/curve/two
370
371         Relative paths are interpreted relative to the location of the
372         playlist file.
373
374         Examples
375         --------
376
377         >>> from .engine import CommandMessage
378
379         >>> root_path = os.path.sep + 'path'
380         >>> p = FilePlaylist(drivers=[],
381         ...                  path=os.path.join(root_path, 'to','playlist'))
382         >>> p.info['note'] = 'An example playlist'
383         >>> c = Curve(os.path.join(root_path, 'to', 'curve', 'one'))
384         >>> c.info['note'] = 'The first curve'
385         >>> p.append_curve(c)
386         >>> c = Curve(os.path.join(root_path, 'to', 'curve', 'two'))
387         >>> c.info['attr with spaces'] = 'The second curve\\nwith endlines'
388         >>> c.command_stack.extend([
389         ...         CommandMessage('command A', {'arg 0':0, 'arg 1':'X'}),
390         ...         CommandMessage('command B', {'arg 0':1, 'curve':c}),
391         ...         ])
392         >>> p.append_curve(c)
393         >>> print p.flatten()  # doctest: +REPORT_UDIFF
394         # Hooke playlist version 0.2
395         !!python/object/new:hooke.playlist.FilePlaylist
396         listitems:
397         - !!python/object:hooke.curve.Curve
398           info: {note: The first curve}
399           name: one
400           path: curve/one
401         - &id001 !!python/object:hooke.curve.Curve
402           command_stack: !!python/object/new:hooke.command_stack.CommandStack
403             listitems:
404             - !!python/object:hooke.engine.CommandMessage
405               arguments: {arg 0: 0, arg 1: X}
406               command: command A
407               explicit_user_call: true
408             - !!python/object:hooke.engine.CommandMessage
409               arguments:
410                 arg 0: 1
411                 curve: *id001
412               command: command B
413               explicit_user_call: true
414           info: {attr with spaces: 'The second curve
415         <BLANKLINE>
416               with endlines'}
417           name: two
418           path: curve/two
419         state:
420           _base_path: /path/to
421           info: {note: An example playlist}
422           name: playlist.hkp
423           path: /path/to/playlist.hkp
424           version: '0.2'
425         <BLANKLINE>
426         >>> p.relative_curve_paths = False
427         >>> print p.flatten()  # doctest: +REPORT_UDIFF
428         # Hooke playlist version 0.2
429         !!python/object/new:hooke.playlist.FilePlaylist
430         listitems:
431         - !!python/object:hooke.curve.Curve
432           info: {note: The first curve}
433           name: one
434           path: /path/to/curve/one
435         - &id001 !!python/object:hooke.curve.Curve
436           command_stack: !!python/object/new:hooke.command_stack.CommandStack
437             listitems:
438             - !!python/object:hooke.engine.CommandMessage
439               arguments: {arg 0: 0, arg 1: X}
440               command: command A
441               explicit_user_call: true
442             - !!python/object:hooke.engine.CommandMessage
443               arguments:
444                 arg 0: 1
445                 curve: *id001
446               command: command B
447               explicit_user_call: true
448           info: {attr with spaces: 'The second curve
449         <BLANKLINE>
450               with endlines'}
451           name: two
452           path: /path/to/curve/two
453         state:
454           _base_path: /path/to
455           info: {note: An example playlist}
456           name: playlist.hkp
457           path: /path/to/playlist.hkp
458           relative_curve_paths: false
459           version: '0.2'
460         <BLANKLINE>
461         """
462         rcp = self._relative_curve_paths
463         self._relative_curve_paths = self.relative_curve_paths
464         self.update_curve_paths()
465         self._relative_curve_paths = rcp
466         digest = self._digest
467         self._digest = None  # don't save the digest (recursive file).
468         yaml_string = yaml.dump(self, allow_unicode=True)
469         self._digest = digest
470         self.update_curve_paths()
471         return ('# Hooke playlist version %s\n' % self.version) + yaml_string
472
473     def save(self, path=None, makedirs=True):
474         """Saves the playlist to a YAML file.
475         """
476         self.set_path(path)
477         dirname = os.path.dirname(self.path) or '.'
478         if makedirs == True and not os.path.isdir(dirname):
479             os.makedirs(dirname)
480         with open(self.path, 'w') as f:
481             f.write(self.flatten())
482             self._digest = self.digest()
483
484
485 def from_string(string):
486     u"""Load a playlist from a string.
487
488     Examples
489     --------
490
491     Minimal example.
492
493     >>> string = '''# Hooke playlist version 0.2
494     ... !!python/object/new:hooke.playlist.FilePlaylist
495     ... state:
496     ...   version: '0.2'
497     ... listitems:
498     ... - !!python/object:hooke.curve.Curve
499     ...   path: curve/one
500     ... - !!python/object:hooke.curve.Curve
501     ...   path: curve/two
502     ... '''
503     >>> p = from_string(string)
504     >>> p.set_path('/path/to/playlist')
505     >>> for curve in p:
506     ...     print curve.name, curve.path
507     one /path/to/curve/one
508     two /path/to/curve/two
509
510     More complicated example.
511
512     >>> string = '''# Hooke playlist version 0.2
513     ... !!python/object/new:hooke.playlist.FilePlaylist
514     ... listitems:
515     ... - !!python/object:hooke.curve.Curve
516     ...   info: {note: The first curve}
517     ...   name: one
518     ...   path: /path/to/curve/one
519     ... - &id001 !!python/object:hooke.curve.Curve
520     ...   command_stack: !!python/object/new:hooke.command_stack.CommandStack
521     ...     listitems:
522     ...     - !!python/object:hooke.engine.CommandMessage
523     ...       arguments: {arg 0: 0, arg 1: X}
524     ...       command: command A
525     ...     - !!python/object:hooke.engine.CommandMessage
526     ...       arguments:
527     ...         arg 0: 1
528     ...         curve: *id001
529     ...       command: command B
530     ...   info: {attr with spaces: 'The second curve
531     ... 
532     ...       with endlines'}
533     ...   name: two
534     ...   path: /path/to/curve/two
535     ... state:
536     ...   _base_path: /path/to
537     ...   _index: 1
538     ...   info: {note: An example playlist}
539     ...   name: playlist.hkp
540     ...   path: /path/to/playlist.hkp
541     ...   version: '0.2'
542     ... '''
543     >>> p = from_string(string)
544     >>> p.set_path('/path/to/playlist')
545     >>> p._index
546     1
547     >>> p.info
548     {'note': 'An example playlist'}
549     >>> for curve in p:
550     ...     print curve.name, curve.path
551     one /path/to/curve/one
552     two /path/to/curve/two
553     >>> p[-1].info['attr with spaces']
554     'The second curve\\nwith endlines'
555     >>> type(p[-1].command_stack)
556     <class 'hooke.command_stack.CommandStack'>
557     >>> p[0].command_stack
558     []
559     >>> type(p[0].command_stack)
560     <class 'hooke.command_stack.CommandStack'>
561     >>> p[-1].command_stack  # doctest: +NORMALIZE_WHITESPACE
562     [<CommandMessage command A {arg 0: 0, arg 1: X}>,
563      <CommandMessage command B {arg 0: 1, curve: <Curve two>}>]
564     >>> type(p[1].command_stack)
565     <class 'hooke.command_stack.CommandStack'>
566     >>> c2 = p[-1]
567     >>> c2.command_stack[-1].arguments['curve'] == c2
568     True
569     """
570     return yaml.load(string)
571
572 def load(path=None, drivers=None, identify=True, hooke=None):
573     """Load a playlist from a file.
574     """
575     path = os.path.expanduser(playlist_path(path))
576     with open(path, 'r') as f:
577         text = f.read()
578     playlist = from_string(text)
579     playlist.set_path(path)
580     playlist._digest = playlist.digest()
581     if drivers != None:
582         playlist.drivers = drivers
583     playlist.set_path(path)
584     for curve in playlist:
585         curve.set_hooke(hooke)
586         if identify == True:
587             curve.identify(playlist.drivers)
588     return playlist
589
590
591 class Playlists (NoteIndexList):
592     """A :class:`NoteIndexList` of :class:`FilePlaylist`\s.
593     """
594     pass