third_party/cython/src/Cython/Compiler/Tests/TestVisitor.py

   1 from Cython.Compiler.ModuleNode import ModuleNode
   2 from Cython.Compiler.Symtab import ModuleScope
   3 from Cython.TestUtils import TransformTest
   4 from Cython.Compiler.Visitor import MethodDispatcherTransform
   5 from Cython.Compiler.ParseTreeTransforms import (
   6     NormalizeTree, AnalyseDeclarationsTransform,
   7     AnalyseExpressionsTransform, InterpretCompilerDirectives)
   8
   9
  10 class TestMethodDispatcherTransform(TransformTest):
  11     _tree = None
  12
  13     def _build_tree(self):
  14         if self._tree is None:
  15             context = None
  16
  17             def fake_module(node):
  18                 scope = ModuleScope('test', None, None)
  19                 return ModuleNode(node.pos, doc=None, body=node,
  20                                   scope=scope, full_module_name='test',
  21                                   directive_comments={})
  22             pipeline = [
  23                 fake_module,
  24                 NormalizeTree(context),
  25                 InterpretCompilerDirectives(context, {}),
  26                 AnalyseDeclarationsTransform(context),
  27                 AnalyseExpressionsTransform(context),
  28             ]
  29             self._tree = self.run_pipeline(pipeline, u"""
  30                 cdef bytes s = b'asdfg'
  31                 cdef dict d = {1:2}
  32                 x = s * 3
  33                 d.get('test')
  34             """)
  35         return self._tree
  36
  37     def test_builtin_method(self):
  38         calls = [0]
  39         class Test(MethodDispatcherTransform):
  40             def _handle_simple_method_dict_get(self, node, func, args, unbound):
  41                 calls[0] += 1
  42                 return node
  43
  44         tree = self._build_tree()
  45         Test(None)(tree)
  46         self.assertEqual(1, calls[0])
  47
  48     def test_binop_method(self):
  49         calls = {'bytes': 0, 'object': 0}
  50         class Test(MethodDispatcherTransform):
  51             def _handle_simple_method_bytes___mul__(self, node, func, args, unbound):
  52                 calls['bytes'] += 1
  53                 return node
  54             def _handle_simple_method_object___mul__(self, node, func, args, unbound):
  55                 calls['object'] += 1
  56                 return node
  57
  58         tree = self._build_tree()
  59         Test(None)(tree)
  60         self.assertEqual(1, calls['bytes'])
  61         self.assertEqual(0, calls['object'])