1b2732e9dSimarom"""Python version-independent methods for C/Python buffers.
2b2732e9dSimarom
3b2732e9dSimaromThis file was copied and adapted from mpi4py.
4b2732e9dSimarom
5b2732e9dSimaromAuthors
6b2732e9dSimarom-------
7b2732e9dSimarom* MinRK
8b2732e9dSimarom"""
9b2732e9dSimarom
10b2732e9dSimarom#-----------------------------------------------------------------------------
11b2732e9dSimarom#  Copyright (c) 2010 Lisandro Dalcin
12b2732e9dSimarom#  All rights reserved.
13b2732e9dSimarom#  Used under BSD License: http://www.opensource.org/licenses/bsd-license.php
14b2732e9dSimarom#
15b2732e9dSimarom#  Retrieval:
16b2732e9dSimarom#  Jul 23, 2010 18:00 PST (r539)
17b2732e9dSimarom#  http://code.google.com/p/mpi4py/source/browse/trunk/src/MPI/asbuffer.pxi
18b2732e9dSimarom#
19b2732e9dSimarom#  Modifications from original:
20b2732e9dSimarom#  Copyright (c) 2010-2012 Brian Granger, Min Ragan-Kelley
21b2732e9dSimarom#
22b2732e9dSimarom#  Distributed under the terms of the New BSD License.  The full license is in
23b2732e9dSimarom#  the file COPYING.BSD, distributed as part of this software.
24b2732e9dSimarom#-----------------------------------------------------------------------------
25b2732e9dSimarom
26b2732e9dSimarom
27b2732e9dSimarom#-----------------------------------------------------------------------------
28b2732e9dSimarom# Python includes.
29b2732e9dSimarom#-----------------------------------------------------------------------------
30b2732e9dSimarom
31b2732e9dSimarom# get version-independent aliases:
32b2732e9dSimaromcdef extern from "pyversion_compat.h":
33b2732e9dSimarom    pass
34b2732e9dSimarom
35b2732e9dSimarom# Python 3 buffer interface (PEP 3118)
36b2732e9dSimaromcdef extern from "Python.h":
37b2732e9dSimarom    int PY_MAJOR_VERSION
38b2732e9dSimarom    int PY_MINOR_VERSION
39b2732e9dSimarom    ctypedef int Py_ssize_t
40b2732e9dSimarom    ctypedef struct PyMemoryViewObject:
41b2732e9dSimarom        pass
42b2732e9dSimarom    ctypedef struct Py_buffer:
43b2732e9dSimarom        void *buf
44b2732e9dSimarom        Py_ssize_t len
45b2732e9dSimarom        int readonly
46b2732e9dSimarom        char *format
47b2732e9dSimarom        int ndim
48b2732e9dSimarom        Py_ssize_t *shape
49b2732e9dSimarom        Py_ssize_t *strides
50b2732e9dSimarom        Py_ssize_t *suboffsets
51b2732e9dSimarom        Py_ssize_t itemsize
52b2732e9dSimarom        void *internal
53b2732e9dSimarom    cdef enum:
54b2732e9dSimarom        PyBUF_SIMPLE
55b2732e9dSimarom        PyBUF_WRITABLE
56b2732e9dSimarom        PyBUF_FORMAT
57b2732e9dSimarom        PyBUF_ANY_CONTIGUOUS
58b2732e9dSimarom    int  PyObject_CheckBuffer(object)
59b2732e9dSimarom    int  PyObject_GetBuffer(object, Py_buffer *, int) except -1
60b2732e9dSimarom    void PyBuffer_Release(Py_buffer *)
61b2732e9dSimarom    
62b2732e9dSimarom    int PyBuffer_FillInfo(Py_buffer *view, object obj, void *buf,
63b2732e9dSimarom                Py_ssize_t len, int readonly, int infoflags) except -1
64b2732e9dSimarom    object PyMemoryView_FromBuffer(Py_buffer *info)
65b2732e9dSimarom    
66b2732e9dSimarom    object PyMemoryView_FromObject(object)
67b2732e9dSimarom
68b2732e9dSimarom# Python 2 buffer interface (legacy)
69b2732e9dSimaromcdef extern from "Python.h":
70b2732e9dSimarom    ctypedef void const_void "const void"
71b2732e9dSimarom    Py_ssize_t Py_END_OF_BUFFER
72b2732e9dSimarom    int PyObject_CheckReadBuffer(object)
73b2732e9dSimarom    int PyObject_AsReadBuffer (object, const_void **, Py_ssize_t *) except -1
74b2732e9dSimarom    int PyObject_AsWriteBuffer(object, void **, Py_ssize_t *) except -1
75b2732e9dSimarom    
76b2732e9dSimarom    object PyBuffer_FromMemory(void *ptr, Py_ssize_t s)
77b2732e9dSimarom    object PyBuffer_FromReadWriteMemory(void *ptr, Py_ssize_t s)
78b2732e9dSimarom
79b2732e9dSimarom    object PyBuffer_FromObject(object, Py_ssize_t offset, Py_ssize_t size)
80b2732e9dSimarom    object PyBuffer_FromReadWriteObject(object, Py_ssize_t offset, Py_ssize_t size)
81b2732e9dSimarom
82b2732e9dSimarom
83b2732e9dSimarom#-----------------------------------------------------------------------------
84b2732e9dSimarom# asbuffer: C buffer from python object
85b2732e9dSimarom#-----------------------------------------------------------------------------
86b2732e9dSimarom
87b2732e9dSimarom
88b2732e9dSimaromcdef inline int memoryview_available():
89b2732e9dSimarom    return PY_MAJOR_VERSION >= 3 or (PY_MAJOR_VERSION >=2 and PY_MINOR_VERSION >= 7)
90b2732e9dSimarom
91b2732e9dSimaromcdef inline int oldstyle_available():
92b2732e9dSimarom    return PY_MAJOR_VERSION < 3
93b2732e9dSimarom
94b2732e9dSimarom
95b2732e9dSimaromcdef inline int check_buffer(object ob):
96b2732e9dSimarom    """Version independent check for whether an object is a buffer.
97b2732e9dSimarom    
98b2732e9dSimarom    Parameters
99b2732e9dSimarom    ----------
100b2732e9dSimarom    object : object
101b2732e9dSimarom        Any Python object
102b2732e9dSimarom
103b2732e9dSimarom    Returns
104b2732e9dSimarom    -------
105b2732e9dSimarom    int : 0 if no buffer interface, 3 if newstyle buffer interface, 2 if oldstyle.
106b2732e9dSimarom    """
107b2732e9dSimarom    if PyObject_CheckBuffer(ob):
108b2732e9dSimarom        return 3
109b2732e9dSimarom    if oldstyle_available():
110b2732e9dSimarom        return PyObject_CheckReadBuffer(ob) and 2
111b2732e9dSimarom    return 0
112b2732e9dSimarom
113b2732e9dSimarom
114b2732e9dSimaromcdef inline object asbuffer(object ob, int writable, int format,
115b2732e9dSimarom                            void **base, Py_ssize_t *size,
116b2732e9dSimarom                            Py_ssize_t *itemsize):
117b2732e9dSimarom    """Turn an object into a C buffer in a Python version-independent way.
118b2732e9dSimarom    
119b2732e9dSimarom    Parameters
120b2732e9dSimarom    ----------
121b2732e9dSimarom    ob : object
122b2732e9dSimarom        The object to be turned into a buffer.
123b2732e9dSimarom        Must provide a Python Buffer interface
124b2732e9dSimarom    writable : int
125b2732e9dSimarom        Whether the resulting buffer should be allowed to write
126b2732e9dSimarom        to the object.
127b2732e9dSimarom    format : int
128b2732e9dSimarom        The format of the buffer.  See Python buffer docs.
129b2732e9dSimarom    base : void **
130b2732e9dSimarom        The pointer that will be used to store the resulting C buffer.
131b2732e9dSimarom    size : Py_ssize_t *
132b2732e9dSimarom        The size of the buffer(s).
133b2732e9dSimarom    itemsize : Py_ssize_t *
134b2732e9dSimarom        The size of an item, if the buffer is non-contiguous.
135b2732e9dSimarom    
136b2732e9dSimarom    Returns
137b2732e9dSimarom    -------
138b2732e9dSimarom    An object describing the buffer format. Generally a str, such as 'B'.
139b2732e9dSimarom    """
140b2732e9dSimarom
141b2732e9dSimarom    cdef void *bptr = NULL
142b2732e9dSimarom    cdef Py_ssize_t blen = 0, bitemlen = 0
143b2732e9dSimarom    cdef Py_buffer view
144b2732e9dSimarom    cdef int flags = PyBUF_SIMPLE
145b2732e9dSimarom    cdef int mode = 0
146b2732e9dSimarom    
147b2732e9dSimarom    bfmt = None
148b2732e9dSimarom
149b2732e9dSimarom    mode = check_buffer(ob)
150b2732e9dSimarom    if mode == 0:
151b2732e9dSimarom        raise TypeError("%r does not provide a buffer interface."%ob)
152b2732e9dSimarom
153b2732e9dSimarom    if mode == 3:
154b2732e9dSimarom        flags = PyBUF_ANY_CONTIGUOUS
155b2732e9dSimarom        if writable:
156b2732e9dSimarom            flags |= PyBUF_WRITABLE
157b2732e9dSimarom        if format:
158b2732e9dSimarom            flags |= PyBUF_FORMAT
159b2732e9dSimarom        PyObject_GetBuffer(ob, &view, flags)
160b2732e9dSimarom        bptr = view.buf
161b2732e9dSimarom        blen = view.len
162b2732e9dSimarom        if format:
163b2732e9dSimarom            if view.format != NULL:
164b2732e9dSimarom                bfmt = view.format
165b2732e9dSimarom                bitemlen = view.itemsize
166b2732e9dSimarom        PyBuffer_Release(&view)
167b2732e9dSimarom    else: # oldstyle
168b2732e9dSimarom        if writable:
169b2732e9dSimarom            PyObject_AsWriteBuffer(ob, &bptr, &blen)
170b2732e9dSimarom        else:
171b2732e9dSimarom            PyObject_AsReadBuffer(ob, <const_void **>&bptr, &blen)
172b2732e9dSimarom        if format:
173b2732e9dSimarom            try: # numpy.ndarray
174b2732e9dSimarom                dtype = ob.dtype
175b2732e9dSimarom                bfmt = dtype.char
176b2732e9dSimarom                bitemlen = dtype.itemsize
177b2732e9dSimarom            except AttributeError:
178b2732e9dSimarom                try: # array.array
179b2732e9dSimarom                    bfmt = ob.typecode
180b2732e9dSimarom                    bitemlen = ob.itemsize
181b2732e9dSimarom                except AttributeError:
182b2732e9dSimarom                    if isinstance(ob, bytes):
183b2732e9dSimarom                        bfmt = b"B"
184b2732e9dSimarom                        bitemlen = 1
185b2732e9dSimarom                    else:
186b2732e9dSimarom                        # nothing found
187b2732e9dSimarom                        bfmt = None
188b2732e9dSimarom                        bitemlen = 0
189b2732e9dSimarom    if base: base[0] = <void *>bptr
190b2732e9dSimarom    if size: size[0] = <Py_ssize_t>blen
191b2732e9dSimarom    if itemsize: itemsize[0] = <Py_ssize_t>bitemlen
192b2732e9dSimarom    
193b2732e9dSimarom    if PY_MAJOR_VERSION >= 3 and bfmt is not None:
194b2732e9dSimarom        return bfmt.decode('ascii')
195b2732e9dSimarom    return bfmt
196b2732e9dSimarom
197b2732e9dSimarom
198b2732e9dSimaromcdef inline object asbuffer_r(object ob, void **base, Py_ssize_t *size):
199b2732e9dSimarom    """Wrapper for standard calls to asbuffer with a readonly buffer."""
200b2732e9dSimarom    asbuffer(ob, 0, 0, base, size, NULL)
201b2732e9dSimarom    return ob
202b2732e9dSimarom
203b2732e9dSimarom
204b2732e9dSimaromcdef inline object asbuffer_w(object ob, void **base, Py_ssize_t *size):
205b2732e9dSimarom    """Wrapper for standard calls to asbuffer with a writable buffer."""
206b2732e9dSimarom    asbuffer(ob, 1, 0, base, size, NULL)
207b2732e9dSimarom    return ob
208b2732e9dSimarom
209b2732e9dSimarom#------------------------------------------------------------------------------
210b2732e9dSimarom# frombuffer: python buffer/view from C buffer
211b2732e9dSimarom#------------------------------------------------------------------------------
212b2732e9dSimarom
213b2732e9dSimarom
214b2732e9dSimaromcdef inline object frombuffer_3(void *ptr, Py_ssize_t s, int readonly):
215b2732e9dSimarom    """Python 3 version of frombuffer.
216b2732e9dSimarom
217b2732e9dSimarom    This is the Python 3 model, but will work on Python >= 2.6. Currently,
218b2732e9dSimarom    we use it only on >= 3.0.
219b2732e9dSimarom    """
220b2732e9dSimarom    cdef Py_buffer pybuf
221b2732e9dSimarom    cdef Py_ssize_t *shape = [s]
222b2732e9dSimarom    cdef str astr=""
223b2732e9dSimarom    PyBuffer_FillInfo(&pybuf, astr, ptr, s, readonly, PyBUF_SIMPLE)
224b2732e9dSimarom    pybuf.format = "B"
225b2732e9dSimarom    pybuf.shape = shape
226b2732e9dSimarom    return PyMemoryView_FromBuffer(&pybuf)
227b2732e9dSimarom
228b2732e9dSimarom
229b2732e9dSimaromcdef inline object frombuffer_2(void *ptr, Py_ssize_t s, int readonly):
230b2732e9dSimarom    """Python 2 version of frombuffer. 
231b2732e9dSimarom
232b2732e9dSimarom    This must be used for Python <= 2.6, but we use it for all Python < 3.
233b2732e9dSimarom    """
234b2732e9dSimarom    
235b2732e9dSimarom    if oldstyle_available():
236b2732e9dSimarom        if readonly:
237b2732e9dSimarom            return PyBuffer_FromMemory(ptr, s)
238b2732e9dSimarom        else:
239b2732e9dSimarom            return PyBuffer_FromReadWriteMemory(ptr, s)
240b2732e9dSimarom    else:
241b2732e9dSimarom        raise NotImplementedError("Old style buffers not available.")
242b2732e9dSimarom
243b2732e9dSimarom
244b2732e9dSimaromcdef inline object frombuffer(void *ptr, Py_ssize_t s, int readonly):
245b2732e9dSimarom    """Create a Python Buffer/View of a C array. 
246b2732e9dSimarom    
247b2732e9dSimarom    Parameters
248b2732e9dSimarom    ----------
249b2732e9dSimarom    ptr : void *
250b2732e9dSimarom        Pointer to the array to be copied.
251b2732e9dSimarom    s : size_t
252b2732e9dSimarom        Length of the buffer.
253b2732e9dSimarom    readonly : int
254b2732e9dSimarom        whether the resulting object should be allowed to write to the buffer.
255b2732e9dSimarom    
256b2732e9dSimarom    Returns
257b2732e9dSimarom    -------
258b2732e9dSimarom    Python Buffer/View of the C buffer.
259b2732e9dSimarom    """
260b2732e9dSimarom    # oldstyle first priority for now
261b2732e9dSimarom    if oldstyle_available():
262b2732e9dSimarom        return frombuffer_2(ptr, s, readonly)
263b2732e9dSimarom    else:
264b2732e9dSimarom        return frombuffer_3(ptr, s, readonly)
265b2732e9dSimarom
266b2732e9dSimarom
267b2732e9dSimaromcdef inline object frombuffer_r(void *ptr, Py_ssize_t s):
268b2732e9dSimarom    """Wrapper for readonly view frombuffer."""
269b2732e9dSimarom    return frombuffer(ptr, s, 1)
270b2732e9dSimarom
271b2732e9dSimarom
272b2732e9dSimaromcdef inline object frombuffer_w(void *ptr, Py_ssize_t s):
273b2732e9dSimarom    """Wrapper for writable view frombuffer."""
274b2732e9dSimarom    return frombuffer(ptr, s, 0)
275b2732e9dSimarom
276b2732e9dSimarom#------------------------------------------------------------------------------
277b2732e9dSimarom# viewfromobject: python buffer/view from python object, refcounts intact
278b2732e9dSimarom# frombuffer(asbuffer(obj)) would lose track of refs
279b2732e9dSimarom#------------------------------------------------------------------------------
280b2732e9dSimarom
281b2732e9dSimaromcdef inline object viewfromobject(object obj, int readonly):
282b2732e9dSimarom    """Construct a Python Buffer/View object from another Python object.
283b2732e9dSimarom
284b2732e9dSimarom    This work in a Python version independent manner.
285b2732e9dSimarom    
286b2732e9dSimarom    Parameters
287b2732e9dSimarom    ----------
288b2732e9dSimarom    obj : object
289b2732e9dSimarom        The input object to be cast as a buffer
290b2732e9dSimarom    readonly : int
291b2732e9dSimarom        Whether the result should be prevented from overwriting the original.
292b2732e9dSimarom    
293b2732e9dSimarom    Returns
294b2732e9dSimarom    -------
295b2732e9dSimarom    Buffer/View of the original object.
296b2732e9dSimarom    """
297b2732e9dSimarom    if not memoryview_available():
298b2732e9dSimarom        if readonly:
299b2732e9dSimarom            return PyBuffer_FromObject(obj, 0, Py_END_OF_BUFFER)
300b2732e9dSimarom        else:
301b2732e9dSimarom            return PyBuffer_FromReadWriteObject(obj, 0, Py_END_OF_BUFFER)
302b2732e9dSimarom    else:
303b2732e9dSimarom        return PyMemoryView_FromObject(obj)
304b2732e9dSimarom
305b2732e9dSimarom
306b2732e9dSimaromcdef inline object viewfromobject_r(object obj):
307b2732e9dSimarom    """Wrapper for readonly viewfromobject."""
308b2732e9dSimarom    return viewfromobject(obj, 1)
309b2732e9dSimarom
310b2732e9dSimarom
311b2732e9dSimaromcdef inline object viewfromobject_w(object obj):
312b2732e9dSimarom    """Wrapper for writable viewfromobject."""
313b2732e9dSimarom    return viewfromobject(obj, 0)
314