├── runtime.txt ├── requirements.txt ├── NOTICE ├── pyproject.toml ├── test ├── __init__.py ├── default_allocator.py ├── null_allocator.py ├── page_aligned_allocator.py ├── debug_allocator.py └── aligned_allocator.py ├── .github └── workflows │ └── pypi.yml ├── README.md ├── setup.py ├── NumPy-Allocator.ipynb ├── LICENSE └── numpy_allocator.c /runtime.txt: -------------------------------------------------------------------------------- 1 | python-3.8 2 | -------------------------------------------------------------------------------- /requirements.txt: -------------------------------------------------------------------------------- 1 | hypothesis 2 | pytest 3 | -------------------------------------------------------------------------------- /NOTICE: -------------------------------------------------------------------------------- 1 | inaccel/numpy-allocator 2 | Copyright © 2018-2023 InAccel 3 | 4 | This product includes software developed at 5 | InAccel (https://inaccel.com). 6 | -------------------------------------------------------------------------------- /pyproject.toml: -------------------------------------------------------------------------------- 1 | [build-system] 2 | requires = [ 3 | "numpy>=1.22", 4 | "setuptools", 5 | "setuptools_scm", 6 | "wheel" 7 | ] 8 | 9 | [tool.cibuildwheel] 10 | build = [ 11 | "cp3*" 12 | ] 13 | -------------------------------------------------------------------------------- /test/__init__.py: -------------------------------------------------------------------------------- 1 | from .aligned_allocator import aligned_allocator 2 | from .debug_allocator import debug_allocator 3 | from .default_allocator import default_allocator 4 | from .null_allocator import null_allocator 5 | from .page_aligned_allocator import page_aligned_allocator 6 | -------------------------------------------------------------------------------- /test/default_allocator.py: -------------------------------------------------------------------------------- 1 | import numpy_allocator 2 | 3 | 4 | class default_allocator(metaclass=numpy_allocator.type): 5 | _handler_ = numpy_allocator.default_handler 6 | 7 | 8 | def main(): 9 | import numpy as np 10 | 11 | with default_allocator: 12 | print(np.core.multiarray.get_handler_name(), 13 | np.core.multiarray.get_handler_version()) 14 | 15 | 16 | if __name__ == '__main__': 17 | main() 18 | -------------------------------------------------------------------------------- /test/null_allocator.py: -------------------------------------------------------------------------------- 1 | import numpy_allocator 2 | 3 | 4 | class null_allocator(metaclass=numpy_allocator.type): 5 | 6 | def _calloc_(nelem, elsize): 7 | return None 8 | 9 | def _free_(ptr, size): 10 | pass 11 | 12 | def _malloc_(size): 13 | return None 14 | 15 | def _realloc_(ptr, new_size): 16 | return None 17 | 18 | 19 | def main(): 20 | import numpy as np 21 | 22 | with np.testing.assert_raises(MemoryError): 23 | with null_allocator: 24 | np.ndarray(()) 25 | 26 | 27 | if __name__ == '__main__': 28 | main() 29 | -------------------------------------------------------------------------------- /.github/workflows/pypi.yml: -------------------------------------------------------------------------------- 1 | name: PyPI 2 | on: 3 | push: 4 | tags: 5 | - v*.*.* 6 | jobs: 7 | pipeline: 8 | runs-on: ubuntu-latest 9 | steps: 10 | - name: Checkout 11 | uses: actions/checkout@v4 12 | - name: Set up Python 13 | uses: actions/setup-python@v5 14 | with: 15 | python-version: 3.x 16 | - name: Create source distribution 17 | run: pipx run build --sdist 18 | - name: Create (binary) wheel distributions 19 | uses: pypa/cibuildwheel@v2.21.2 20 | with: 21 | output-dir: dist 22 | - name: Publish distributions to PyPI 23 | uses: pypa/gh-action-pypi-publish@release/v1 24 | with: 25 | password: ${{ secrets.PYPI }} 26 | -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | ## Memory management in [NumPy](https://numpy.org)* 2 | 3 | [![Binder](https://mybinder.org/badge_logo.svg)](https://mybinder.org/v2/gh/inaccel/numpy-allocator/master?labpath=NumPy-Allocator.ipynb) 4 | [![PyPI version](https://badge.fury.io/py/numpy-allocator.svg)](https://badge.fury.io/py/numpy-allocator) 5 | 6 | **NumPy is a trademark owned by [NumFOCUS](https://numfocus.org).* 7 | 8 | #### Customize Memory Allocators 9 | 10 | Α metaclass is used to override the internal data memory routines. The metaclass has four optional fields: 11 | 12 | ```python 13 | >>> import ctypes 14 | >>> import ctypes.util 15 | >>> import numpy_allocator 16 | >>> my = ctypes.CDLL(ctypes.util.find_library('my')) 17 | >>> class my_allocator(metaclass=numpy_allocator.type): 18 | ... _calloc_ = ctypes.addressof(my.calloc_func) 19 | ... _free_ = ctypes.addressof(my.free_func) 20 | ... _malloc_ = ctypes.addressof(my.malloc_func) 21 | ... _realloc_ = ctypes.addressof(my.realloc_func) 22 | ... 23 | ``` 24 | 25 | #### An example using the allocator 26 | 27 | ```python 28 | >>> import numpy as np 29 | >>> with my_allocator: 30 | ... a = np.array([1, 2, 3]) 31 | ... 32 | >>> my_allocator.handles(a) 33 | True 34 | ``` 35 | -------------------------------------------------------------------------------- /test/page_aligned_allocator.py: -------------------------------------------------------------------------------- 1 | from ctypes import CDLL, c_int, c_size_t, c_void_p 2 | from ctypes.util import find_library 3 | from mmap import PAGESIZE 4 | import numpy_allocator 5 | 6 | std = CDLL(find_library('c')) 7 | 8 | std.free.argtypes = [c_void_p] 9 | std.free.restype = None 10 | 11 | std.memalign.argtypes = [c_size_t, c_size_t] 12 | std.memalign.restype = c_void_p 13 | 14 | std.memcpy.argtypes = [c_void_p, c_void_p, c_size_t] 15 | std.memcpy.restype = c_void_p 16 | 17 | std.memset.argtypes = [c_void_p, c_int, c_size_t] 18 | std.memset.restype = c_void_p 19 | 20 | std.realloc.argtypes = [c_void_p, c_size_t] 21 | std.realloc.restype = c_void_p 22 | 23 | 24 | class page_aligned_allocator(metaclass=numpy_allocator.type): 25 | 26 | def _calloc_(nelem, elsize): 27 | result = std.memalign(PAGESIZE, nelem * elsize) 28 | if result: 29 | result = std.memset(result, 0, nelem * elsize) 30 | return result 31 | 32 | def _malloc_(size): 33 | return std.memalign(PAGESIZE, size) 34 | 35 | def _realloc_(ptr, new_size): 36 | result = std.realloc(ptr, new_size) 37 | if result and result % PAGESIZE != 0: 38 | tmp = result 39 | result = std.memalign(PAGESIZE, new_size) 40 | if result: 41 | result = std.memcpy(result, tmp, new_size) 42 | std.free(tmp) 43 | return result 44 | 45 | 46 | def main(): 47 | import numpy as np 48 | 49 | with page_aligned_allocator: 50 | print(page_aligned_allocator) 51 | 52 | np.core.test() 53 | 54 | 55 | if __name__ == '__main__': 56 | main() 57 | -------------------------------------------------------------------------------- /setup.py: -------------------------------------------------------------------------------- 1 | from setuptools import Extension, setup 2 | 3 | import numpy 4 | 5 | 6 | def README(): 7 | with open('README.md') as md: 8 | return md.read() 9 | 10 | 11 | setup( 12 | name='numpy-allocator', 13 | use_scm_version=True, 14 | description='Configurable memory allocations', 15 | long_description=README(), 16 | long_description_content_type='text/markdown', 17 | author='InAccel', 18 | author_email='info@inaccel.com', 19 | url='https://github.com/inaccel/numpy-allocator', 20 | ext_modules=[ 21 | Extension( 22 | name='numpy_allocator', 23 | sources=[ 24 | 'numpy_allocator.c', 25 | ], 26 | include_dirs=[ 27 | numpy.get_include(), 28 | ], 29 | ), 30 | ], 31 | classifiers=[ 32 | 'Development Status :: 5 - Production/Stable', 33 | 'License :: OSI Approved :: Apache Software License', 34 | 'Operating System :: POSIX :: Linux', 35 | 'Programming Language :: Python', 36 | 'Programming Language :: Python :: 3', 37 | 'Programming Language :: Python :: 3.8', 38 | 'Programming Language :: Python :: 3.9', 39 | 'Programming Language :: Python :: 3.10', 40 | 'Programming Language :: Python :: 3.11', 41 | 'Programming Language :: Python :: 3.12', 42 | 'Programming Language :: Python :: 3.13', 43 | 'Programming Language :: Python :: Implementation :: CPython', 44 | ], 45 | license='Apache-2.0', 46 | platforms=[ 47 | 'Linux', 48 | ], 49 | install_requires=[ 50 | 'numpy>=1.22', 51 | ], 52 | python_requires='>=3.8', 53 | ) 54 | -------------------------------------------------------------------------------- /test/debug_allocator.py: -------------------------------------------------------------------------------- 1 | from ctypes import CDLL, c_size_t, c_void_p 2 | from ctypes.util import find_library 3 | import numpy_allocator 4 | 5 | std = CDLL(find_library('c')) 6 | 7 | std.calloc.argtypes = [c_size_t, c_size_t] 8 | std.calloc.restype = c_void_p 9 | 10 | std.free.argtypes = [c_void_p] 11 | std.free.restype = None 12 | 13 | std.malloc.argtypes = [c_size_t] 14 | std.malloc.restype = c_void_p 15 | 16 | std.realloc.argtypes = [c_void_p, c_size_t] 17 | std.realloc.restype = c_void_p 18 | 19 | 20 | class debug_allocator(metaclass=numpy_allocator.type): 21 | 22 | def _calloc_(nelem, elsize): 23 | result = std.calloc(nelem, elsize) 24 | if result: 25 | print('0x%x calloc(%d, %d)' % (result, nelem, elsize)) 26 | else: 27 | print('calloc(%d, %d)' % (nelem, elsize)) 28 | return result 29 | 30 | def _free_(ptr, size): 31 | std.free(ptr) 32 | print('free(0x%x)' % ptr) 33 | 34 | def _malloc_(size): 35 | result = std.malloc(size) 36 | if result: 37 | print('0x%x malloc(%d)' % (result, size)) 38 | else: 39 | print('malloc(%d)' % size) 40 | return result 41 | 42 | def _realloc_(ptr, new_size): 43 | result = std.realloc(ptr, new_size) 44 | if result: 45 | print('0x%x realloc(0x%x, %d)' % (result, ptr, new_size)) 46 | else: 47 | print('realloc(0x%x, %d)' % (ptr, new_size)) 48 | return result 49 | 50 | 51 | def main(): 52 | import numpy as np 53 | 54 | with debug_allocator: 55 | a = np.arange(15).reshape((3, 5)) 56 | 57 | assert debug_allocator.handles(a) 58 | 59 | b = np.array([6, 7, 8]) 60 | 61 | assert not debug_allocator.handles(b) 62 | 63 | 64 | if __name__ == '__main__': 65 | main() 66 | -------------------------------------------------------------------------------- /test/aligned_allocator.py: -------------------------------------------------------------------------------- 1 | from ctypes import CDLL, c_int, c_size_t, c_void_p 2 | from ctypes.util import find_library 3 | import numpy_allocator 4 | 5 | std = CDLL(find_library('c')) 6 | 7 | std.free.argtypes = [c_void_p] 8 | std.free.restype = None 9 | 10 | std.memalign.argtypes = [c_size_t, c_size_t] 11 | std.memalign.restype = c_void_p 12 | 13 | std.memcpy.argtypes = [c_void_p, c_void_p, c_size_t] 14 | std.memcpy.restype = c_void_p 15 | 16 | std.memset.argtypes = [c_void_p, c_int, c_size_t] 17 | std.memset.restype = c_void_p 18 | 19 | std.realloc.argtypes = [c_void_p, c_size_t] 20 | std.realloc.restype = c_void_p 21 | 22 | 23 | class aligned_allocator(numpy_allocator.object): 24 | 25 | def __init__(self, alignment): 26 | self.alignment = alignment 27 | 28 | def __str__(self): 29 | return '{}({})'.format(self.__class__.__name__, self.alignment) 30 | 31 | def _calloc_(self, nelem, elsize): 32 | result = std.memalign(self.alignment, nelem * elsize) 33 | if result: 34 | result = std.memset(result, 0, nelem * elsize) 35 | return result 36 | 37 | def _malloc_(self, size): 38 | return std.memalign(self.alignment, size) 39 | 40 | def _realloc_(self, ptr, new_size): 41 | result = std.realloc(ptr, new_size) 42 | if result and result % self.alignment != 0: 43 | tmp = result 44 | result = std.memalign(self.alignment, new_size) 45 | if result: 46 | result = std.memcpy(result, tmp, new_size) 47 | std.free(tmp) 48 | return result 49 | 50 | 51 | def main(): 52 | from mmap import PAGESIZE 53 | import numpy as np 54 | 55 | with aligned_allocator(PAGESIZE) as page_aligned_allocator: 56 | print(page_aligned_allocator) 57 | 58 | np.core.test() 59 | 60 | 61 | if __name__ == '__main__': 62 | main() 63 | -------------------------------------------------------------------------------- /NumPy-Allocator.ipynb: -------------------------------------------------------------------------------- 1 | { 2 | "cells": [ 3 | { 4 | "cell_type": "markdown", 5 | "metadata": {}, 6 | "source": [ 7 | "### Aligned Allocator [✍](test/aligned_allocator.py)" 8 | ] 9 | }, 10 | { 11 | "cell_type": "code", 12 | "execution_count": null, 13 | "metadata": {}, 14 | "outputs": [], 15 | "source": [ 16 | "%%python\n", 17 | "from test import aligned_allocator\n", 18 | "\n", 19 | "from mmap import PAGESIZE\n", 20 | "import numpy as np\n", 21 | "\n", 22 | "with aligned_allocator(PAGESIZE) as page_aligned_allocator:\n", 23 | " print(page_aligned_allocator)\n", 24 | "\n", 25 | " np.core.test()" 26 | ] 27 | }, 28 | { 29 | "cell_type": "markdown", 30 | "metadata": {}, 31 | "source": [ 32 | "### Debug Allocator [✍](test/debug_allocator.py)" 33 | ] 34 | }, 35 | { 36 | "cell_type": "code", 37 | "execution_count": null, 38 | "metadata": {}, 39 | "outputs": [], 40 | "source": [ 41 | "%%python\n", 42 | "from test import debug_allocator\n", 43 | "\n", 44 | "import numpy as np\n", 45 | "\n", 46 | "with debug_allocator:\n", 47 | " a = np.arange(15).reshape((3, 5))\n", 48 | "\n", 49 | "assert debug_allocator.handles(a)\n", 50 | "\n", 51 | "b = np.array([6, 7, 8])\n", 52 | "\n", 53 | "assert not debug_allocator.handles(b)" 54 | ] 55 | }, 56 | { 57 | "cell_type": "markdown", 58 | "metadata": {}, 59 | "source": [ 60 | "### Default Allocator [✍](test/default_allocator.py)" 61 | ] 62 | }, 63 | { 64 | "cell_type": "code", 65 | "execution_count": null, 66 | "metadata": {}, 67 | "outputs": [], 68 | "source": [ 69 | "%%python\n", 70 | "from test import default_allocator\n", 71 | "\n", 72 | "import numpy as np\n", 73 | "\n", 74 | "with default_allocator:\n", 75 | " print(np.core.multiarray.get_handler_name(),\n", 76 | " np.core.multiarray.get_handler_version())" 77 | ] 78 | }, 79 | { 80 | "cell_type": "markdown", 81 | "metadata": {}, 82 | "source": [ 83 | "### Null Allocator [✍](test/null_allocator.py)" 84 | ] 85 | }, 86 | { 87 | "cell_type": "code", 88 | "execution_count": null, 89 | "metadata": {}, 90 | "outputs": [], 91 | "source": [ 92 | "%%python\n", 93 | "from test import null_allocator\n", 94 | "\n", 95 | "import numpy as np\n", 96 | "\n", 97 | "with np.testing.assert_raises(MemoryError):\n", 98 | " with null_allocator:\n", 99 | " np.ndarray(())" 100 | ] 101 | }, 102 | { 103 | "cell_type": "markdown", 104 | "metadata": {}, 105 | "source": [ 106 | "### Page Aligned Allocator [✍](test/page_aligned_allocator.py)" 107 | ] 108 | }, 109 | { 110 | "cell_type": "code", 111 | "execution_count": null, 112 | "metadata": {}, 113 | "outputs": [], 114 | "source": [ 115 | "%%python\n", 116 | "from test import page_aligned_allocator\n", 117 | "\n", 118 | "import numpy as np\n", 119 | "\n", 120 | "with page_aligned_allocator:\n", 121 | " print(page_aligned_allocator)\n", 122 | "\n", 123 | " np.core.test()" 124 | ] 125 | } 126 | ], 127 | "metadata": { 128 | "kernelspec": { 129 | "display_name": "Python 3", 130 | "language": "python", 131 | "name": "python3" 132 | } 133 | }, 134 | "nbformat": 4, 135 | "nbformat_minor": 5 136 | } 137 | -------------------------------------------------------------------------------- /LICENSE: -------------------------------------------------------------------------------- 1 | 2 | Apache License 3 | Version 2.0, January 2004 4 | http://www.apache.org/licenses/ 5 | 6 | TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION 7 | 8 | 1. Definitions. 9 | 10 | "License" shall mean the terms and conditions for use, reproduction, 11 | and distribution as defined by Sections 1 through 9 of this document. 12 | 13 | "Licensor" shall mean the copyright owner or entity authorized by 14 | the copyright owner that is granting the License. 15 | 16 | "Legal Entity" shall mean the union of the acting entity and all 17 | other entities that control, are controlled by, or are under common 18 | control with that entity. For the purposes of this definition, 19 | "control" means (i) the power, direct or indirect, to cause the 20 | direction or management of such entity, whether by contract or 21 | otherwise, or (ii) ownership of fifty percent (50%) or more of the 22 | outstanding shares, or (iii) beneficial ownership of such entity. 23 | 24 | "You" (or "Your") shall mean an individual or Legal Entity 25 | exercising permissions granted by this License. 26 | 27 | "Source" form shall mean the preferred form for making modifications, 28 | including but not limited to software source code, documentation 29 | source, and configuration files. 30 | 31 | "Object" form shall mean any form resulting from mechanical 32 | transformation or translation of a Source form, including but 33 | not limited to compiled object code, generated documentation, 34 | and conversions to other media types. 35 | 36 | "Work" shall mean the work of authorship, whether in Source or 37 | Object form, made available under the License, as indicated by a 38 | copyright notice that is included in or attached to the work 39 | (an example is provided in the Appendix below). 40 | 41 | "Derivative Works" shall mean any work, whether in Source or Object 42 | form, that is based on (or derived from) the Work and for which the 43 | editorial revisions, annotations, elaborations, or other modifications 44 | represent, as a whole, an original work of authorship. For the purposes 45 | of this License, Derivative Works shall not include works that remain 46 | separable from, or merely link (or bind by name) to the interfaces of, 47 | the Work and Derivative Works thereof. 48 | 49 | "Contribution" shall mean any work of authorship, including 50 | the original version of the Work and any modifications or additions 51 | to that Work or Derivative Works thereof, that is intentionally 52 | submitted to Licensor for inclusion in the Work by the copyright owner 53 | or by an individual or Legal Entity authorized to submit on behalf of 54 | the copyright owner. For the purposes of this definition, "submitted" 55 | means any form of electronic, verbal, or written communication sent 56 | to the Licensor or its representatives, including but not limited to 57 | communication on electronic mailing lists, source code control systems, 58 | and issue tracking systems that are managed by, or on behalf of, the 59 | Licensor for the purpose of discussing and improving the Work, but 60 | excluding communication that is conspicuously marked or otherwise 61 | designated in writing by the copyright owner as "Not a Contribution." 62 | 63 | "Contributor" shall mean Licensor and any individual or Legal Entity 64 | on behalf of whom a Contribution has been received by Licensor and 65 | subsequently incorporated within the Work. 66 | 67 | 2. Grant of Copyright License. Subject to the terms and conditions of 68 | this License, each Contributor hereby grants to You a perpetual, 69 | worldwide, non-exclusive, no-charge, royalty-free, irrevocable 70 | copyright license to reproduce, prepare Derivative Works of, 71 | publicly display, publicly perform, sublicense, and distribute the 72 | Work and such Derivative Works in Source or Object form. 73 | 74 | 3. Grant of Patent License. Subject to the terms and conditions of 75 | this License, each Contributor hereby grants to You a perpetual, 76 | worldwide, non-exclusive, no-charge, royalty-free, irrevocable 77 | (except as stated in this section) patent license to make, have made, 78 | use, offer to sell, sell, import, and otherwise transfer the Work, 79 | where such license applies only to those patent claims licensable 80 | by such Contributor that are necessarily infringed by their 81 | Contribution(s) alone or by combination of their Contribution(s) 82 | with the Work to which such Contribution(s) was submitted. If You 83 | institute patent litigation against any entity (including a 84 | cross-claim or counterclaim in a lawsuit) alleging that the Work 85 | or a Contribution incorporated within the Work constitutes direct 86 | or contributory patent infringement, then any patent licenses 87 | granted to You under this License for that Work shall terminate 88 | as of the date such litigation is filed. 89 | 90 | 4. Redistribution. You may reproduce and distribute copies of the 91 | Work or Derivative Works thereof in any medium, with or without 92 | modifications, and in Source or Object form, provided that You 93 | meet the following conditions: 94 | 95 | (a) You must give any other recipients of the Work or 96 | Derivative Works a copy of this License; and 97 | 98 | (b) You must cause any modified files to carry prominent notices 99 | stating that You changed the files; and 100 | 101 | (c) You must retain, in the Source form of any Derivative Works 102 | that You distribute, all copyright, patent, trademark, and 103 | attribution notices from the Source form of the Work, 104 | excluding those notices that do not pertain to any part of 105 | the Derivative Works; and 106 | 107 | (d) If the Work includes a "NOTICE" text file as part of its 108 | distribution, then any Derivative Works that You distribute must 109 | include a readable copy of the attribution notices contained 110 | within such NOTICE file, excluding those notices that do not 111 | pertain to any part of the Derivative Works, in at least one 112 | of the following places: within a NOTICE text file distributed 113 | as part of the Derivative Works; within the Source form or 114 | documentation, if provided along with the Derivative Works; or, 115 | within a display generated by the Derivative Works, if and 116 | wherever such third-party notices normally appear. The contents 117 | of the NOTICE file are for informational purposes only and 118 | do not modify the License. You may add Your own attribution 119 | notices within Derivative Works that You distribute, alongside 120 | or as an addendum to the NOTICE text from the Work, provided 121 | that such additional attribution notices cannot be construed 122 | as modifying the License. 123 | 124 | You may add Your own copyright statement to Your modifications and 125 | may provide additional or different license terms and conditions 126 | for use, reproduction, or distribution of Your modifications, or 127 | for any such Derivative Works as a whole, provided Your use, 128 | reproduction, and distribution of the Work otherwise complies with 129 | the conditions stated in this License. 130 | 131 | 5. Submission of Contributions. Unless You explicitly state otherwise, 132 | any Contribution intentionally submitted for inclusion in the Work 133 | by You to the Licensor shall be under the terms and conditions of 134 | this License, without any additional terms or conditions. 135 | Notwithstanding the above, nothing herein shall supersede or modify 136 | the terms of any separate license agreement you may have executed 137 | with Licensor regarding such Contributions. 138 | 139 | 6. Trademarks. This License does not grant permission to use the trade 140 | names, trademarks, service marks, or product names of the Licensor, 141 | except as required for reasonable and customary use in describing the 142 | origin of the Work and reproducing the content of the NOTICE file. 143 | 144 | 7. Disclaimer of Warranty. Unless required by applicable law or 145 | agreed to in writing, Licensor provides the Work (and each 146 | Contributor provides its Contributions) on an "AS IS" BASIS, 147 | WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or 148 | implied, including, without limitation, any warranties or conditions 149 | of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A 150 | PARTICULAR PURPOSE. You are solely responsible for determining the 151 | appropriateness of using or redistributing the Work and assume any 152 | risks associated with Your exercise of permissions under this License. 153 | 154 | 8. Limitation of Liability. In no event and under no legal theory, 155 | whether in tort (including negligence), contract, or otherwise, 156 | unless required by applicable law (such as deliberate and grossly 157 | negligent acts) or agreed to in writing, shall any Contributor be 158 | liable to You for damages, including any direct, indirect, special, 159 | incidental, or consequential damages of any character arising as a 160 | result of this License or out of the use or inability to use the 161 | Work (including but not limited to damages for loss of goodwill, 162 | work stoppage, computer failure or malfunction, or any and all 163 | other commercial damages or losses), even if such Contributor 164 | has been advised of the possibility of such damages. 165 | 166 | 9. Accepting Warranty or Additional Liability. While redistributing 167 | the Work or Derivative Works thereof, You may choose to offer, 168 | and charge a fee for, acceptance of support, warranty, indemnity, 169 | or other liability obligations and/or rights consistent with this 170 | License. However, in accepting such obligations, You may act only 171 | on Your own behalf and on Your sole responsibility, not on behalf 172 | of any other Contributor, and only if You agree to indemnify, 173 | defend, and hold each Contributor harmless for any liability 174 | incurred by, or claims asserted against, such Contributor by reason 175 | of your accepting any such warranty or additional liability. 176 | 177 | END OF TERMS AND CONDITIONS 178 | -------------------------------------------------------------------------------- /numpy_allocator.c: -------------------------------------------------------------------------------- 1 | #define NPY_NO_DEPRECATED_API NPY_1_22_API_VERSION 2 | #define NPY_TARGET_VERSION NPY_1_22_API_VERSION 3 | 4 | #include 5 | #include 6 | 7 | typedef struct { 8 | void *calloc; 9 | void *free; 10 | void *malloc; 11 | void *realloc; 12 | } PyDataMem_Funcs; 13 | 14 | static void *call_realloc(PyObject *_realloc_, void *ptr, size_t new_size) { 15 | PyObject *_ptr; 16 | if (!ptr) { 17 | Py_INCREF(Py_None); 18 | _ptr = Py_None; 19 | } else { 20 | _ptr = PyLong_FromVoidPtr(ptr); 21 | if (!_ptr) { 22 | return NULL; 23 | } 24 | } 25 | PyObject *_new_size = PyLong_FromSize_t(new_size); 26 | if (!_new_size) { 27 | Py_DECREF(_ptr); 28 | 29 | return NULL; 30 | } 31 | PyObject *_new_ptr = PyObject_CallFunctionObjArgs(_realloc_, _ptr, _new_size, NULL); 32 | Py_DECREF(_new_size); 33 | Py_DECREF(_ptr); 34 | if (!_new_ptr) { 35 | return NULL; 36 | } 37 | void *new_ptr = NULL; 38 | if (_new_ptr != Py_None) { 39 | new_ptr = PyLong_AsVoidPtr(_new_ptr); 40 | } 41 | Py_DECREF(_new_ptr); 42 | return new_ptr; 43 | } 44 | 45 | static void *safe_realloc(void *ctx, void *ptr, size_t new_size) { 46 | PyGILState_STATE state = PyGILState_Ensure(); 47 | PyObject *type; 48 | PyObject *value; 49 | PyObject *traceback; 50 | PyErr_Fetch(&type, &value, &traceback); 51 | void *new_ptr = call_realloc((PyObject *) ((PyDataMem_Funcs *) ctx)->realloc, ptr, new_size); 52 | if (PyErr_Occurred()) { 53 | PyErr_WriteUnraisable((PyObject *) ((PyDataMem_Funcs *) ctx)->realloc); 54 | } 55 | PyErr_Restore(type, value, traceback); 56 | PyGILState_Release(state); 57 | return new_ptr; 58 | } 59 | 60 | typedef void *(PyDataMem_ReallocFunc)(void *ptr, size_t new_size); 61 | 62 | static void *unsafe_realloc(void *ctx, void *ptr, size_t new_size) { 63 | void **_realloc_ = ((PyDataMem_Funcs *) ctx)->realloc; 64 | return ((PyDataMem_ReallocFunc *) *_realloc_)(ptr, new_size); 65 | } 66 | 67 | static void *default_realloc(void *ctx, void *ptr, size_t new_size) { 68 | return realloc(ptr, new_size); 69 | } 70 | 71 | static void *call_malloc(PyObject *_malloc_, size_t size) { 72 | PyObject *_size = PyLong_FromSize_t(size); 73 | if (!_size) { 74 | return NULL; 75 | } 76 | PyObject *_ptr = PyObject_CallFunctionObjArgs(_malloc_, _size, NULL); 77 | Py_DECREF(_size); 78 | if (!_ptr) { 79 | return NULL; 80 | } 81 | void *ptr = NULL; 82 | if (_ptr != Py_None) { 83 | ptr = PyLong_AsVoidPtr(_ptr); 84 | } 85 | Py_DECREF(_ptr); 86 | return ptr; 87 | } 88 | 89 | static void *safe_malloc(void *ctx, size_t size) { 90 | PyGILState_STATE state = PyGILState_Ensure(); 91 | PyObject *type; 92 | PyObject *value; 93 | PyObject *traceback; 94 | PyErr_Fetch(&type, &value, &traceback); 95 | void *ptr = call_malloc((PyObject *) ((PyDataMem_Funcs *) ctx)->malloc, size); 96 | if (PyErr_Occurred()) { 97 | PyErr_WriteUnraisable((PyObject *) ((PyDataMem_Funcs *) ctx)->malloc); 98 | } 99 | PyErr_Restore(type, value, traceback); 100 | PyGILState_Release(state); 101 | return ptr; 102 | } 103 | 104 | typedef void *(PyDataMem_MallocFunc)(size_t size); 105 | 106 | static void *unsafe_malloc(void *ctx, size_t size) { 107 | void **_malloc_ = ((PyDataMem_Funcs *) ctx)->malloc; 108 | return ((PyDataMem_MallocFunc *) *_malloc_)(size); 109 | } 110 | 111 | static void *default_malloc(void *ctx, size_t size) { 112 | return malloc(size); 113 | } 114 | 115 | static void call_free(PyObject *_free_, void *ptr, size_t size) { 116 | PyObject *_ptr; 117 | if (!ptr) { 118 | Py_INCREF(Py_None); 119 | _ptr = Py_None; 120 | } else { 121 | _ptr = PyLong_FromVoidPtr(ptr); 122 | if (!_ptr) { 123 | return; 124 | } 125 | } 126 | PyObject *_size = PyLong_FromSize_t(size); 127 | if (!_size) { 128 | Py_DECREF(_ptr); 129 | 130 | return; 131 | } 132 | Py_XDECREF(PyObject_CallFunctionObjArgs(_free_, _ptr, _size, NULL)); 133 | Py_DECREF(_size); 134 | Py_DECREF(_ptr); 135 | } 136 | 137 | static void safe_free(void *ctx, void *ptr, size_t size) { 138 | PyGILState_STATE state = PyGILState_Ensure(); 139 | PyObject *type; 140 | PyObject *value; 141 | PyObject *traceback; 142 | PyErr_Fetch(&type, &value, &traceback); 143 | call_free((PyObject *) ((PyDataMem_Funcs *) ctx)->free, ptr, size); 144 | if (PyErr_Occurred()) { 145 | PyErr_WriteUnraisable((PyObject *) ((PyDataMem_Funcs *) ctx)->free); 146 | } 147 | PyErr_Restore(type, value, traceback); 148 | PyGILState_Release(state); 149 | } 150 | 151 | typedef void (PyDataMem_FreeFunc)(void *ptr, size_t size); 152 | 153 | static void unsafe_free(void *ctx, void *ptr, size_t size) { 154 | void **_free_ = ((PyDataMem_Funcs *) ctx)->free; 155 | ((PyDataMem_FreeFunc *) *_free_)(ptr, size); 156 | } 157 | 158 | static void default_free(void *ctx, void *ptr, size_t size) { 159 | free(ptr); 160 | } 161 | 162 | static void *call_calloc(PyObject *_calloc_, size_t nelem, size_t elsize) { 163 | PyObject *_nelem = PyLong_FromSize_t(nelem); 164 | if (!_nelem) { 165 | return NULL; 166 | } 167 | PyObject *_elsize = PyLong_FromSize_t(elsize); 168 | if (!_elsize) { 169 | Py_DECREF(_nelem); 170 | 171 | return NULL; 172 | } 173 | PyObject *_ptr = PyObject_CallFunctionObjArgs(_calloc_, _nelem, _elsize, NULL); 174 | Py_DECREF(_elsize); 175 | Py_DECREF(_nelem); 176 | if (!_ptr) { 177 | return NULL; 178 | } 179 | void *ptr = NULL; 180 | if (_ptr != Py_None) { 181 | ptr = PyLong_AsVoidPtr(_ptr); 182 | } 183 | Py_DECREF(_ptr); 184 | return ptr; 185 | } 186 | 187 | static void *safe_calloc(void *ctx, size_t nelem, size_t elsize) { 188 | PyGILState_STATE state = PyGILState_Ensure(); 189 | PyObject *type; 190 | PyObject *value; 191 | PyObject *traceback; 192 | PyErr_Fetch(&type, &value, &traceback); 193 | void *ptr = call_calloc((PyObject *) ((PyDataMem_Funcs *) ctx)->calloc, nelem, elsize); 194 | if (PyErr_Occurred()) { 195 | PyErr_WriteUnraisable((PyObject *) ((PyDataMem_Funcs *) ctx)->calloc); 196 | } 197 | PyErr_Restore(type, value, traceback); 198 | PyGILState_Release(state); 199 | return ptr; 200 | } 201 | 202 | typedef void *(PyDataMem_CallocFunc)(size_t nelem, size_t elsize); 203 | 204 | static void *unsafe_calloc(void *ctx, size_t nelem, size_t elsize) { 205 | void **_calloc_ = ((PyDataMem_Funcs *) ctx)->calloc; 206 | return ((PyDataMem_CallocFunc *) *_calloc_)(nelem, elsize); 207 | } 208 | 209 | static void *default_calloc(void *ctx, size_t nelem, size_t elsize) { 210 | return calloc(nelem, elsize); 211 | } 212 | 213 | static void handler_destructor(PyObject *handler) { 214 | PyDataMem_Handler *mem_handler = (PyDataMem_Handler *) PyCapsule_GetPointer(handler, "mem_handler"); 215 | if (!mem_handler) { 216 | return; 217 | } 218 | 219 | Py_XDECREF(((PyDataMem_Funcs *) mem_handler->allocator.ctx)->realloc); 220 | 221 | Py_XDECREF(((PyDataMem_Funcs *) mem_handler->allocator.ctx)->malloc); 222 | 223 | Py_XDECREF(((PyDataMem_Funcs *) mem_handler->allocator.ctx)->free); 224 | 225 | Py_XDECREF(((PyDataMem_Funcs *) mem_handler->allocator.ctx)->calloc); 226 | 227 | free(mem_handler->allocator.ctx); 228 | 229 | free(mem_handler); 230 | } 231 | 232 | static PyObject *handler(PyObject *allocator, PyObject *args) { 233 | if (PyObject_HasAttrString(allocator, "_handler_")) { 234 | return PyObject_GetAttrString(allocator, "_handler_"); 235 | } else { 236 | PyDataMem_Handler *mem_handler = (PyDataMem_Handler *) calloc(1, sizeof(PyDataMem_Handler)); 237 | if (!mem_handler) { 238 | PyErr_NoMemory(); 239 | 240 | return NULL; 241 | } 242 | 243 | mem_handler->allocator.ctx = calloc(1, sizeof(PyDataMem_Funcs)); 244 | if (!mem_handler->allocator.ctx) { 245 | free(mem_handler); 246 | 247 | PyErr_NoMemory(); 248 | 249 | return NULL; 250 | } 251 | 252 | PyObject *handler = PyCapsule_New(mem_handler, "mem_handler", handler_destructor); 253 | if (!handler) { 254 | free(mem_handler->allocator.ctx); 255 | 256 | free(mem_handler); 257 | 258 | return NULL; 259 | } 260 | 261 | PyObject *name = PyObject_Str(allocator); 262 | if (!name) { 263 | Py_DECREF(handler); 264 | 265 | return NULL; 266 | } 267 | strncpy(mem_handler->name, PyUnicode_AsUTF8(name), sizeof(((PyDataMem_Handler *) NULL)->name) - 1); 268 | Py_DECREF(name); 269 | 270 | mem_handler->version = 1; 271 | 272 | if (PyObject_HasAttrString(allocator, "_calloc_")) { 273 | PyObject *_calloc_ = PyObject_GetAttrString(allocator, "_calloc_"); 274 | if (!_calloc_) { 275 | Py_DECREF(handler); 276 | 277 | return NULL; 278 | } else if (_calloc_ == Py_None) { 279 | Py_DECREF(_calloc_); 280 | mem_handler->allocator.calloc = default_calloc; 281 | } else if (PyLong_Check(_calloc_)) { 282 | void *ptr = PyLong_AsVoidPtr(_calloc_); 283 | Py_DECREF(_calloc_); 284 | if (!ptr) { 285 | Py_DECREF(handler); 286 | 287 | return NULL; 288 | } 289 | ((PyDataMem_Funcs *) mem_handler->allocator.ctx)->calloc = ptr; 290 | mem_handler->allocator.calloc = unsafe_calloc; 291 | } else if (PyCallable_Check(_calloc_)) { 292 | ((PyDataMem_Funcs *) mem_handler->allocator.ctx)->calloc = _calloc_; 293 | mem_handler->allocator.calloc = safe_calloc; 294 | } else { 295 | Py_DECREF(handler); 296 | 297 | PyErr_SetString(PyExc_TypeError, "_calloc_ must be an integer address or a callable"); 298 | return NULL; 299 | } 300 | } else { 301 | mem_handler->allocator.calloc = default_calloc; 302 | } 303 | 304 | if (PyObject_HasAttrString(allocator, "_free_")) { 305 | PyObject *_free_ = PyObject_GetAttrString(allocator, "_free_"); 306 | if (!_free_) { 307 | Py_DECREF(handler); 308 | 309 | return NULL; 310 | } else if (_free_ == Py_None) { 311 | Py_DECREF(_free_); 312 | mem_handler->allocator.free = default_free; 313 | } else if (PyLong_Check(_free_)) { 314 | void *ptr = PyLong_AsVoidPtr(_free_); 315 | Py_DECREF(_free_); 316 | if (!ptr) { 317 | Py_DECREF(handler); 318 | 319 | return NULL; 320 | } 321 | ((PyDataMem_Funcs *) mem_handler->allocator.ctx)->free = ptr; 322 | mem_handler->allocator.free = unsafe_free; 323 | } else if (PyCallable_Check(_free_)) { 324 | ((PyDataMem_Funcs *) mem_handler->allocator.ctx)->free = _free_; 325 | mem_handler->allocator.free = safe_free; 326 | } else { 327 | Py_DECREF(handler); 328 | 329 | PyErr_SetString(PyExc_TypeError, "_free_ must be an integer address or a callable"); 330 | return NULL; 331 | } 332 | } else { 333 | mem_handler->allocator.free = default_free; 334 | } 335 | 336 | if (PyObject_HasAttrString(allocator, "_malloc_")) { 337 | PyObject *_malloc_ = PyObject_GetAttrString(allocator, "_malloc_"); 338 | if (!_malloc_) { 339 | Py_DECREF(handler); 340 | 341 | return NULL; 342 | } else if (_malloc_ == Py_None) { 343 | Py_DECREF(_malloc_); 344 | mem_handler->allocator.malloc = default_malloc; 345 | } else if (PyLong_Check(_malloc_)) { 346 | void *ptr = PyLong_AsVoidPtr(_malloc_); 347 | Py_DECREF(_malloc_); 348 | if (!ptr) { 349 | Py_DECREF(handler); 350 | 351 | return NULL; 352 | } 353 | ((PyDataMem_Funcs *) mem_handler->allocator.ctx)->malloc = ptr; 354 | mem_handler->allocator.malloc = unsafe_malloc; 355 | } else if (PyCallable_Check(_malloc_)) { 356 | ((PyDataMem_Funcs *) mem_handler->allocator.ctx)->malloc = _malloc_; 357 | mem_handler->allocator.malloc = safe_malloc; 358 | } else { 359 | Py_DECREF(handler); 360 | 361 | PyErr_SetString(PyExc_TypeError, "_malloc_ must be an integer address or a callable"); 362 | return NULL; 363 | } 364 | } else { 365 | mem_handler->allocator.malloc = default_malloc; 366 | } 367 | 368 | if (PyObject_HasAttrString(allocator, "_realloc_")) { 369 | PyObject *_realloc_ = PyObject_GetAttrString(allocator, "_realloc_"); 370 | if (!_realloc_) { 371 | Py_DECREF(handler); 372 | 373 | return NULL; 374 | } else if (_realloc_ == Py_None) { 375 | Py_DECREF(_realloc_); 376 | mem_handler->allocator.realloc = default_realloc; 377 | } else if (PyLong_Check(_realloc_)) { 378 | void *ptr = PyLong_AsVoidPtr(_realloc_); 379 | Py_DECREF(_realloc_); 380 | if (!ptr) { 381 | Py_DECREF(handler); 382 | 383 | return NULL; 384 | } 385 | ((PyDataMem_Funcs *) mem_handler->allocator.ctx)->realloc = ptr; 386 | mem_handler->allocator.realloc = unsafe_realloc; 387 | } else if (PyCallable_Check(_realloc_)) { 388 | ((PyDataMem_Funcs *) mem_handler->allocator.ctx)->realloc = _realloc_; 389 | mem_handler->allocator.realloc = safe_realloc; 390 | } else { 391 | Py_DECREF(handler); 392 | 393 | PyErr_SetString(PyExc_TypeError, "_realloc_ must be an integer address or a callable"); 394 | return NULL; 395 | } 396 | } else { 397 | mem_handler->allocator.realloc = default_realloc; 398 | } 399 | 400 | if (PyObject_SetAttrString(allocator, "_handler_", handler)) { 401 | Py_DECREF(handler); 402 | 403 | return NULL; 404 | } 405 | 406 | return handler; 407 | } 408 | } 409 | 410 | static PyObject *handles(PyObject *allocator, PyObject *array) { 411 | if (!PyArray_Check(array)) { 412 | PyErr_SetString(PyExc_TypeError, "argument must be an ndarray"); 413 | return NULL; 414 | } 415 | 416 | while (array && PyArray_Check(array)) { 417 | if (PyArray_CHKFLAGS((PyArrayObject *) array, NPY_ARRAY_OWNDATA)) { 418 | PyObject *array_handler = PyArray_HANDLER((PyArrayObject *) array); 419 | if (!array_handler) { 420 | PyErr_SetString(PyExc_RuntimeError, "no memory handler found but OWNDATA flag set"); 421 | return NULL; 422 | } 423 | 424 | PyObject *allocator_handler = handler(allocator, array); 425 | if (!allocator_handler) { 426 | return NULL; 427 | } 428 | Py_DECREF(allocator_handler); 429 | 430 | if (array_handler != allocator_handler) { 431 | Py_RETURN_FALSE; 432 | } 433 | 434 | Py_RETURN_TRUE; 435 | } 436 | 437 | array = PyArray_BASE((PyArrayObject *) array); 438 | } 439 | 440 | Py_RETURN_FALSE; 441 | } 442 | 443 | static PyObject *var; 444 | 445 | static PyObject *PyContextVar_Pop(PyObject *var) { 446 | PyObject *list; 447 | if (PyContextVar_Get(var, NULL, &list)) { 448 | return NULL; 449 | } 450 | 451 | PyObject *capsule = PySequence_GetItem(list, PySequence_Size(list) - 1); 452 | if (!capsule) { 453 | Py_DECREF(list); 454 | 455 | return NULL; 456 | } 457 | 458 | int error = PySequence_DelItem(list, PySequence_Size(list) - 1); 459 | Py_DECREF(list); 460 | if (error) { 461 | Py_DECREF(capsule); 462 | 463 | return NULL; 464 | } 465 | 466 | return capsule; 467 | } 468 | 469 | static PyObject *__exit__(PyObject *allocator, PyObject *args) { 470 | PyObject *new_handler = PyContextVar_Pop(var); 471 | if (!new_handler) { 472 | return NULL; 473 | } 474 | 475 | PyObject *old_handler = PyDataMem_SetHandler(new_handler); 476 | Py_DECREF(new_handler); 477 | if (!old_handler) { 478 | return NULL; 479 | } 480 | Py_DECREF(old_handler); 481 | 482 | Py_RETURN_NONE; 483 | } 484 | 485 | static int PyContextVar_Push(PyObject *var, PyObject *capsule) { 486 | PyObject *list; 487 | if (PyContextVar_Get(var, NULL, &list)) { 488 | return -1; 489 | } 490 | 491 | int error = PyList_Append(list, capsule); 492 | Py_DECREF(list); 493 | if (error) { 494 | return -1; 495 | } 496 | 497 | return 0; 498 | } 499 | 500 | static PyObject *__enter__(PyObject *allocator, PyObject *args) { 501 | PyObject *new_handler = handler(allocator, args); 502 | if (!new_handler) { 503 | return NULL; 504 | } 505 | 506 | PyObject *old_handler = PyDataMem_SetHandler(new_handler); 507 | Py_DECREF(new_handler); 508 | if (!old_handler) { 509 | return NULL; 510 | } 511 | 512 | int error = PyContextVar_Push(var, old_handler); 513 | Py_DECREF(old_handler); 514 | if (error) { 515 | return NULL; 516 | } 517 | 518 | Py_INCREF(allocator); 519 | return allocator; 520 | } 521 | 522 | static PyMethodDef tp_methods[] = { 523 | {"__enter__", __enter__, METH_NOARGS, NULL}, 524 | {"__exit__", __exit__, METH_VARARGS, NULL}, 525 | {"handler", handler, METH_NOARGS, NULL}, 526 | {"handles", handles, METH_O, NULL}, 527 | {NULL, NULL, 0, NULL}, 528 | }; 529 | 530 | static PyObject *tp_str(PyObject *allocator) { 531 | PyObject *__name__ = PyObject_GetAttrString(allocator, "__name__"); 532 | if (!__name__) { 533 | return NULL; 534 | } 535 | 536 | PyObject *allocator_str = PyObject_Str(__name__); 537 | Py_DECREF(__name__); 538 | return allocator_str; 539 | } 540 | 541 | static PyTypeObject type = { 542 | PyVarObject_HEAD_INIT(NULL, 0) 543 | .tp_name = "numpy_allocator.type", 544 | .tp_str = tp_str, 545 | .tp_flags = Py_TPFLAGS_BASETYPE | Py_TPFLAGS_DEFAULT, 546 | .tp_methods = tp_methods, 547 | }; 548 | 549 | static PyTypeObject object = { 550 | PyVarObject_HEAD_INIT(NULL, 0) 551 | .tp_name = "numpy_allocator.object", 552 | .tp_str = tp_str, 553 | .tp_flags = Py_TPFLAGS_BASETYPE | Py_TPFLAGS_DEFAULT, 554 | .tp_methods = tp_methods, 555 | }; 556 | 557 | static int exec_module(PyObject *module) { 558 | PyObject *list = PyList_New(0); 559 | if (!list) { 560 | return -1; 561 | } 562 | 563 | var = PyContextVar_New("var", list); 564 | Py_DECREF(list); 565 | if (!var) { 566 | return -1; 567 | } 568 | 569 | object.tp_base = &PyBaseObject_Type; 570 | object.tp_new = PyBaseObject_Type.tp_new; 571 | if (PyType_Ready(&object)) { 572 | Py_DECREF(var); 573 | 574 | return -1; 575 | } 576 | 577 | Py_INCREF(&object); 578 | 579 | if (PyModule_AddObject(module, "object", (PyObject *) &object)) { 580 | Py_DECREF(&object); 581 | 582 | Py_DECREF(var); 583 | 584 | return -1; 585 | } 586 | 587 | type.tp_base = &PyType_Type; 588 | type.tp_new = PyType_Type.tp_new; 589 | if (PyType_Ready(&type)) { 590 | Py_DECREF(&object); 591 | 592 | Py_DECREF(var); 593 | 594 | return -1; 595 | } 596 | 597 | Py_INCREF(&type); 598 | 599 | if (PyModule_AddObject(module, "type", (PyObject *) &type)) { 600 | Py_DECREF(&type); 601 | 602 | Py_DECREF(&object); 603 | 604 | Py_DECREF(var); 605 | 606 | return -1; 607 | } 608 | 609 | if (PyObject_SetAttrString(module, "default_handler", PyDataMem_DefaultHandler)) { 610 | Py_DECREF(&type); 611 | 612 | Py_DECREF(&object); 613 | 614 | Py_DECREF(var); 615 | 616 | return -1; 617 | } 618 | 619 | return 0; 620 | } 621 | 622 | static PyModuleDef_Slot m_slots[] = { 623 | {Py_mod_exec, exec_module}, 624 | {0, NULL}, 625 | }; 626 | 627 | static PyObject *set_handler(PyObject *module, PyObject *handler) { 628 | if (handler == Py_None) { 629 | return PyDataMem_SetHandler(NULL); 630 | } else { 631 | return PyDataMem_SetHandler(handler); 632 | } 633 | } 634 | 635 | static PyObject *get_handler(PyObject *module, PyObject *args) { 636 | PyObject *array = NULL; 637 | if (!PyArg_ParseTuple(args, "|O:get_handler", &array)) { 638 | return NULL; 639 | } 640 | 641 | if (array) { 642 | if (!PyArray_Check(array)) { 643 | PyErr_SetString(PyExc_TypeError, "if supplied, argument must be an ndarray"); 644 | return NULL; 645 | } 646 | 647 | while (array && PyArray_Check(array)) { 648 | if (PyArray_CHKFLAGS((PyArrayObject *) array, NPY_ARRAY_OWNDATA)) { 649 | PyObject *array_handler = PyArray_HANDLER((PyArrayObject *) array); 650 | if (!array_handler) { 651 | PyErr_SetString(PyExc_RuntimeError, "no memory handler found but OWNDATA flag set"); 652 | return NULL; 653 | } 654 | 655 | Py_INCREF(array_handler); 656 | return array_handler; 657 | } 658 | 659 | array = PyArray_BASE((PyArrayObject *) array); 660 | } 661 | 662 | Py_RETURN_NONE; 663 | } else { 664 | return PyDataMem_GetHandler(); 665 | } 666 | } 667 | 668 | static PyMethodDef m_methods[] = { 669 | {"get_handler", get_handler, METH_VARARGS, NULL}, 670 | {"set_handler", set_handler, METH_O, NULL}, 671 | {NULL, NULL, 0, NULL}, 672 | }; 673 | 674 | static PyModuleDef def = { 675 | PyModuleDef_HEAD_INIT, 676 | .m_name = "numpy_allocator", 677 | .m_methods = m_methods, 678 | .m_slots = m_slots, 679 | }; 680 | 681 | PyMODINIT_FUNC PyInit_numpy_allocator(void) { 682 | import_array(); 683 | 684 | return PyModuleDef_Init(&def); 685 | } 686 | --------------------------------------------------------------------------------