Add hooke.playlist.Playlists for better playlist serialization.
[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 . import curve as curve
34 from .util.itertools import reverse_enumerate
35
36
37 class NoteIndexList (list):
38     """A list that keeps track of a "current" item and additional notes.
39
40     :attr:`index` (i.e. "bookmark") is the index of the currently
41     current curve.  Also keep a :class:`dict` of additional information
42     (:attr:`info`).
43     """
44     def __init__(self, name=None):
45         super(NoteIndexList, self).__init__()
46         self.name = name
47         self.info = {}
48         self._index = 0
49         self._set_ignored_attrs()
50
51     def __str__(self):
52         return str(self.__unicode__())
53
54     def __unicode__(self):
55         return u'<%s %s>' % (self.__class__.__name__, self.name)
56
57     def __repr__(self):
58         return self.__str__()
59
60     def _set_ignored_attrs(self):
61         self._ignored_attrs = ['_ignored_attrs', '_default_attrs']
62         self._default_attrs = {
63             'info': {},
64             }
65
66     def __getstate__(self):
67         state = dict(self.__dict__)
68         for key in self._ignored_attrs:
69             if key in state:
70                 del(state[key])
71         for key,value in self._default_attrs.items():
72             if key in state and state[key] == value:
73                 del(state[key])
74         assert 'items' not in state
75         state['items'] = []
76         self._assert_clean_state(self, state)
77         for item in self:  # save curves and their attributes
78             item_state = self._item_getstate(item)
79             self._assert_clean_state(item, item_state)
80             state['items'].append(item_state)
81         return state
82
83     def __setstate__(self, state):
84         self._set_ignored_attrs()
85         for key,value in self._default_attrs.items():
86             setattr(self, key, value)
87         for key,value in state.items():
88             if key == 'items':
89                 continue
90             setattr(self, key, value)
91         for item_state in state['items']:
92             self.append(self._item_setstate(item_state))
93
94     def _item_getstate(self, item):
95         return item
96
97     def _item_setstate(self, state):
98         return state
99
100     def _assert_clean_state(self, owner, state):
101         for k,v in state.items():
102             if k == 'drivers':  # HACK.  Need better driver serialization.
103                 continue
104             try:
105                 yaml.safe_dump((k,v))
106             except RepresenterError, e:
107                 raise NotImplementedError(
108                     'cannot convert %s.%s = %s (%s) to safe YAML'
109                     % (owner.__class__.__name__, k, v, type(v)))
110
111     def _setup_item(self, item):
112         """Perform any required initialization before returning an item.
113         """
114         pass
115
116     def index(self, value=None, *args, **kwargs):
117         """Extend `list.index`, returning the current index if `value`
118         is `None`.
119         """
120         if value == None:
121             return self._index
122         return super(NoteIndexList, self).index(value, *args, **kwargs)
123
124     def current(self):
125         if len(self) == 0:
126             return None
127         item = self[self._index]
128         self._setup_item(item)
129         return item
130
131     def jump(self, index):
132         if len(self) == 0:
133             self._index = 0
134         else:
135             self._index = index % len(self)
136
137     def next(self):
138         self.jump(self._index + 1)
139
140     def previous(self):
141         self.jump(self._index - 1)
142
143     def items(self, reverse=False):
144         """Iterate through `self` calling `_setup_item` on each item
145         before yielding.
146
147         Notes
148         -----
149         Updates :attr:`_index` during the iteration so
150         :func:`~hooke.plugin.curve.current_curve_callback` works as
151         expected in :class:`~hooke.command.Command`\s called from
152         :class:`~hooke.plugin.playlist.ApplyCommand`.  After the
153         iteration completes, :attr:`_index` is restored to its
154         original value.
155         """
156         index = self._index
157         items = self
158         if reverse == True:
159             items = reverse_enumerate(self)
160         else:
161             items = enumerate(self)
162         for i,item in items:
163             self._index = i
164             self._setup_item(item)
165             yield item
166         self._index = index
167
168     def filter(self, keeper_fn=lambda item:True, load_curves=True,
169                *args, **kwargs):
170         c = copy.deepcopy(self)
171         if load_curves == True:
172             items = c.items(reverse=True)
173         else:
174             items = reversed(c)
175         for item in items: 
176             if keeper_fn(item, *args, **kwargs) != True:
177                 c.remove(item)
178         try: # attempt to maintain the same current item
179             c._index = c.index(self.current())
180         except ValueError:
181             c._index = 0
182         return c
183
184
185 class Playlist (NoteIndexList):
186     """A :class:`NoteIndexList` of :class:`hooke.curve.Curve`\s.
187
188     Keeps a list of :attr:`drivers` for loading curves.
189     """
190     def __init__(self, drivers, name=None):
191         super(Playlist, self).__init__(name=name)
192         self.drivers = drivers
193         self._max_loaded = 100 # curves to hold in memory simultaneously.
194
195     def _set_ignored_attrs(self):
196         super(Playlist, self)._set_ignored_attrs()
197         self._ignored_attrs.extend([
198                 '_item_ignored_attrs', '_item_default_attrs',
199                 '_loaded'])
200         self._item_ignored_attrs = []
201         self._item_default_attrs = {
202             'command_stack': [],
203             'data': None,
204             'driver': None,
205             'info': {},
206             'name': None,
207             }
208         self._loaded = [] # List of loaded curves, see :meth:`._setup_item`.
209
210     def _item_getstate(self, item):
211         assert isinstance(item, curve.Curve), type(item)
212         state = item.__getstate__()
213         for key in self._item_ignored_attrs:
214             if key in state:
215                 del(state[key])
216         for key,value in self._item_default_attrs.items():
217             if key in state and state[key] == value:
218                 del(state[key])
219         return state
220
221     def _item_setstate(self, state):
222         for key,value in self._item_default_attrs.items():
223             if key not in state:
224                 state[key] = value
225         item = curve.Curve(path=None)
226         item.__setstate__(state)
227         return item
228
229     def append_curve_by_path(self, path, info=None, identify=True, hooke=None):
230         path = os.path.normpath(path)
231         c = curve.Curve(path, info=info)
232         c.set_hooke(hooke)
233         if identify == True:
234             c.identify(self.drivers)
235         self.append(c)
236         return c
237
238     def _setup_item(self, curve):
239         if curve != None and curve not in self._loaded:
240             if curve not in self:
241                 self.append(curve)
242             if curve.driver == None:
243                 c.identify(self.drivers)
244             if curve.data == None:
245                 curve.load()
246             self._loaded.append(curve)
247             if len(self._loaded) > self._max_loaded:
248                 oldest = self._loaded.pop(0)
249                 oldest.unload()
250
251
252 class FilePlaylist (Playlist):
253     """A file-backed :class:`Playlist`.
254     """
255     version = '0.2'
256
257     def __init__(self, drivers, name=None, path=None):
258         super(FilePlaylist, self).__init__(drivers, name)
259         self.path = self._base_path = None
260         self.set_path(path)
261         self._relative_curve_paths = True
262
263     def _set_ignored_attrs(self):
264         super(FilePlaylist, self)._set_ignored_attrs()
265         self._ignored_attrs.append('_digest')
266         self._digest = None
267
268     def __getstate__(self):
269         state = super(FilePlaylist, self).__getstate__()
270         assert 'version' not in state, state
271         state['version'] = self.version
272         return state
273
274     def __setstate__(self, state):
275         assert('version') in state, state
276         version = state.pop('version')
277         assert version == FilePlaylist.version, (
278             'invalid version %s (%s) != %s (%s)'
279             % (version, type(version),
280                FilePlaylist.version, type(FilePlaylist.version)))
281         super(FilePlaylist, self).__setstate__(state)
282
283     def _item_getstate(self, item):
284         state = super(FilePlaylist, self)._item_getstate(item)
285         if state.get('path', None) != None:
286             path = os.path.abspath(os.path.expanduser(state['path']))
287             if self._relative_curve_paths == True:
288                 path = os.path.relpath(path, self._base_path)
289             state['path'] = path
290         return state
291
292     def _item_setstate(self, state):
293         item = super(FilePlaylist, self)._item_setstate(state)
294         if 'path' in state:
295             item.set_path(os.path.join(self._base_path, state['path']))
296         return item
297
298     def set_path(self, path):
299         if path == None:
300             if self._base_path == None:
301                 self._base_path = os.getcwd()
302         else:
303             if not path.endswith('.hkp'):
304                 path += '.hkp'
305             self.path = path
306             self._base_path = os.path.dirname(os.path.abspath(
307                 os.path.expanduser(self.path)))
308             if self.name == None:
309                 self.name = os.path.basename(path)
310
311     def append_curve_by_path(self, path, *args, **kwargs):
312         if self._base_path != None:
313             path = os.path.join(self._base_path, path)
314         super(FilePlaylist, self).append_curve_by_path(path, *args, **kwargs)
315
316     def is_saved(self):
317         return self.digest() == self._digest
318
319     def digest(self):
320         r"""Compute the sha1 digest of the flattened playlist
321         representation.
322
323         Examples
324         --------
325
326         >>> root_path = os.path.sep + 'path'
327         >>> p = FilePlaylist(drivers=[],
328         ...                  path=os.path.join(root_path, 'to','playlist'))
329         >>> p.info['note'] = 'An example playlist'
330         >>> c = curve.Curve(os.path.join(root_path, 'to', 'curve', 'one'))
331         >>> c.info['note'] = 'The first curve'
332         >>> p.append(c)
333         >>> c = curve.Curve(os.path.join(root_path, 'to', 'curve', 'two'))
334         >>> c.info['note'] = 'The second curve'
335         >>> p.append(c)
336         >>> p.digest()
337         '\xa1\x1ax\xb1|\x84uA\xe4\x1d\xbf`\x004|\x82\xc2\xdd\xc1\x9e'
338         """
339         string = self.flatten()
340         return hashlib.sha1(string).digest()
341
342     def flatten(self):
343         """Create a string representation of the playlist.
344
345         A playlist is a YAML document with the following minimal syntax::
346
347             version: '0.2'
348             items:
349             - path: picoforce.000
350             - path: picoforce.001
351
352         Relative paths are interpreted relative to the location of the
353         playlist file.
354
355         Examples
356         --------
357
358         >>> from .engine import CommandMessage
359
360         >>> root_path = os.path.sep + 'path'
361         >>> p = FilePlaylist(drivers=[],
362         ...                  path=os.path.join(root_path, 'to','playlist'))
363         >>> p.info['note'] = 'An example playlist'
364         >>> c = curve.Curve(os.path.join(root_path, 'to', 'curve', 'one'))
365         >>> c.info['note'] = 'The first curve'
366         >>> p.append(c)
367         >>> c = curve.Curve(os.path.join(root_path, 'to', 'curve', 'two'))
368         >>> c.info['attr with spaces'] = 'The second curve\\nwith endlines'
369         >>> c.command_stack.extend([
370         ...         CommandMessage('command A', {'arg 0':0, 'arg 1':'X'}),
371         ...         CommandMessage('command B', {'arg 0':1, 'arg 1':'Y'}),
372         ...         ])
373         >>> p.append(c)
374         >>> print p.flatten()  # doctest: +REPORT_UDIFF
375         # Hooke playlist version 0.2
376         _base_path: /path/to
377         _index: 0
378         _max_loaded: 100
379         _relative_curve_paths: true
380         drivers: []
381         info: {note: An example playlist}
382         items:
383         - info: {note: The first curve}
384           name: one
385           path: curve/one
386         - command_stack:
387           - arguments: {arg 0: 0, arg 1: X}
388             command: command A
389           - arguments: {arg 0: 1, arg 1: Y}
390             command: command B
391           info: {attr with spaces: 'The second curve
392         <BLANKLINE>
393               with endlines'}
394           name: two
395           path: curve/two
396         name: playlist.hkp
397         path: /path/to/playlist.hkp
398         version: '0.2'
399         <BLANKLINE>
400         >>> p._relative_curve_paths = False
401         >>> print p.flatten()  # doctest: +REPORT_UDIFF
402         # Hooke playlist version 0.2
403         _base_path: /path/to
404         _index: 0
405         _max_loaded: 100
406         _relative_curve_paths: false
407         drivers: []
408         info: {note: An example playlist}
409         items:
410         - info: {note: The first curve}
411           name: one
412           path: /path/to/curve/one
413         - command_stack:
414           - arguments: {arg 0: 0, arg 1: X}
415             command: command A
416           - arguments: {arg 0: 1, arg 1: Y}
417             command: command B
418           info: {attr with spaces: 'The second curve
419         <BLANKLINE>
420               with endlines'}
421           name: two
422           path: /path/to/curve/two
423         name: playlist.hkp
424         path: /path/to/playlist.hkp
425         version: '0.2'
426         <BLANKLINE>
427         """
428         yaml_string = yaml.dump(self.__getstate__(), allow_unicode=True)
429         return ('# Hooke playlist version %s\n' % self.version) + yaml_string
430
431     def from_string(self, string):
432         u"""Load a playlist from a string.
433
434         Examples
435         --------
436
437         Minimal example.
438
439         >>> string = '''# Hooke playlist version 0.2
440         ... version: '0.2'
441         ... items:
442         ... - path: picoforce.000
443         ... - path: picoforce.001
444         ... '''
445         >>> p = FilePlaylist(drivers=[],
446         ...                 path=os.path.join('/path', 'to', 'my', 'playlist'))
447         >>> p.from_string(string)
448         >>> for curve in p:
449         ...     print curve.path
450         /path/to/my/picoforce.000
451         /path/to/my/picoforce.001
452
453         More complicated example.
454
455         >>> string = '''# Hooke playlist version 0.2
456         ... _base_path: /path/to
457         ... _digest: null
458         ... _index: 1
459         ... _max_loaded: 100
460         ... _relative_curve_paths: true
461         ... info: {note: An example playlist}
462         ... items:
463         ... - info: {note: The first curve}
464         ...   path: curve/one
465         ... - command_stack:
466         ...   - arguments: {arg 0: 0, arg 1: X}
467         ...     command: command A
468         ...   - arguments: {arg 0: 1, arg 1: Y}
469         ...     command: command B
470         ...   info: {attr with spaces: 'The second curve
471         ... 
472         ...       with endlines'}
473         ...   name: two
474         ...   path: curve/two
475         ... name: playlist.hkp
476         ... path: /path/to/playlist.hkp
477         ... version: '0.2'
478         ... '''
479         >>> p = FilePlaylist(drivers=[],
480         ...                  path=os.path.join('path', 'to', 'my', 'playlist'))
481         >>> p.from_string(string)
482         >>> p._index
483         1
484         >>> p.info
485         {'note': 'An example playlist'}
486         >>> for curve in p:
487         ...     print curve.name, curve.path
488         one /path/to/curve/one
489         two /path/to/curve/two
490         >>> p[-1].info['attr with spaces']
491         'The second curve\\nwith endlines'
492         >>> type(p[-1].command_stack)
493         <class 'hooke.command_stack.CommandStack'>
494         >>> p[-1].command_stack  # doctest: +NORMALIZE_WHITESPACE
495         [<CommandMessage command A {arg 0: 0, arg 1: X}>,
496          <CommandMessage command B {arg 0: 1, arg 1: Y}>]
497         """
498         state = yaml.load(string)
499         self.__setstate__(state)
500
501     def save(self, path=None, makedirs=True):
502         """Saves the playlist to a YAML file.
503         """
504         self.set_path(path)
505         dirname = os.path.dirname(self.path) or '.'
506         if makedirs == True and not os.path.isdir(dirname):
507             os.makedirs(dirname)
508         with open(self.path, 'w') as f:
509             f.write(self.flatten())
510             self._digest = self.digest()
511
512     def load(self, path=None, identify=True, hooke=None):
513         """Load a playlist from a file.
514         """
515         self.set_path(path)
516         with open(self.path, 'r') as f:
517             text = f.read()
518         self.from_string(text)
519         self._digest = self.digest()
520         for curve in self:
521             curve.set_hooke(hooke)
522             if identify == True:
523                 curve.identify(self.drivers)
524
525
526 class Playlists (NoteIndexList):
527     """A :class:`NoteIndexList` of :class:`FilePlaylist`\s.
528     """
529     def __init__(self, *arg, **kwargs):
530         super(Playlists, self).__init__(*arg, **kwargs)
531
532     def _item_getstate(self, item):
533         assert isinstance(item, FilePlaylist), type(item)
534         return item.__getstate__()
535
536     def _item_setstate(self, state):
537         item = FilePlaylist(drivers=[])
538         item.__setstate__(state)
539         return item