enthought / mayavi
1
# Author: Prabhu Ramachandran
2
# Copyright (c) 2004-2020, Enthought, Inc.
3
# License: BSD Style.
4

5 4
"""This module generates the class hierarchy for any given Python
6
modules.  This can be used (among other things) to generate the
7
traitified VTK classes in the correct order.
8

9
"""
10

11 4
import sys
12

13 4
if sys.version_info[0] > 2:
14 4
    import builtins
15
else:
16 0
    import __builtin__ as builtins
17

18

19 4
class TreeNode:
20
    """Represents a node in the class tree.
21

22
    Stores information on the sub and super classes of a particular
23
    class.  It also stores the inheritance level of the class inside
24
    the inheritance tree (essentially how many levels of inheritance
25
    are there below this class).  This inheritance level is computed
26
    when the `get_level` method is called.  The `get_level` method
27
    works only when the parent information is complete.
28

29
    """
30

31 4
    def __init__(self, klass):
32
        """Given a class, create a node in the tree.
33

34
        Parameters
35
        ----------
36
        - klass : `class`
37

38
          The class which is represented as a node in the tree.
39

40
        """
41 4
        self.klass = klass
42 4
        self.name = klass.__name__
43 4
        self.children = []
44 4
        self.parents = []
45
        # Uninitialized level is set to None
46 4
        self.level = None
47

48 4
    def add_parent(self, parent):
49
        """Add a parent node."""
50 4
        assert isinstance(parent, TreeNode)
51 4
        if parent not in self.parents:
52 4
            self.parents.append(parent)
53

54 4
    def add_child(self, child):
55
        """Add a child node. """
56 4
        assert isinstance(child, TreeNode)
57 4
        if child not in self.children:
58 4
            self.children.append(child)
59

60 4
    def get_level(self):
61
        """Returns the inheritance level of the node (an int).  If the
62
        level has not been set, the method computes it.  Note however,
63
        that this computation will fail if the parent information is
64
        incorrect.
65

66
        """
67 4
        if not self.level:
68 4
            if self.parents:
69 4
                self.level = max([x.get_level() for x in self.parents]) + 1
70
            else:
71 4
                self.level = 0
72 4
        return self.level
73

74 4
    def get_ancestors(self):
75
        """Returns a list of ancestor nodes from which this class has
76
        descended.
77

78
        """
79 4
        def _get_ancestors(node, ancestors):
80 4
            ancestors.extend(node.parents)
81 4
            for p in node.parents:
82 4
                _get_ancestors(p, ancestors)
83 4
        ancestors = []
84 4
        _get_ancestors(self, ancestors)
85 4
        return ancestors
86

87

88 4
class ClassTree:
89
    """Contains and generates all the class tree information.
90

91
    On initialization of the instance, nothing is done.  The classes
92
    are obtained using the list of modules (or a single module) that
93
    is used to initialize the instance.  One must then call the
94
    `create` method to generate the tree structure.  The instance of
95
    the class also can be treated as an iterator which iterates over
96
    the nodes of the tree.
97

98
    There are two ways in which the tree hierarchy is stored.  A
99
    dictionary mapping class names to the tree node and a tree
100
    represented as a list of lists containing the nodes.  The tree is
101
    organized based on a concept of an inheritance level.  A class
102
    that has no parent classes (no base classes) is said to be at
103
    level zero.  If a class inherits successively from 7 classes, it
104
    is at level 6.  An example of inheritance for a vtkFoo class is
105
    given below:
106

107
      vtkFoo -> vtkBar -> vtkObject -> vtkObjectBase
108

109
    Here, vtkObjectBase has an inheritance level of 0 and vtkFoo a
110
    level of 3.  One can traverse the tree by using the level as an
111
    index and find all the classes at a particular level.
112

113
    Here is some example usage of this class::
114

115
        >>> import vtk
116
        >>> t = ClassTree(vtk)
117
        >>> t.create()
118
        >>> print t.get_node('vtkObject').name
119
        vtkObject
120
        >>> print t.get_node('vtkObject').parents[0].name
121
        vtkObjectBase
122
        >>> print len(t.tree[0])
123
        1
124
        >>> t.tree[0][0].name
125
        vtkObjectBase
126

127
    """
128

129 4
    def __init__(self, modules):
130
        """Initialize the instance with the given modules.
131

132
        Parameters
133
        ----------
134

135
        - modules : `sequence` of modules or a module
136

137
          This is either a single module or a sequence of modules.
138
          The instance uses these list of modules to generate the
139
          class tree.
140

141
        """
142 4
        self.nodes = {}
143 4
        self.tree = [[]]
144 4
        if not hasattr(modules, '__iter__'):
145 4
            self.modules = [modules]
146
        else:
147 0
            self.modules = modules
148

149 4
    def __iter__(self):
150 4
        return iter(self.nodes.values())
151

152 4
    def _generate_hierarchy(self, klass):
153
        """Does the hard work of generating the class hierarchy."""
154 4
        node = self.get_node(klass.__name__, create=1)
155 4
        for base in klass.__bases__:
156 4
            base_node = self.get_node_from_class(base, create=1)
157 4
            node.add_parent(base_node)
158 4
            base_node.add_child(node)
159 4
            self._generate_hierarchy(base)
160

161 4
    def get_class(self, name):
162
        """Given a class name in the given modules returns the class."""
163 4
        klass = None
164 4
        for m in self.modules:
165 4
            if hasattr(m, name):
166 4
                return getattr(m, name)
167 4
        if hasattr(builtins, name):
168 4
            klass = getattr(builtins, name)
169 4
        if not klass:
170 4
            try:
171 4
                klass = self.nodes[name].klass
172 0
            except KeyError:
173 0
                raise KeyError("Cannot find class of name %s"%name)
174 4
        return klass
175

176 4
    def add_node(self, klass):
177
        """Create a node for the given class."""
178 4
        name = klass.__name__
179 4
        if not name in self.nodes:
180 4
            node = TreeNode(klass)
181 4
            self.nodes[name] = node
182 4
            return node
183

184 4
    def get_node(self, name, create=0):
185
        """Get a node of the given name.
186

187
        Parameters
188
        ----------
189

190
        - name : `str`
191

192
          Name of the node to get.
193

194
        - create : `boolean`
195

196
          If True, a new node will be added if no node of the given
197
          name is available.  Defaults to False.
198

199
        Returns
200
        -------
201

202
        - `TreeNode`
203

204
        """
205 4
        if name in self.nodes:
206 4
            return self.nodes[name]
207 4
        elif create:
208 4
            return self.add_node(self.get_class(name))
209

210 4
    def get_node_from_class(self, cls, create=0):
211
        """Get a node of the given class.
212

213
        Parameters
214
        ----------
215

216
        - cls : `class`
217

218
          Class of the node to get.
219

220
        - create : `boolean`
221

222
          If True, a new node will be added if no node of the given
223
          name is available.  Defaults to False.
224

225
        Returns
226
        -------
227

228
        - `TreeNode`
229

230
        """
231 4
        name = cls.__name__
232 4
        if name in self.nodes:
233 4
            return self.nodes[name]
234 4
        elif create:
235 4
            return self.add_node(cls)
236

237 4
    def create(self, class_names=None):
238
        """This method generates the class tree given an optional list
239
        of class names.
240

241
        Parameters
242
        ----------
243

244
        - class_names - `list` of `str`
245

246
          An optional list of names of the classes to generate the
247
          tree for.  Defaults to None where the class list is computed
248
          from the modules.
249

250
        """
251 4
        if class_names is None:
252 4
            class_names = []
253 4
            for m in self.modules:
254 4
                class_names.extend(dir(m))
255

256
        # Generate the nodes.
257 4
        for name in class_names:
258 4
            if ('.' in name):
259
                # With VTK 6.x and above there are strange names
260
                # in the vtk module which we ignore.
261 2
                continue
262 4
            klass = self.get_class(name)
263 4
            if klass and hasattr(klass, '__bases__'):
264 4
                self._generate_hierarchy(klass)
265

266
        # Compute the inheritance level and store the nodes in the tree.
267 4
        for node in self:
268 4
            d = node.get_level()
269 4
            while len(self.tree) <= d:
270 4
                self.tree.append([])
271 4
            self.tree[d].append(node)
272

273
        # Sort the nodes alphabetically.
274 4
        for nodes in self.tree:
275 4
            nodes.sort(key=lambda x:x.name)

Read our documentation on viewing source code .

Loading