Clean up code and fix tests.

vetschn · Nov 10, 2021 · c57c7e7 · c57c7e7
1 parent ee1bdc1
commit c57c7e7
Show file tree

Hide file tree

Showing 2 changed files with 27 additions and 11 deletions.
diff --git a/mayavi/filters/user_defined.py b/mayavi/filters/user_defined.py
@@ -3,7 +3,7 @@
 # License: BSD Style.
 
 # Enthought library imports.
-from tvtk.tools.tvtk_doc import TVTKFilterChooser, TVTK_FILTERS
+from tvtk.tools.tvtk_doc import TVTKFilterChooser, get_tvtk_filters
 
 # Local imports.
 from mayavi.filters.filter_base import FilterBase
@@ -73,11 +73,11 @@ def _choose_filter(self):
     def _check_object(self, obj):
         if obj is None:
             return False
-        if obj.__class__.__name__ in TVTK_FILTERS:
+        tvtk_filters = get_tvtk_filters()
+        if obj.__class__.__name__ in tvtk_filters:
             return True
         return False
 
     def _filter_changed(self, old, new):
         self.name = 'UserDefined:%s'%new.__class__.__name__
         super(UserDefined, self)._filter_changed(old, new)
-
diff --git a/tvtk/tools/tvtk_doc.py b/tvtk/tools/tvtk_doc.py
@@ -148,6 +148,26 @@ def _setup_tvtk_names():
         TVTK_CLASSES, TVTK_SOURCES, TVTK_FILTERS, TVTK_SINKS = r
 
 
+def get_tvtk_classes():
+    _setup_tvtk_names()
+    return TVTK_CLASSES
+
+
+def get_tvtk_sources():
+    _setup_tvtk_names()
+    return TVTK_SOURCES
+
+
+def get_tvtk_filters():
+    _setup_tvtk_names()
+    return TVTK_FILTERS
+
+
+def get_tvtk_sinks():
+    _setup_tvtk_names()
+    return TVTK_SINKS
+
+
 # ##############################################################################
 # `DocSearch` class.
 # ##############################################################################
@@ -395,8 +415,7 @@ def _search_changed(self, value):
             self.completions = result[:self.n_completion]
 
     def _available_default(self):
-        _setup_tvtk_names()
-        return TVTK_CLASSES
+        return get_tvtk_classes()
 
 
 ################################################################################
@@ -406,8 +425,7 @@ class TVTKSourceChooser(TVTKClassChooser):
     available = List
 
     def _available_default(self):
-        _setup_tvtk_names()
-        return TVTK_SOURCES
+        return get_tvtk_sources()
 
 
 # ##############################################################################
@@ -417,8 +435,7 @@ class TVTKFilterChooser(TVTKClassChooser):
     available = List
 
     def _available_default(self):
-        _setup_tvtk_names()
-        return TVTK_FILTERS
+        return get_tvtk_filters()
 
 
 # ##############################################################################
@@ -428,8 +445,7 @@ class TVTKSinkChooser(TVTKClassChooser):
     available = List
 
     def _available_default(self):
-        _setup_tvtk_names()
-        return TVTK_SINKS
+        return get_tvtk_sinks()
 
 
 def main():