9ee31259f780e633eb0a25afc0148bd46a244901
[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.dump((k,v))
106             except (TypeError, RepresenterError), e:
107                 raise NotImplementedError(
108                     'cannot convert %s.%s = %s (%s) to 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: !!python/object/new:hooke.command_stack.CommandStack
387             listitems:
388             - !!python/object:hooke.engine.CommandMessage
389               arguments: {arg 0: 0, arg 1: X}
390               command: command A
391             - !!python/object:hooke.engine.CommandMessage
392               arguments: {arg 0: 1, arg 1: Y}
393               command: command B
394           info: {attr with spaces: 'The second curve
395         <BLANKLINE>
396               with endlines'}
397           name: two
398           path: curve/two
399         name: playlist.hkp
400         path: /path/to/playlist.hkp
401         version: '0.2'
402         <BLANKLINE>
403         >>> p._relative_curve_paths = False
404         >>> print p.flatten()  # doctest: +REPORT_UDIFF
405         # Hooke playlist version 0.2
406         _base_path: /path/to
407         _index: 0
408         _max_loaded: 100
409         _relative_curve_paths: false
410         drivers: []
411         info: {note: An example playlist}
412         items:
413         - info: {note: The first curve}
414           name: one
415           path: /path/to/curve/one
416         - command_stack: !!python/object/new:hooke.command_stack.CommandStack
417             listitems:
418             - !!python/object:hooke.engine.CommandMessage
419               arguments: {arg 0: 0, arg 1: X}
420               command: command A
421             - !!python/object:hooke.engine.CommandMessage
422               arguments: {arg 0: 1, arg 1: Y}
423               command: command B
424           info: {attr with spaces: 'The second curve
425         <BLANKLINE>
426               with endlines'}
427           name: two
428           path: /path/to/curve/two
429         name: playlist.hkp
430         path: /path/to/playlist.hkp
431         version: '0.2'
432         <BLANKLINE>
433         """
434         yaml_string = yaml.dump(self.__getstate__(), allow_unicode=True)
435         return ('# Hooke playlist version %s\n' % self.version) + yaml_string
436
437     def from_string(self, string):
438         u"""Load a playlist from a string.
439
440         Examples
441         --------
442
443         Minimal example.
444
445         >>> string = '''# Hooke playlist version 0.2
446         ... version: '0.2'
447         ... items:
448         ... - path: picoforce.000
449         ... - path: picoforce.001
450         ... '''
451         >>> p = FilePlaylist(drivers=[],
452         ...                 path=os.path.join('/path', 'to', 'my', 'playlist'))
453         >>> p.from_string(string)
454         >>> for curve in p:
455         ...     print curve.path
456         /path/to/my/picoforce.000
457         /path/to/my/picoforce.001
458
459         More complicated example.
460
461         >>> string = '''# Hooke playlist version 0.2
462         ... _base_path: /path/to
463         ... _digest: null
464         ... _index: 1
465         ... _max_loaded: 100
466         ... _relative_curve_paths: true
467         ... info: {note: An example playlist}
468         ... items:
469         ... - info: {note: The first curve}
470         ...   path: curve/one
471         ... - command_stack: !!python/object/new:hooke.command_stack.CommandStack
472         ...      listitems:
473         ...      - !!python/object:hooke.engine.CommandMessage
474         ...        arguments: {arg 0: 0, arg 1: X}
475         ...        command: command A
476         ...      - !!python/object:hooke.engine.CommandMessage
477         ...        arguments: {arg 0: 1, arg 1: Y}
478         ...        command: command B
479         ...   info: {attr with spaces: 'The second curve
480         ... 
481         ...       with endlines'}
482         ...   name: two
483         ...   path: curve/two
484         ... name: playlist.hkp
485         ... path: /path/to/playlist.hkp
486         ... version: '0.2'
487         ... '''
488         >>> p = FilePlaylist(drivers=[],
489         ...                  path=os.path.join('path', 'to', 'my', 'playlist'))
490         >>> p.from_string(string)
491         >>> p._index
492         1
493         >>> p.info
494         {'note': 'An example playlist'}
495         >>> for curve in p:
496         ...     print curve.name, curve.path
497         one /path/to/curve/one
498         two /path/to/curve/two
499         >>> p[-1].info['attr with spaces']
500         'The second curve\\nwith endlines'
501         >>> type(p[-1].command_stack)
502         <class 'hooke.command_stack.CommandStack'>
503         >>> p[-1].command_stack  # doctest: +NORMALIZE_WHITESPACE
504         [<CommandMessage command A {arg 0: 0, arg 1: X}>,
505          <CommandMessage command B {arg 0: 1, arg 1: Y}>]
506         """
507         state = yaml.load(string)
508         self.__setstate__(state)
509
510     def save(self, path=None, makedirs=True):
511         """Saves the playlist to a YAML file.
512         """
513         self.set_path(path)
514         dirname = os.path.dirname(self.path) or '.'
515         if makedirs == True and not os.path.isdir(dirname):
516             os.makedirs(dirname)
517         with open(self.path, 'w') as f:
518             f.write(self.flatten())
519             self._digest = self.digest()
520
521     def load(self, path=None, identify=True, hooke=None):
522         """Load a playlist from a file.
523         """
524         self.set_path(path)
525         with open(self.path, 'r') as f:
526             text = f.read()
527         self.from_string(text)
528         self._digest = self.digest()
529         for curve in self:
530             curve.set_hooke(hooke)
531             if identify == True:
532                 curve.identify(self.drivers)
533
534
535 class Playlists (NoteIndexList):
536     """A :class:`NoteIndexList` of :class:`FilePlaylist`\s.
537     """
538     def __init__(self, *arg, **kwargs):
539         super(Playlists, self).__init__(*arg, **kwargs)
540
541     def _item_getstate(self, item):
542         assert isinstance(item, FilePlaylist), type(item)
543         return item.__getstate__()
544
545     def _item_setstate(self, state):
546         item = FilePlaylist(drivers=[])
547         item.__setstate__(state)
548         return item