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