Added hooke.util.yaml fixing YAML/NumPy type issues (by dropping data).
[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 RepresenterError, e:
107                 raise NotImplementedError(
108                     'cannot convert %s.%s = %s (%s) to YAML\n%s'
109                     % (owner.__class__.__name__, k, v, type(v), e))
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 = ['data']
201         self._item_default_attrs = {
202             'command_stack': [],
203             'driver': None,
204             'info': {},
205             'name': None,
206             }
207         self._loaded = [] # List of loaded curves, see :meth:`._setup_item`.
208
209     def _item_getstate(self, item):
210         assert isinstance(item, curve.Curve), type(item)
211         state = item.__getstate__()
212         for key in self._item_ignored_attrs:
213             if key in state:
214                 del(state[key])
215         for key,value in self._item_default_attrs.items():
216             if key in state and state[key] == value:
217                 del(state[key])
218         return state
219
220     def _item_setstate(self, state):
221         for key,value in self._item_default_attrs.items():
222             if key not in state:
223                 state[key] = value
224         item = curve.Curve(path=None)
225         item.__setstate__(state)
226         return item
227
228     def append_curve_by_path(self, path, info=None, identify=True, hooke=None):
229         path = os.path.normpath(path)
230         c = curve.Curve(path, info=info)
231         c.set_hooke(hooke)
232         if identify == True:
233             c.identify(self.drivers)
234         self.append(c)
235         return c
236
237     def _setup_item(self, curve):
238         if curve != None and curve not in self._loaded:
239             if curve not in self:
240                 self.append(curve)
241             if curve.driver == None:
242                 c.identify(self.drivers)
243             if curve.data == None:
244                 curve.load()
245             self._loaded.append(curve)
246             if len(self._loaded) > self._max_loaded:
247                 oldest = self._loaded.pop(0)
248                 oldest.unload()
249
250
251 class FilePlaylist (Playlist):
252     """A file-backed :class:`Playlist`.
253     """
254     version = '0.2'
255
256     def __init__(self, drivers, name=None, path=None):
257         super(FilePlaylist, self).__init__(drivers, name)
258         self.path = self._base_path = None
259         self.set_path(path)
260         self._relative_curve_paths = True
261
262     def _set_ignored_attrs(self):
263         super(FilePlaylist, self)._set_ignored_attrs()
264         self._ignored_attrs.append('_digest')
265         self._digest = None
266
267     def __getstate__(self):
268         state = super(FilePlaylist, self).__getstate__()
269         assert 'version' not in state, state
270         state['version'] = self.version
271         return state
272
273     def __setstate__(self, state):
274         assert('version') in state, state
275         version = state.pop('version')
276         assert version == FilePlaylist.version, (
277             'invalid version %s (%s) != %s (%s)'
278             % (version, type(version),
279                FilePlaylist.version, type(FilePlaylist.version)))
280         super(FilePlaylist, self).__setstate__(state)
281
282     def _item_getstate(self, item):
283         state = super(FilePlaylist, self)._item_getstate(item)
284         if state.get('path', None) != None:
285             path = os.path.abspath(os.path.expanduser(state['path']))
286             if self._relative_curve_paths == True:
287                 path = os.path.relpath(path, self._base_path)
288             state['path'] = path
289         return state
290
291     def _item_setstate(self, state):
292         item = super(FilePlaylist, self)._item_setstate(state)
293         if 'path' in state:
294             item.set_path(os.path.join(self._base_path, state['path']))
295         return item
296
297     def set_path(self, path):
298         if path == None:
299             if self._base_path == None:
300                 self._base_path = os.getcwd()
301         else:
302             if not path.endswith('.hkp'):
303                 path += '.hkp'
304             self.path = path
305             self._base_path = os.path.dirname(os.path.abspath(
306                 os.path.expanduser(self.path)))
307             if self.name == None:
308                 self.name = os.path.basename(path)
309
310     def append_curve_by_path(self, path, *args, **kwargs):
311         if self._base_path != None:
312             path = os.path.join(self._base_path, path)
313         super(FilePlaylist, self).append_curve_by_path(path, *args, **kwargs)
314
315     def is_saved(self):
316         return self.digest() == self._digest
317
318     def digest(self):
319         r"""Compute the sha1 digest of the flattened playlist
320         representation.
321
322         Examples
323         --------
324
325         >>> root_path = os.path.sep + 'path'
326         >>> p = FilePlaylist(drivers=[],
327         ...                  path=os.path.join(root_path, 'to','playlist'))
328         >>> p.info['note'] = 'An example playlist'
329         >>> c = curve.Curve(os.path.join(root_path, 'to', 'curve', 'one'))
330         >>> c.info['note'] = 'The first curve'
331         >>> p.append(c)
332         >>> c = curve.Curve(os.path.join(root_path, 'to', 'curve', 'two'))
333         >>> c.info['note'] = 'The second curve'
334         >>> p.append(c)
335         >>> p.digest()
336         '\xa1\x1ax\xb1|\x84uA\xe4\x1d\xbf`\x004|\x82\xc2\xdd\xc1\x9e'
337         """
338         string = self.flatten()
339         return hashlib.sha1(string).digest()
340
341     def flatten(self):
342         """Create a string representation of the playlist.
343
344         A playlist is a YAML document with the following minimal syntax::
345
346             version: '0.2'
347             items:
348             - path: picoforce.000
349             - path: picoforce.001
350
351         Relative paths are interpreted relative to the location of the
352         playlist file.
353
354         Examples
355         --------
356
357         >>> from .engine import CommandMessage
358
359         >>> root_path = os.path.sep + 'path'
360         >>> p = FilePlaylist(drivers=[],
361         ...                  path=os.path.join(root_path, 'to','playlist'))
362         >>> p.info['note'] = 'An example playlist'
363         >>> c = curve.Curve(os.path.join(root_path, 'to', 'curve', 'one'))
364         >>> c.info['note'] = 'The first curve'
365         >>> p.append(c)
366         >>> c = curve.Curve(os.path.join(root_path, 'to', 'curve', 'two'))
367         >>> c.info['attr with spaces'] = 'The second curve\\nwith endlines'
368         >>> c.command_stack.extend([
369         ...         CommandMessage('command A', {'arg 0':0, 'arg 1':'X'}),
370         ...         CommandMessage('command B', {'arg 0':1, 'arg 1':'Y'}),
371         ...         ])
372         >>> p.append(c)
373         >>> print p.flatten()  # doctest: +REPORT_UDIFF
374         # Hooke playlist version 0.2
375         _base_path: /path/to
376         _index: 0
377         _max_loaded: 100
378         _relative_curve_paths: true
379         drivers: []
380         info: {note: An example playlist}
381         items:
382         - info: {note: The first curve}
383           name: one
384           path: curve/one
385         - command_stack: !!python/object/new:hooke.command_stack.CommandStack
386             listitems:
387             - !!python/object:hooke.engine.CommandMessage
388               arguments: {arg 0: 0, arg 1: X}
389               command: command A
390             - !!python/object:hooke.engine.CommandMessage
391               arguments: {arg 0: 1, arg 1: Y}
392               command: command B
393           info: {attr with spaces: 'The second curve
394         <BLANKLINE>
395               with endlines'}
396           name: two
397           path: curve/two
398         name: playlist.hkp
399         path: /path/to/playlist.hkp
400         version: '0.2'
401         <BLANKLINE>
402         >>> p._relative_curve_paths = False
403         >>> print p.flatten()  # doctest: +REPORT_UDIFF
404         # Hooke playlist version 0.2
405         _base_path: /path/to
406         _index: 0
407         _max_loaded: 100
408         _relative_curve_paths: false
409         drivers: []
410         info: {note: An example playlist}
411         items:
412         - info: {note: The first curve}
413           name: one
414           path: /path/to/curve/one
415         - command_stack: !!python/object/new:hooke.command_stack.CommandStack
416             listitems:
417             - !!python/object:hooke.engine.CommandMessage
418               arguments: {arg 0: 0, arg 1: X}
419               command: command A
420             - !!python/object:hooke.engine.CommandMessage
421               arguments: {arg 0: 1, arg 1: Y}
422               command: command B
423           info: {attr with spaces: 'The second curve
424         <BLANKLINE>
425               with endlines'}
426           name: two
427           path: /path/to/curve/two
428         name: playlist.hkp
429         path: /path/to/playlist.hkp
430         version: '0.2'
431         <BLANKLINE>
432         """
433         yaml_string = yaml.dump(self.__getstate__(), allow_unicode=True)
434         return ('# Hooke playlist version %s\n' % self.version) + yaml_string
435
436     def from_string(self, string):
437         u"""Load a playlist from a string.
438
439         Examples
440         --------
441
442         Minimal example.
443
444         >>> string = '''# Hooke playlist version 0.2
445         ... version: '0.2'
446         ... items:
447         ... - path: picoforce.000
448         ... - path: picoforce.001
449         ... '''
450         >>> p = FilePlaylist(drivers=[],
451         ...                 path=os.path.join('/path', 'to', 'my', 'playlist'))
452         >>> p.from_string(string)
453         >>> for curve in p:
454         ...     print curve.path
455         /path/to/my/picoforce.000
456         /path/to/my/picoforce.001
457
458         More complicated example.
459
460         >>> string = '''# Hooke playlist version 0.2
461         ... _base_path: /path/to
462         ... _digest: null
463         ... _index: 1
464         ... _max_loaded: 100
465         ... _relative_curve_paths: true
466         ... info: {note: An example playlist}
467         ... items:
468         ... - info: {note: The first curve}
469         ...   path: curve/one
470         ... - command_stack: !!python/object/new:hooke.command_stack.CommandStack
471         ...      listitems:
472         ...      - !!python/object:hooke.engine.CommandMessage
473         ...        arguments: {arg 0: 0, arg 1: X}
474         ...        command: command A
475         ...      - !!python/object:hooke.engine.CommandMessage
476         ...        arguments: {arg 0: 1, arg 1: Y}
477         ...        command: command B
478         ...   info: {attr with spaces: 'The second curve
479         ... 
480         ...       with endlines'}
481         ...   name: two
482         ...   path: curve/two
483         ... name: playlist.hkp
484         ... path: /path/to/playlist.hkp
485         ... version: '0.2'
486         ... '''
487         >>> p = FilePlaylist(drivers=[],
488         ...                  path=os.path.join('path', 'to', 'my', 'playlist'))
489         >>> p.from_string(string)
490         >>> p._index
491         1
492         >>> p.info
493         {'note': 'An example playlist'}
494         >>> for curve in p:
495         ...     print curve.name, curve.path
496         one /path/to/curve/one
497         two /path/to/curve/two
498         >>> p[-1].info['attr with spaces']
499         'The second curve\\nwith endlines'
500         >>> type(p[-1].command_stack)
501         <class 'hooke.command_stack.CommandStack'>
502         >>> p[-1].command_stack  # doctest: +NORMALIZE_WHITESPACE
503         [<CommandMessage command A {arg 0: 0, arg 1: X}>,
504          <CommandMessage command B {arg 0: 1, arg 1: Y}>]
505         """
506         state = yaml.load(string)
507         self.__setstate__(state)
508
509     def save(self, path=None, makedirs=True):
510         """Saves the playlist to a YAML file.
511         """
512         self.set_path(path)
513         dirname = os.path.dirname(self.path) or '.'
514         if makedirs == True and not os.path.isdir(dirname):
515             os.makedirs(dirname)
516         with open(self.path, 'w') as f:
517             f.write(self.flatten())
518             self._digest = self.digest()
519
520     def load(self, path=None, identify=True, hooke=None):
521         """Load a playlist from a file.
522         """
523         self.set_path(path)
524         with open(self.path, 'r') as f:
525             text = f.read()
526         self.from_string(text)
527         self._digest = self.digest()
528         for curve in self:
529             curve.set_hooke(hooke)
530             if identify == True:
531                 curve.identify(self.drivers)
532
533
534 class Playlists (NoteIndexList):
535     """A :class:`NoteIndexList` of :class:`FilePlaylist`\s.
536     """
537     def __init__(self, *arg, **kwargs):
538         super(Playlists, self).__init__(*arg, **kwargs)
539
540     def _item_getstate(self, item):
541         assert isinstance(item, FilePlaylist), type(item)
542         return item.__getstate__()
543
544     def _item_setstate(self, state):
545         item = FilePlaylist(drivers=[])
546         item.__setstate__(state)
547         return item