1""" 2Add container iteration powers to wrapped C++ classes 3""" 4 5from pybindgen.typehandlers.base import ForwardWrapperBase 6from pybindgen.typehandlers import codesink 7from pybindgen.pytypeobject import PyTypeObject 8from pybindgen import utils 9 10 11class IterNextWrapper(ForwardWrapperBase): 12 ''' 13 tp_iternext wrapper 14 ''' 15 16 HAVE_RETURN_VALUE = True 17 18 def __init__(self, container): 19 """ 20 value_type -- a ReturnValue object handling the value type; 21 container -- the L{Container} 22 """ 23 super(IterNextWrapper, self).__init__( 24 None, [], "return NULL;", "return NULL;", no_c_retval=True) 25 assert isinstance(container, CppClassContainerTraits) 26 self.container = container 27 self.c_function_name = "_wrap_%s__tp_iternext" % (self.container.iter_pystruct) 28 self.iter_variable_name = None 29 self.reset_code_generation_state() 30 31 def reset_code_generation_state(self): 32 super(IterNextWrapper, self).reset_code_generation_state() 33 self.iter_variable_name = self.declarations.declare_variable( 34 "%s::%s" % (self.container.cppclass.full_name, self.container.iterator_type), 'iter') 35 36 def generate_call(self): 37 self.before_call.write_code("%s = *self->iterator;" % (self.iter_variable_name,)) 38 self.before_call.write_error_check( 39 "%s == self->container->obj->%s()" % (self.iter_variable_name, self.container.end_method), 40 "PyErr_SetNone(PyExc_StopIteration);") 41 self.before_call.write_code("++(*self->iterator);") 42 if self.container.key_type is None: 43 self.container.value_type.value = "(*%s)" % self.iter_variable_name 44 self.container.value_type.convert_c_to_python(self) 45 else: 46 self.container.value_type.value = "%s->second" % self.iter_variable_name 47 self.container.value_type.convert_c_to_python(self) 48 self.container.key_type.value = "%s->first" % self.iter_variable_name 49 self.container.key_type.convert_c_to_python(self) 50 51 def generate(self, code_sink): 52 """ 53 code_sink -- a CodeSink instance that will receive the generated code 54 """ 55 56 tmp_sink = codesink.MemoryCodeSink() 57 self.generate_body(tmp_sink) 58 code_sink.writeln("static PyObject* %s(%s *self)" % (self.c_function_name, 59 self.container.iter_pystruct)) 60 code_sink.writeln('{') 61 code_sink.indent() 62 tmp_sink.flush_to(code_sink) 63 code_sink.unindent() 64 code_sink.writeln('}') 65 66 67class CppClassContainerTraits(object): 68 def __init__(self, cppclass, value_type, begin_method='begin', end_method='end', iterator_type='iterator', is_mapping=False): 69 """ 70 :param cppclass: the L{CppClass} object that receives the container traits 71 72 :param value_type: a ReturnValue of the element type: note, 73 for mapping containers, value_type is a tuple with two 74 ReturnValue's: (key, element). 75 """ 76 self.cppclass = cppclass 77 self.begin_method = begin_method 78 self.end_method = end_method 79 self.iterator_type = iterator_type 80 81 self.iter_pytype = PyTypeObject() 82 self._iter_pystruct = None 83 84 if is_mapping: 85 (key_type, value_type) = value_type 86 self.key_type = utils.eval_retval(key_type, self) 87 self.value_type = utils.eval_retval(value_type, self) 88 else: 89 self.key_type = None 90 self.value_type = utils.eval_retval(value_type, self) 91 92 def get_iter_pystruct(self): 93 return "%s_Iter" % self.cppclass.pystruct 94 iter_pystruct = property(get_iter_pystruct) 95 96 def get_iter_pytypestruct(self): 97 return "%s_IterType" % self.cppclass.pystruct 98 iter_pytypestruct = property(get_iter_pytypestruct) 99 100 101 def generate_forward_declarations(self, code_sink, dummy_module): 102 """ 103 Generates forward declarations for the instance and type 104 structures. 105 """ 106 107 # container iterator pystruct 108 code_sink.writeln(''' 109typedef struct { 110 PyObject_HEAD 111 %s *container; 112 %s::%s *iterator; 113} %s; 114 ''' % (self.cppclass.pystruct, self.cppclass.full_name, self.iterator_type, self.iter_pystruct)) 115 116 code_sink.writeln() 117 code_sink.writeln('extern PyTypeObject %s;' % (self.iter_pytypestruct,)) 118 code_sink.writeln() 119 120 def get_iter_python_name(self): 121 return "%sIter" % self.cppclass.get_python_name() 122 123 def get_iter_python_full_name(self, module): 124 if self.cppclass.outer_class is None: 125 mod_path = module.get_module_path() 126 mod_path.append(self.get_iter_python_name()) 127 return '.'.join(mod_path) 128 else: 129 return '%s.%s' % (self.cppclass.outer_class.pytype.slots['tp_name'], 130 self.get_iter_python_name()) 131 132 133 def generate(self, code_sink, module, docstring=None): 134 """Generates the class to a code sink""" 135 136 ## --- register the iter type in the module --- 137 module.after_init.write_code("/* Register the '%s' class iterator*/" % self.cppclass.full_name) 138 module.after_init.write_error_check('PyType_Ready(&%s)' % (self.iter_pytypestruct,)) 139 140 if self.cppclass.outer_class is None: 141 module.after_init.write_code( 142 'PyModule_AddObject(m, (char *) \"%s\", (PyObject *) &%s);' % ( 143 self.get_iter_python_name(), self.iter_pytypestruct)) 144 else: 145 module.after_init.write_code( 146 'PyDict_SetItemString((PyObject*) %s.tp_dict, (char *) \"%s\", (PyObject *) &%s);' % ( 147 self.cppclass.outer_class.pytypestruct, self.cppclass.get_iter_python_name(), self.iter_pytypestruct)) 148 149 self._generate_gc_methods(code_sink) 150 self._generate_destructor(code_sink) 151 self._generate_iter_methods(code_sink) 152 self._generate_type_structure(code_sink, module, docstring) 153 154 def _generate_type_structure(self, code_sink, module, docstring): 155 """generate the type structure""" 156 self.iter_pytype.slots.setdefault("tp_basicsize", "sizeof(%s)" % (self.iter_pystruct,)) 157 self.iter_pytype.slots.setdefault("tp_flags", ("Py_TPFLAGS_DEFAULT|Py_TPFLAGS_HAVE_GC")) 158 self.iter_pytype.slots.setdefault("typestruct", self.iter_pytypestruct) 159 self.iter_pytype.slots.setdefault("tp_name", self.get_iter_python_full_name(module)) 160 if docstring: 161 self.iter_pytype.slots.setdefault("tp_doc", '"%s"' % docstring) 162 self.iter_pytype.generate(code_sink) 163 164 def _get_iter_delete_code(self): 165 delete_code = ("delete self->iterator;\n" 166 " self->iterator = NULL;\n") 167 return delete_code 168 169 def _get_container_delete_code(self): 170 delete_code = ("delete self->obj;\n" 171 " self->obj = NULL;\n") 172 return delete_code 173 174 def _generate_gc_methods(self, code_sink): 175 """Generate tp_clear and tp_traverse""" 176 177 ## --- iterator tp_clear --- 178 tp_clear_function_name = "%s__tp_clear" % (self.iter_pystruct,) 179 self.iter_pytype.slots.setdefault("tp_clear", tp_clear_function_name ) 180 181 code_sink.writeln(r''' 182static void 183%s(%s *self) 184{ 185 Py_CLEAR(self->container); 186 %s 187} 188''' % (tp_clear_function_name, self.iter_pystruct, self._get_iter_delete_code())) 189 190 ## --- iterator tp_traverse --- 191 tp_traverse_function_name = "%s__tp_traverse" % (self.iter_pystruct,) 192 self.iter_pytype.slots.setdefault("tp_traverse", tp_traverse_function_name ) 193 194 code_sink.writeln(r''' 195static int 196%s(%s *self, visitproc visit, void *arg) 197{ 198 Py_VISIT((PyObject *) self->container); 199 return 0; 200} 201''' % (tp_traverse_function_name, self.iter_pystruct)) 202 203 204 def _generate_destructor(self, code_sink): 205 """Generate a tp_dealloc function and register it in the type""" 206 207 # -- iterator -- 208 iter_tp_dealloc_function_name = "_wrap_%s__tp_dealloc" % (self.iter_pystruct,) 209 code_sink.writeln(r''' 210static void 211%s(%s *self) 212{ 213 Py_CLEAR(self->container); 214 %s 215 Py_TYPE(self)->tp_free((PyObject*)self); 216} 217''' % (iter_tp_dealloc_function_name, self.iter_pystruct, self._get_iter_delete_code())) 218 219 self.iter_pytype.slots.setdefault("tp_dealloc", iter_tp_dealloc_function_name ) 220 221 222 def _generate_iter_methods(self, code_sink): 223 224 container_tp_iter_function_name = "_wrap_%s__tp_iter" % (self.cppclass.pystruct,) 225 iterator_tp_iter_function_name = "_wrap_%s__tp_iter" % (self.iter_pystruct,) 226 subst_vars = { 227 'CONTAINER_ITER_FUNC': container_tp_iter_function_name, 228 'ITERATOR_ITER_FUNC': iterator_tp_iter_function_name, 229 'PYSTRUCT': self.cppclass.pystruct, 230 'ITER_PYSTRUCT': self.iter_pystruct, 231 'ITER_PYTYPESTRUCT': self.iter_pytypestruct, 232 'CTYPE': self.cppclass.full_name, 233 'BEGIN_METHOD': self.begin_method, 234 'ITERATOR_TYPE': self.iterator_type, 235 } 236 # -- container -- 237 code_sink.writeln(r''' 238static PyObject* 239%(CONTAINER_ITER_FUNC)s(%(PYSTRUCT)s *self) 240{ 241 %(ITER_PYSTRUCT)s *iter = PyObject_GC_New(%(ITER_PYSTRUCT)s, &%(ITER_PYTYPESTRUCT)s); 242 Py_INCREF(self); 243 iter->container = self; 244 iter->iterator = new %(CTYPE)s::%(ITERATOR_TYPE)s(self->obj->%(BEGIN_METHOD)s()); 245 return (PyObject*) iter; 246} 247''' % subst_vars) 248 249 self.cppclass.pytype.slots.setdefault("tp_iter", container_tp_iter_function_name) 250 251 252 # -- iterator -- 253 container_tp_iter_function_name = "_wrap_%s__tp_iter" % (self.cppclass.pystruct,) 254 code_sink.writeln(r''' 255static PyObject* 256%(ITERATOR_ITER_FUNC)s(%(ITER_PYSTRUCT)s *self) 257{ 258 Py_INCREF(self); 259 return (PyObject*) self; 260} 261''' % subst_vars) 262 263 self.iter_pytype.slots.setdefault("tp_iter", iterator_tp_iter_function_name) 264 265 # -- iterator tp_iternext 266 iternext = IterNextWrapper(self) 267 iternext.generate(code_sink) 268 self.iter_pytype.slots.setdefault("tp_iternext", iternext.c_function_name) 269 270 271