Visitor.py 29 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609610611612613614615616617618619620621622623624625626627628629630631632633634635636637638639640641642643644645646647648649650651652653654655656657658659660661662663664665666667668669670671672673674675676677678679680681682683684685686687688689690691692693694695696697698699700701702703704705706707708709710711712713714715716717718719720721722723724725726727728729730731732733734735736737738739740741742743744745746747748749750751752753754755756757758759760761762763764765766767768769770771772773774775776777778779780781782783784785786787788789790791792793794795796797798799800801802803804805806807808809810811812813814815816817818819820821822823824825826827828829830831832833834835836837838839840
  1. # cython: infer_types=True
  2. # cython: language_level=3
  3. # cython: auto_pickle=False
  4. #
  5. # Tree visitor and transform framework
  6. #
  7. from __future__ import absolute_import, print_function
  8. import sys
  9. import inspect
  10. from . import TypeSlots
  11. from . import Builtin
  12. from . import Nodes
  13. from . import ExprNodes
  14. from . import Errors
  15. from . import DebugFlags
  16. from . import Future
  17. import cython
  18. cython.declare(_PRINTABLE=tuple)
  19. if sys.version_info[0] >= 3:
  20. _PRINTABLE = (bytes, str, int, float)
  21. else:
  22. _PRINTABLE = (str, unicode, long, int, float)
  23. class TreeVisitor(object):
  24. """
  25. Base class for writing visitors for a Cython tree, contains utilities for
  26. recursing such trees using visitors. Each node is
  27. expected to have a child_attrs iterable containing the names of attributes
  28. containing child nodes or lists of child nodes. Lists are not considered
  29. part of the tree structure (i.e. contained nodes are considered direct
  30. children of the parent node).
  31. visit_children visits each of the children of a given node (see the visit_children
  32. documentation). When recursing the tree using visit_children, an attribute
  33. access_path is maintained which gives information about the current location
  34. in the tree as a stack of tuples: (parent_node, attrname, index), representing
  35. the node, attribute and optional list index that was taken in each step in the path to
  36. the current node.
  37. Example:
  38. >>> class SampleNode(object):
  39. ... child_attrs = ["head", "body"]
  40. ... def __init__(self, value, head=None, body=None):
  41. ... self.value = value
  42. ... self.head = head
  43. ... self.body = body
  44. ... def __repr__(self): return "SampleNode(%s)" % self.value
  45. ...
  46. >>> tree = SampleNode(0, SampleNode(1), [SampleNode(2), SampleNode(3)])
  47. >>> class MyVisitor(TreeVisitor):
  48. ... def visit_SampleNode(self, node):
  49. ... print("in %s %s" % (node.value, self.access_path))
  50. ... self.visitchildren(node)
  51. ... print("out %s" % node.value)
  52. ...
  53. >>> MyVisitor().visit(tree)
  54. in 0 []
  55. in 1 [(SampleNode(0), 'head', None)]
  56. out 1
  57. in 2 [(SampleNode(0), 'body', 0)]
  58. out 2
  59. in 3 [(SampleNode(0), 'body', 1)]
  60. out 3
  61. out 0
  62. """
  63. def __init__(self):
  64. super(TreeVisitor, self).__init__()
  65. self.dispatch_table = {}
  66. self.access_path = []
  67. def dump_node(self, node):
  68. ignored = list(node.child_attrs or []) + [
  69. u'child_attrs', u'pos', u'gil_message', u'cpp_message', u'subexprs']
  70. values = []
  71. pos = getattr(node, 'pos', None)
  72. if pos:
  73. source = pos[0]
  74. if source:
  75. import os.path
  76. source = os.path.basename(source.get_description())
  77. values.append(u'%s:%s:%s' % (source, pos[1], pos[2]))
  78. attribute_names = dir(node)
  79. for attr in attribute_names:
  80. if attr in ignored:
  81. continue
  82. if attr.startswith('_') or attr.endswith('_'):
  83. continue
  84. try:
  85. value = getattr(node, attr)
  86. except AttributeError:
  87. continue
  88. if value is None or value == 0:
  89. continue
  90. elif isinstance(value, list):
  91. value = u'[...]/%d' % len(value)
  92. elif not isinstance(value, _PRINTABLE):
  93. continue
  94. else:
  95. value = repr(value)
  96. values.append(u'%s = %s' % (attr, value))
  97. return u'%s(%s)' % (node.__class__.__name__, u',\n '.join(values))
  98. def _find_node_path(self, stacktrace):
  99. import os.path
  100. last_traceback = stacktrace
  101. nodes = []
  102. while hasattr(stacktrace, 'tb_frame'):
  103. frame = stacktrace.tb_frame
  104. node = frame.f_locals.get(u'self')
  105. if isinstance(node, Nodes.Node):
  106. code = frame.f_code
  107. method_name = code.co_name
  108. pos = (os.path.basename(code.co_filename),
  109. frame.f_lineno)
  110. nodes.append((node, method_name, pos))
  111. last_traceback = stacktrace
  112. stacktrace = stacktrace.tb_next
  113. return (last_traceback, nodes)
  114. def _raise_compiler_error(self, child, e):
  115. trace = ['']
  116. for parent, attribute, index in self.access_path:
  117. node = getattr(parent, attribute)
  118. if index is None:
  119. index = ''
  120. else:
  121. node = node[index]
  122. index = u'[%d]' % index
  123. trace.append(u'%s.%s%s = %s' % (
  124. parent.__class__.__name__, attribute, index,
  125. self.dump_node(node)))
  126. stacktrace, called_nodes = self._find_node_path(sys.exc_info()[2])
  127. last_node = child
  128. for node, method_name, pos in called_nodes:
  129. last_node = node
  130. trace.append(u"File '%s', line %d, in %s: %s" % (
  131. pos[0], pos[1], method_name, self.dump_node(node)))
  132. raise Errors.CompilerCrash(
  133. getattr(last_node, 'pos', None), self.__class__.__name__,
  134. u'\n'.join(trace), e, stacktrace)
  135. @cython.final
  136. def find_handler(self, obj):
  137. # to resolve, try entire hierarchy
  138. cls = type(obj)
  139. pattern = "visit_%s"
  140. mro = inspect.getmro(cls)
  141. for mro_cls in mro:
  142. handler_method = getattr(self, pattern % mro_cls.__name__, None)
  143. if handler_method is not None:
  144. return handler_method
  145. print(type(self), cls)
  146. if self.access_path:
  147. print(self.access_path)
  148. print(self.access_path[-1][0].pos)
  149. print(self.access_path[-1][0].__dict__)
  150. raise RuntimeError("Visitor %r does not accept object: %s" % (self, obj))
  151. def visit(self, obj):
  152. return self._visit(obj)
  153. @cython.final
  154. def _visit(self, obj):
  155. try:
  156. try:
  157. handler_method = self.dispatch_table[type(obj)]
  158. except KeyError:
  159. handler_method = self.find_handler(obj)
  160. self.dispatch_table[type(obj)] = handler_method
  161. return handler_method(obj)
  162. except Errors.CompileError:
  163. raise
  164. except Errors.AbortError:
  165. raise
  166. except Exception as e:
  167. if DebugFlags.debug_no_exception_intercept:
  168. raise
  169. self._raise_compiler_error(obj, e)
  170. @cython.final
  171. def _visitchild(self, child, parent, attrname, idx):
  172. self.access_path.append((parent, attrname, idx))
  173. result = self._visit(child)
  174. self.access_path.pop()
  175. return result
  176. def visitchildren(self, parent, attrs=None):
  177. return self._visitchildren(parent, attrs)
  178. @cython.final
  179. @cython.locals(idx=int)
  180. def _visitchildren(self, parent, attrs):
  181. """
  182. Visits the children of the given parent. If parent is None, returns
  183. immediately (returning None).
  184. The return value is a dictionary giving the results for each
  185. child (mapping the attribute name to either the return value
  186. or a list of return values (in the case of multiple children
  187. in an attribute)).
  188. """
  189. if parent is None: return None
  190. result = {}
  191. for attr in parent.child_attrs:
  192. if attrs is not None and attr not in attrs: continue
  193. child = getattr(parent, attr)
  194. if child is not None:
  195. if type(child) is list:
  196. childretval = [self._visitchild(x, parent, attr, idx) for idx, x in enumerate(child)]
  197. else:
  198. childretval = self._visitchild(child, parent, attr, None)
  199. assert not isinstance(childretval, list), 'Cannot insert list here: %s in %r' % (attr, parent)
  200. result[attr] = childretval
  201. return result
  202. class VisitorTransform(TreeVisitor):
  203. """
  204. A tree transform is a base class for visitors that wants to do stream
  205. processing of the structure (rather than attributes etc.) of a tree.
  206. It implements __call__ to simply visit the argument node.
  207. It requires the visitor methods to return the nodes which should take
  208. the place of the visited node in the result tree (which can be the same
  209. or one or more replacement). Specifically, if the return value from
  210. a visitor method is:
  211. - [] or None; the visited node will be removed (set to None if an attribute and
  212. removed if in a list)
  213. - A single node; the visited node will be replaced by the returned node.
  214. - A list of nodes; the visited nodes will be replaced by all the nodes in the
  215. list. This will only work if the node was already a member of a list; if it
  216. was not, an exception will be raised. (Typically you want to ensure that you
  217. are within a StatListNode or similar before doing this.)
  218. """
  219. def visitchildren(self, parent, attrs=None, exclude=None):
  220. # generic def entry point for calls from Python subclasses
  221. if exclude is not None:
  222. attrs = self._select_attrs(parent.child_attrs if attrs is None else attrs, exclude)
  223. return self._process_children(parent, attrs)
  224. @cython.final
  225. def _select_attrs(self, attrs, exclude):
  226. return [name for name in attrs if name not in exclude]
  227. @cython.final
  228. def _process_children(self, parent, attrs=None):
  229. # fast cdef entry point for calls from Cython subclasses
  230. result = self._visitchildren(parent, attrs)
  231. for attr, newnode in result.items():
  232. if type(newnode) is list:
  233. newnode = self._flatten_list(newnode)
  234. setattr(parent, attr, newnode)
  235. return result
  236. @cython.final
  237. def _flatten_list(self, orig_list):
  238. # Flatten the list one level and remove any None
  239. newlist = []
  240. for x in orig_list:
  241. if x is not None:
  242. if type(x) is list:
  243. newlist.extend(x)
  244. else:
  245. newlist.append(x)
  246. return newlist
  247. def recurse_to_children(self, node):
  248. self._process_children(node)
  249. return node
  250. def __call__(self, root):
  251. return self._visit(root)
  252. class CythonTransform(VisitorTransform):
  253. """
  254. Certain common conventions and utilities for Cython transforms.
  255. - Sets up the context of the pipeline in self.context
  256. - Tracks directives in effect in self.current_directives
  257. """
  258. def __init__(self, context):
  259. super(CythonTransform, self).__init__()
  260. self.context = context
  261. def __call__(self, node):
  262. from . import ModuleNode
  263. if isinstance(node, ModuleNode.ModuleNode):
  264. self.current_directives = node.directives
  265. return super(CythonTransform, self).__call__(node)
  266. def visit_CompilerDirectivesNode(self, node):
  267. old = self.current_directives
  268. self.current_directives = node.directives
  269. self._process_children(node)
  270. self.current_directives = old
  271. return node
  272. def visit_Node(self, node):
  273. self._process_children(node)
  274. return node
  275. class ScopeTrackingTransform(CythonTransform):
  276. # Keeps track of type of scopes
  277. #scope_type: can be either of 'module', 'function', 'cclass', 'pyclass', 'struct'
  278. #scope_node: the node that owns the current scope
  279. def visit_ModuleNode(self, node):
  280. self.scope_type = 'module'
  281. self.scope_node = node
  282. self._process_children(node)
  283. return node
  284. def visit_scope(self, node, scope_type):
  285. prev = self.scope_type, self.scope_node
  286. self.scope_type = scope_type
  287. self.scope_node = node
  288. self._process_children(node)
  289. self.scope_type, self.scope_node = prev
  290. return node
  291. def visit_CClassDefNode(self, node):
  292. return self.visit_scope(node, 'cclass')
  293. def visit_PyClassDefNode(self, node):
  294. return self.visit_scope(node, 'pyclass')
  295. def visit_FuncDefNode(self, node):
  296. return self.visit_scope(node, 'function')
  297. def visit_CStructOrUnionDefNode(self, node):
  298. return self.visit_scope(node, 'struct')
  299. class EnvTransform(CythonTransform):
  300. """
  301. This transformation keeps a stack of the environments.
  302. """
  303. def __call__(self, root):
  304. self.env_stack = []
  305. self.enter_scope(root, root.scope)
  306. return super(EnvTransform, self).__call__(root)
  307. def current_env(self):
  308. return self.env_stack[-1][1]
  309. def current_scope_node(self):
  310. return self.env_stack[-1][0]
  311. def global_scope(self):
  312. return self.current_env().global_scope()
  313. def enter_scope(self, node, scope):
  314. self.env_stack.append((node, scope))
  315. def exit_scope(self):
  316. self.env_stack.pop()
  317. def visit_FuncDefNode(self, node):
  318. self.enter_scope(node, node.local_scope)
  319. self._process_children(node)
  320. self.exit_scope()
  321. return node
  322. def visit_GeneratorBodyDefNode(self, node):
  323. self._process_children(node)
  324. return node
  325. def visit_ClassDefNode(self, node):
  326. self.enter_scope(node, node.scope)
  327. self._process_children(node)
  328. self.exit_scope()
  329. return node
  330. def visit_CStructOrUnionDefNode(self, node):
  331. self.enter_scope(node, node.scope)
  332. self._process_children(node)
  333. self.exit_scope()
  334. return node
  335. def visit_ScopedExprNode(self, node):
  336. if node.expr_scope:
  337. self.enter_scope(node, node.expr_scope)
  338. self._process_children(node)
  339. self.exit_scope()
  340. else:
  341. self._process_children(node)
  342. return node
  343. def visit_CArgDeclNode(self, node):
  344. # default arguments are evaluated in the outer scope
  345. if node.default:
  346. attrs = [attr for attr in node.child_attrs if attr != 'default']
  347. self._process_children(node, attrs)
  348. self.enter_scope(node, self.current_env().outer_scope)
  349. self.visitchildren(node, ('default',))
  350. self.exit_scope()
  351. else:
  352. self._process_children(node)
  353. return node
  354. class NodeRefCleanupMixin(object):
  355. """
  356. Clean up references to nodes that were replaced.
  357. NOTE: this implementation assumes that the replacement is
  358. done first, before hitting any further references during
  359. normal tree traversal. This needs to be arranged by calling
  360. "self.visitchildren()" at a proper place in the transform
  361. and by ordering the "child_attrs" of nodes appropriately.
  362. """
  363. def __init__(self, *args):
  364. super(NodeRefCleanupMixin, self).__init__(*args)
  365. self._replacements = {}
  366. def visit_CloneNode(self, node):
  367. arg = node.arg
  368. if arg not in self._replacements:
  369. self.visitchildren(arg)
  370. node.arg = self._replacements.get(arg, arg)
  371. return node
  372. def visit_ResultRefNode(self, node):
  373. expr = node.expression
  374. if expr is None or expr not in self._replacements:
  375. self.visitchildren(node)
  376. expr = node.expression
  377. if expr is not None:
  378. node.expression = self._replacements.get(expr, expr)
  379. return node
  380. def replace(self, node, replacement):
  381. self._replacements[node] = replacement
  382. return replacement
  383. find_special_method_for_binary_operator = {
  384. '<': '__lt__',
  385. '<=': '__le__',
  386. '==': '__eq__',
  387. '!=': '__ne__',
  388. '>=': '__ge__',
  389. '>': '__gt__',
  390. '+': '__add__',
  391. '&': '__and__',
  392. '/': '__div__',
  393. '//': '__floordiv__',
  394. '<<': '__lshift__',
  395. '%': '__mod__',
  396. '*': '__mul__',
  397. '|': '__or__',
  398. '**': '__pow__',
  399. '>>': '__rshift__',
  400. '-': '__sub__',
  401. '^': '__xor__',
  402. 'in': '__contains__',
  403. }.get
  404. find_special_method_for_unary_operator = {
  405. 'not': '__not__',
  406. '~': '__inv__',
  407. '-': '__neg__',
  408. '+': '__pos__',
  409. }.get
  410. class MethodDispatcherTransform(EnvTransform):
  411. """
  412. Base class for transformations that want to intercept on specific
  413. builtin functions or methods of builtin types, including special
  414. methods triggered by Python operators. Must run after declaration
  415. analysis when entries were assigned.
  416. Naming pattern for handler methods is as follows:
  417. * builtin functions: _handle_(general|simple|any)_function_NAME
  418. * builtin methods: _handle_(general|simple|any)_method_TYPENAME_METHODNAME
  419. """
  420. # only visit call nodes and Python operations
  421. def visit_GeneralCallNode(self, node):
  422. self._process_children(node)
  423. function = node.function
  424. if not function.type.is_pyobject:
  425. return node
  426. arg_tuple = node.positional_args
  427. if not isinstance(arg_tuple, ExprNodes.TupleNode):
  428. return node
  429. keyword_args = node.keyword_args
  430. if keyword_args and not isinstance(keyword_args, ExprNodes.DictNode):
  431. # can't handle **kwargs
  432. return node
  433. args = arg_tuple.args
  434. return self._dispatch_to_handler(node, function, args, keyword_args)
  435. def visit_SimpleCallNode(self, node):
  436. self._process_children(node)
  437. function = node.function
  438. if function.type.is_pyobject:
  439. arg_tuple = node.arg_tuple
  440. if not isinstance(arg_tuple, ExprNodes.TupleNode):
  441. return node
  442. args = arg_tuple.args
  443. else:
  444. args = node.args
  445. return self._dispatch_to_handler(node, function, args, None)
  446. def visit_PrimaryCmpNode(self, node):
  447. if node.cascade:
  448. # not currently handled below
  449. self._process_children(node)
  450. return node
  451. return self._visit_binop_node(node)
  452. def visit_BinopNode(self, node):
  453. return self._visit_binop_node(node)
  454. def _visit_binop_node(self, node):
  455. self._process_children(node)
  456. # FIXME: could special case 'not_in'
  457. special_method_name = find_special_method_for_binary_operator(node.operator)
  458. if special_method_name:
  459. operand1, operand2 = node.operand1, node.operand2
  460. if special_method_name == '__contains__':
  461. operand1, operand2 = operand2, operand1
  462. elif special_method_name == '__div__':
  463. if Future.division in self.current_env().global_scope().context.future_directives:
  464. special_method_name = '__truediv__'
  465. obj_type = operand1.type
  466. if obj_type.is_builtin_type:
  467. type_name = obj_type.name
  468. else:
  469. type_name = "object" # safety measure
  470. node = self._dispatch_to_method_handler(
  471. special_method_name, None, False, type_name,
  472. node, None, [operand1, operand2], None)
  473. return node
  474. def visit_UnopNode(self, node):
  475. self._process_children(node)
  476. special_method_name = find_special_method_for_unary_operator(node.operator)
  477. if special_method_name:
  478. operand = node.operand
  479. obj_type = operand.type
  480. if obj_type.is_builtin_type:
  481. type_name = obj_type.name
  482. else:
  483. type_name = "object" # safety measure
  484. node = self._dispatch_to_method_handler(
  485. special_method_name, None, False, type_name,
  486. node, None, [operand], None)
  487. return node
  488. ### dispatch to specific handlers
  489. def _find_handler(self, match_name, has_kwargs):
  490. call_type = has_kwargs and 'general' or 'simple'
  491. handler = getattr(self, '_handle_%s_%s' % (call_type, match_name), None)
  492. if handler is None:
  493. handler = getattr(self, '_handle_any_%s' % match_name, None)
  494. return handler
  495. def _delegate_to_assigned_value(self, node, function, arg_list, kwargs):
  496. assignment = function.cf_state[0]
  497. value = assignment.rhs
  498. if value.is_name:
  499. if not value.entry or len(value.entry.cf_assignments) > 1:
  500. # the variable might have been reassigned => play safe
  501. return node
  502. elif value.is_attribute and value.obj.is_name:
  503. if not value.obj.entry or len(value.obj.entry.cf_assignments) > 1:
  504. # the underlying variable might have been reassigned => play safe
  505. return node
  506. else:
  507. return node
  508. return self._dispatch_to_handler(
  509. node, value, arg_list, kwargs)
  510. def _dispatch_to_handler(self, node, function, arg_list, kwargs):
  511. if function.is_name:
  512. # we only consider functions that are either builtin
  513. # Python functions or builtins that were already replaced
  514. # into a C function call (defined in the builtin scope)
  515. if not function.entry:
  516. return node
  517. entry = function.entry
  518. is_builtin = (
  519. entry.is_builtin or
  520. entry is self.current_env().builtin_scope().lookup_here(function.name))
  521. if not is_builtin:
  522. if function.cf_state and function.cf_state.is_single:
  523. # we know the value of the variable
  524. # => see if it's usable instead
  525. return self._delegate_to_assigned_value(
  526. node, function, arg_list, kwargs)
  527. if arg_list and entry.is_cmethod and entry.scope and entry.scope.parent_type.is_builtin_type:
  528. if entry.scope.parent_type is arg_list[0].type:
  529. # Optimised (unbound) method of a builtin type => try to "de-optimise".
  530. return self._dispatch_to_method_handler(
  531. entry.name, self_arg=None, is_unbound_method=True,
  532. type_name=entry.scope.parent_type.name,
  533. node=node, function=function, arg_list=arg_list, kwargs=kwargs)
  534. return node
  535. function_handler = self._find_handler(
  536. "function_%s" % function.name, kwargs)
  537. if function_handler is None:
  538. return self._handle_function(node, function.name, function, arg_list, kwargs)
  539. if kwargs:
  540. return function_handler(node, function, arg_list, kwargs)
  541. else:
  542. return function_handler(node, function, arg_list)
  543. elif function.is_attribute:
  544. attr_name = function.attribute
  545. if function.type.is_pyobject:
  546. self_arg = function.obj
  547. elif node.self and function.entry:
  548. entry = function.entry.as_variable
  549. if not entry or not entry.is_builtin:
  550. return node
  551. # C implementation of a Python builtin method - see if we find further matches
  552. self_arg = node.self
  553. arg_list = arg_list[1:] # drop CloneNode of self argument
  554. else:
  555. return node
  556. obj_type = self_arg.type
  557. is_unbound_method = False
  558. if obj_type.is_builtin_type:
  559. if obj_type is Builtin.type_type and self_arg.is_name and arg_list and arg_list[0].type.is_pyobject:
  560. # calling an unbound method like 'list.append(L,x)'
  561. # (ignoring 'type.mro()' here ...)
  562. type_name = self_arg.name
  563. self_arg = None
  564. is_unbound_method = True
  565. else:
  566. type_name = obj_type.name
  567. else:
  568. type_name = "object" # safety measure
  569. return self._dispatch_to_method_handler(
  570. attr_name, self_arg, is_unbound_method, type_name,
  571. node, function, arg_list, kwargs)
  572. else:
  573. return node
  574. def _dispatch_to_method_handler(self, attr_name, self_arg,
  575. is_unbound_method, type_name,
  576. node, function, arg_list, kwargs):
  577. method_handler = self._find_handler(
  578. "method_%s_%s" % (type_name, attr_name), kwargs)
  579. if method_handler is None:
  580. if (attr_name in TypeSlots.method_name_to_slot
  581. or attr_name == '__new__'):
  582. method_handler = self._find_handler(
  583. "slot%s" % attr_name, kwargs)
  584. if method_handler is None:
  585. return self._handle_method(
  586. node, type_name, attr_name, function,
  587. arg_list, is_unbound_method, kwargs)
  588. if self_arg is not None:
  589. arg_list = [self_arg] + list(arg_list)
  590. if kwargs:
  591. result = method_handler(
  592. node, function, arg_list, is_unbound_method, kwargs)
  593. else:
  594. result = method_handler(
  595. node, function, arg_list, is_unbound_method)
  596. return result
  597. def _handle_function(self, node, function_name, function, arg_list, kwargs):
  598. """Fallback handler"""
  599. return node
  600. def _handle_method(self, node, type_name, attr_name, function,
  601. arg_list, is_unbound_method, kwargs):
  602. """Fallback handler"""
  603. return node
  604. class RecursiveNodeReplacer(VisitorTransform):
  605. """
  606. Recursively replace all occurrences of a node in a subtree by
  607. another node.
  608. """
  609. def __init__(self, orig_node, new_node):
  610. super(RecursiveNodeReplacer, self).__init__()
  611. self.orig_node, self.new_node = orig_node, new_node
  612. def visit_CloneNode(self, node):
  613. if node is self.orig_node:
  614. return self.new_node
  615. if node.arg is self.orig_node:
  616. node.arg = self.new_node
  617. return node
  618. def visit_Node(self, node):
  619. self._process_children(node)
  620. if node is self.orig_node:
  621. return self.new_node
  622. else:
  623. return node
  624. def recursively_replace_node(tree, old_node, new_node):
  625. replace_in = RecursiveNodeReplacer(old_node, new_node)
  626. replace_in(tree)
  627. class NodeFinder(TreeVisitor):
  628. """
  629. Find out if a node appears in a subtree.
  630. """
  631. def __init__(self, node):
  632. super(NodeFinder, self).__init__()
  633. self.node = node
  634. self.found = False
  635. def visit_Node(self, node):
  636. if self.found:
  637. pass # short-circuit
  638. elif node is self.node:
  639. self.found = True
  640. else:
  641. self._visitchildren(node, None)
  642. def tree_contains(tree, node):
  643. finder = NodeFinder(node)
  644. finder.visit(tree)
  645. return finder.found
  646. # Utils
  647. def replace_node(ptr, value):
  648. """Replaces a node. ptr is of the form used on the access path stack
  649. (parent, attrname, listidx|None)
  650. """
  651. parent, attrname, listidx = ptr
  652. if listidx is None:
  653. setattr(parent, attrname, value)
  654. else:
  655. getattr(parent, attrname)[listidx] = value
  656. class PrintTree(TreeVisitor):
  657. """Prints a representation of the tree to standard output.
  658. Subclass and override repr_of to provide more information
  659. about nodes. """
  660. def __init__(self, start=None, end=None):
  661. TreeVisitor.__init__(self)
  662. self._indent = ""
  663. if start is not None or end is not None:
  664. self._line_range = (start or 0, end or 2**30)
  665. else:
  666. self._line_range = None
  667. def indent(self):
  668. self._indent += " "
  669. def unindent(self):
  670. self._indent = self._indent[:-2]
  671. def __call__(self, tree, phase=None):
  672. print("Parse tree dump at phase '%s'" % phase)
  673. self.visit(tree)
  674. return tree
  675. # Don't do anything about process_list, the defaults gives
  676. # nice-looking name[idx] nodes which will visually appear
  677. # under the parent-node, not displaying the list itself in
  678. # the hierarchy.
  679. def visit_Node(self, node):
  680. self._print_node(node)
  681. self.indent()
  682. self.visitchildren(node)
  683. self.unindent()
  684. return node
  685. def visit_CloneNode(self, node):
  686. self._print_node(node)
  687. self.indent()
  688. line = node.pos[1]
  689. if self._line_range is None or self._line_range[0] <= line <= self._line_range[1]:
  690. print("%s- %s: %s" % (self._indent, 'arg', self.repr_of(node.arg)))
  691. self.indent()
  692. self.visitchildren(node.arg)
  693. self.unindent()
  694. self.unindent()
  695. return node
  696. def _print_node(self, node):
  697. line = node.pos[1]
  698. if self._line_range is None or self._line_range[0] <= line <= self._line_range[1]:
  699. if len(self.access_path) == 0:
  700. name = "(root)"
  701. else:
  702. parent, attr, idx = self.access_path[-1]
  703. if idx is not None:
  704. name = "%s[%d]" % (attr, idx)
  705. else:
  706. name = attr
  707. print("%s- %s: %s" % (self._indent, name, self.repr_of(node)))
  708. def repr_of(self, node):
  709. if node is None:
  710. return "(none)"
  711. else:
  712. result = node.__class__.__name__
  713. if isinstance(node, ExprNodes.NameNode):
  714. result += "(type=%s, name=\"%s\")" % (repr(node.type), node.name)
  715. elif isinstance(node, Nodes.DefNode):
  716. result += "(name=\"%s\")" % node.name
  717. elif isinstance(node, ExprNodes.ExprNode):
  718. t = node.type
  719. result += "(type=%s)" % repr(t)
  720. elif node.pos:
  721. pos = node.pos
  722. path = pos[0].get_description()
  723. if '/' in path:
  724. path = path.split('/')[-1]
  725. if '\\' in path:
  726. path = path.split('\\')[-1]
  727. result += "(pos=(%s:%s:%s))" % (path, pos[1], pos[2])
  728. return result
  729. if __name__ == "__main__":
  730. import doctest
  731. doctest.testmod()