Merge a80bdc9 into dd18680

holoviews/core/dimension.py

@@ -7,17 +7,18 @@
 
 import re
 import datetime as dt
+import weakref
 
 from operator import itemgetter
 from collections import defaultdict, Counter
 from itertools import chain
-from functools import reduce
+from functools import reduce, partial
 
 import param
 import numpy as np
 
 from . import util
-from .options import Store, Opts
+from .options import Store, Opts, cleanup_custom_options
 from .pprint import PrettyPrinter
 from .tree import AttrTree
 from .util import basestring, OrderedDict, bytes_to_unicode, unicode
@@ -631,6 +632,7 @@ def __init__(self, data, id=None, plot_id=None, **params):
         may be set to associate some custom options with the object.
         """
         self.data = data
+        self._id = None
         self.id = id
         self._plot_id = plot_id or util.builtins.id(self)
         if isinstance(params.get('label',None), tuple):
@@ -651,6 +653,24 @@ def __init__(self, data, id=None, plot_id=None, **params):
             raise ValueError("Supplied label %r contains invalid characters." %
                              self.label)
 
+    @property
+    def id(self):
+        return self._id
+
+    @id.setter
+    def id(self, opts_id):
+        """Handles tracking and cleanup of custom ids."""
+        old_id = self._id
+        self._id = opts_id
+        if old_id is not None:
+            cleanup_custom_options(old_id)
+        if opts_id is not None and opts_id != old_id:
+            if opts_id not in Store._weakrefs:
+                Store._weakrefs[opts_id] = []
+            ref = weakref.ref(self, partial(cleanup_custom_options, opts_id))
+            Store._weakrefs[opts_id].append(ref)
+
+
     def clone(self, data=None, shared_data=True, new_type=None, link=True,
               *args, **overrides):
         """Clones the object, overriding data and parameters.
@@ -833,14 +853,14 @@ def __getstate__(self):
         "Ensures pickles save options applied to this objects."
         obj_dict = self.__dict__.copy()
         try:
-            if Store.save_option_state and (obj_dict.get('id', None) is not None):
-                custom_key = '_custom_option_%d' % obj_dict['id']
+            if Store.save_option_state and (obj_dict.get('_id', None) is not None):
+                custom_key = '_custom_option_%d' % obj_dict['_id']
                 if custom_key not in obj_dict:
-                    obj_dict[custom_key] = {backend:s[obj_dict['id']]
+                    obj_dict[custom_key] = {backend:s[obj_dict['_id']]
                                             for backend,s in Store._custom_options.items()
-                                            if obj_dict['id'] in s}
+                                            if obj_dict['_id'] in s}
             else:
-                obj_dict['id'] = None
+                obj_dict['_id'] = None
         except:
             self.param.warning("Could not pickle custom style information.")
         return obj_dict
@@ -849,12 +869,15 @@ def __getstate__(self):
     def __setstate__(self, d):
         "Restores options applied to this object."
         d = param_aliases(d)
+
+        # Backwards compatibility for objects before id was made a property
+        opts_id = d['_id'] if '_id' in d else d.pop('id', None)
         try:
             load_options = Store.load_counter_offset is not None
             if load_options:
                 matches = [k for k in d if k.startswith('_custom_option')]
                 for match in matches:
-                    custom_id = int(match.split('_')[-1])
+                    custom_id = int(match.split('_')[-1])+Store.load_counter_offset
                     if not isinstance(d[match], dict):
                         # Backward compatibility before multiple backends
                         backend_info = {'matplotlib':d[match]}
@@ -863,20 +886,22 @@ def __setstate__(self, d):
                     for backend, info in  backend_info.items():
                         if backend not in Store._custom_options:
                             Store._custom_options[backend] = {}
-                        Store._custom_options[backend][Store.load_counter_offset + custom_id] = info
-
+                        Store._custom_options[backend][custom_id] = info
+                    if backend_info:
+                        if custom_id not in Store._weakrefs:
+                            Store._weakrefs[custom_id] = []
+                        ref = weakref.ref(self, partial(cleanup_custom_options, custom_id))
+                        Store._weakrefs[opts_id].append(ref)
                     d.pop(match)
 
-                if d['id'] is not None:
-                    d['id'] += Store.load_counter_offset
-                else:
-                    d['id'] = None
+                if opts_id is not None:
+                    opts_id += Store.load_counter_offset
         except:
             self.param.warning("Could not unpickle custom style information.")
+        d['_id'] = opts_id
         self.__dict__.update(d)
 
 
-
 class Dimensioned(LabelledData):
     """
     Dimensioned is a base class that allows the data contents of a

holoviews/core/options.py

@@ -48,6 +48,38 @@
 from .pprint import InfoPrinter, PrettyPrinter
 
 
+def cleanup_custom_options(id, weakref=None):
+    """
+    Cleans up unused custom trees if all objects referencing the
+    custom id have been garbage collected or tree is otherwise
+    unreferenced.
+    """
+    try:
+        if Store._options_context:
+            return
+        weakrefs = Store._weakrefs.get(id, [])
+        if weakref in weakrefs:
+            weakrefs.remove(weakref)
+        refs = []
+        for wr in list(weakrefs):
+            r = wr()
+            if r is None or r.id != id:
+                weakrefs.remove(wr)
+            else:
+                refs.append(r)
+        if not refs:
+            for bk in Store.loaded_backends():
+                if id in Store._custom_options[bk]:
+                    Store._custom_options[bk].pop(id)
+        if not weakrefs:
+            Store._weakrefs.pop(id, None)
+    except Exception as e:
+        raise Exception('Cleanup of custom options tree with id %s '
+                        'failed with the following exception: %s, '
+                        'an unreferenced orphan tree may persist in '
+                        'memory' % (e, id))
+
+
 class SkipRendering(Exception):
     """
     A SkipRendering exception in the plotting code will make the display
@@ -1195,6 +1227,10 @@ class Store(object):
     # populated for the given backend.
     _options = {}
 
+    # Weakrefs to record objects per id
+    _weakrefs = {}
+    _options_context = False
+
     # A list of hooks to call after registering the plot and style options
     option_setters = []
 
@@ -1545,15 +1581,20 @@ def propagate_ids(cls, obj, match_id, new_id, applied_keys, backend=None):
         matching the applied_keys. This method can only be called if
         there is a tree with a matching id in Store.custom_options
         """
-        if not new_id in Store.custom_options(backend=backend):
-            raise AssertionError("The set_ids method requires "
-                                 "Store.custom_options to contain"
-                                 " a tree with id %d" % new_id)
+        applied = []
         def propagate(o):
             if o.id == match_id or (o.__class__.__name__ == 'DynamicMap'):
                 setattr(o, 'id', new_id)
+                applied.append(o)
         obj.traverse(propagate, specs=set(applied_keys) | {'DynamicMap'})
 
+        # Clean up the custom tree if it was not applied
+        if not new_id in Store.custom_options(backend=backend):
+            raise AssertionError("New option id %d does not match any "
+                                 "option trees in Store.custom_options."
+                                 % new_id)
+        return applied
+
     @classmethod
     def capture_ids(cls, obj):
         """
@@ -1774,12 +1815,14 @@ def options(cls, obj, options=None, **kwargs):
         """
         if (options is None) and kwargs == {}: yield
         else:
+            Store._options_context = True
             optstate = cls.state(obj)
             groups = Store.options().groups.keys()
             options = cls.merge_options(groups, options, **kwargs)
             cls.set_options(obj, options)
             yield
         if options is not None:
+            Store._options_context = True
             cls.state(obj, state=optstate)
 
 
@@ -1869,8 +1912,16 @@ def set_options(cls, obj, options=None, backend=None, **kwargs):
         spec, compositor_applied = cls.expand_compositor_keys(options)
         custom_trees, id_mapping = cls.create_custom_trees(obj, spec)
         cls.update_backends(id_mapping, custom_trees, backend=backend)
-
+        
         # Propagate ids to the objects
+        not_used = []
         for (match_id, new_id) in id_mapping:
-            cls.propagate_ids(obj, match_id, new_id, compositor_applied+list(spec.keys()), backend=backend)
+            applied = cls.propagate_ids(obj, match_id, new_id, compositor_applied+list(spec.keys()), backend=backend)
+            if not applied:
+                not_used.append(new_id)
+
+        # Clean up unused custom option trees
+        for new_id in set(not_used):
+            cleanup_custom_options(new_id)
+
         return obj

holoviews/tests/core/testdimensioned.py

@@ -1,3 +1,6 @@
+import gc
+
+from holoviews.core.spaces import HoloMap
 from holoviews.core.element import Element
 from holoviews.core.options import Store, Keywords, Options, OptionTree
 from ..utils import LoggingComparisonTestCase
@@ -15,12 +18,15 @@ def setUp(self):
         self.current_backend = Store.current_backend
         self.register_custom(TestObj, 'backend_1', ['plot_custom1'])
         self.register_custom(TestObj, 'backend_2', ['plot_custom2'])
-        Store.current_backend = 'backend_1'
+        Store.set_current_backend('backend_1')
 
     def tearDown(self):
+        Store._weakrefs = {}
         Store._options.pop('backend_1')
         Store._options.pop('backend_2')
-        Store.current_backend = self.current_backend
+        Store._custom_options.pop('backend_1')
+        Store._custom_options.pop('backend_2')
+        Store.set_current_backend(self.current_backend)
 
     @classmethod
     def register_custom(cls, objtype, backend, custom_plot=[], custom_style=[]):
@@ -184,3 +190,51 @@ def test_apply_options_when_backend_switched(self):
         assert plot_opts.options == {'plot_opt1': 'D'}
         style_opts = Store.lookup_options('backend_2', obj, 'style')
         assert style_opts.options == {'style_opt1': 'C'}
+
+
+
+class TestOptionsCleanup(CustomBackendTestCase):
+
+    def test_opts_resassignment_cleans_unused_tree(self):
+        obj = TestObj([]).opts(style_opt1='A').opts(plot_opt1='B')
+        custom_options = Store._custom_options['backend_1']
+        self.assertIn(obj.id, custom_options)
+        self.assertEqual(len(custom_options), 1)
+
+    def test_opts_multiple_resassignment_cleans_unused_tree(self):
+        obj = HoloMap({0: TestObj([]), 1: TestObj([])}).opts(style_opt1='A').opts(plot_opt1='B')
+        custom_options = Store._custom_options['backend_1']
+        self.assertIn(obj.last.id, custom_options)
+        self.assertEqual(len(custom_options), 1)
+
+    def test_opts_resassignment_cleans_unused_tree_cross_backend(self):
+        obj = TestObj([]).opts(style_opt1='A').opts(plot_opt1='B', backend='backend_2')
+        custom_options = Store._custom_options['backend_1']
+        self.assertIn(obj.id, custom_options)
+        self.assertEqual(len(custom_options), 1)
+        custom_options = Store._custom_options['backend_2']
+        self.assertIn(obj.id, custom_options)
+        self.assertEqual(len(custom_options), 1)
+
+    def test_garbage_collect_cleans_unused_tree(self):
+        obj = TestObj([]).opts(style_opt1='A')
+        del obj
+        gc.collect()
+        custom_options = Store._custom_options['backend_1']
+        self.assertEqual(len(custom_options), 0)
+
+    def test_partial_garbage_collect_does_not_clear_tree(self):
+        obj = HoloMap({0: TestObj([]), 1: TestObj([])}).opts(style_opt1='A')
+        obj.pop(0)
+        gc.collect()
+        custom_options = Store._custom_options['backend_1']
+        self.assertIn(obj.last.id, custom_options)
+        self.assertEqual(len(custom_options), 1)
+        obj.pop(1)
+        gc.collect()
+        self.assertEqual(len(custom_options), 0)
+
+    def test_opts_clear_cleans_unused_tree(self):
+        TestObj([]).opts(style_opt1='A').opts.clear()
+        custom_options = Store._custom_options['backend_1']
+        self.assertEqual(len(custom_options), 0)