1import os
2import genapi
3
4import numpy_api
5
6from genapi import TypeApi, FunctionApi
7
8h_template = r"""
9#ifdef _UMATHMODULE
10
11extern NPY_NO_EXPORT PyTypeObject PyUFunc_Type;
12
13%s
14
15#else
16
17#if defined(PY_UFUNC_UNIQUE_SYMBOL)
18#define PyUFunc_API PY_UFUNC_UNIQUE_SYMBOL
19#endif
20
21#if defined(NO_IMPORT) || defined(NO_IMPORT_UFUNC)
22extern void **PyUFunc_API;
23#else
24#if defined(PY_UFUNC_UNIQUE_SYMBOL)
25void **PyUFunc_API;
26#else
27static void **PyUFunc_API=NULL;
28#endif
29#endif
30
31%s
32
33static NPY_INLINE int
34_import_umath(void)
35{
36  PyObject *numpy = PyImport_ImportModule("numpy.core._multiarray_umath");
37  PyObject *c_api = NULL;
38
39  if (numpy == NULL) {
40      PyErr_SetString(PyExc_ImportError,
41                      "numpy.core._multiarray_umath failed to import");
42      return -1;
43  }
44  c_api = PyObject_GetAttrString(numpy, "_UFUNC_API");
45  Py_DECREF(numpy);
46  if (c_api == NULL) {
47      PyErr_SetString(PyExc_AttributeError, "_UFUNC_API not found");
48      return -1;
49  }
50
51  if (!PyCapsule_CheckExact(c_api)) {
52      PyErr_SetString(PyExc_RuntimeError, "_UFUNC_API is not PyCapsule object");
53      Py_DECREF(c_api);
54      return -1;
55  }
56  PyUFunc_API = (void **)PyCapsule_GetPointer(c_api, NULL);
57  Py_DECREF(c_api);
58  if (PyUFunc_API == NULL) {
59      PyErr_SetString(PyExc_RuntimeError, "_UFUNC_API is NULL pointer");
60      return -1;
61  }
62  return 0;
63}
64
65#define import_umath() \
66    do {\
67        UFUNC_NOFPE\
68        if (_import_umath() < 0) {\
69            PyErr_Print();\
70            PyErr_SetString(PyExc_ImportError,\
71                    "numpy.core.umath failed to import");\
72            return NULL;\
73        }\
74    } while(0)
75
76#define import_umath1(ret) \
77    do {\
78        UFUNC_NOFPE\
79        if (_import_umath() < 0) {\
80            PyErr_Print();\
81            PyErr_SetString(PyExc_ImportError,\
82                    "numpy.core.umath failed to import");\
83            return ret;\
84        }\
85    } while(0)
86
87#define import_umath2(ret, msg) \
88    do {\
89        UFUNC_NOFPE\
90        if (_import_umath() < 0) {\
91            PyErr_Print();\
92            PyErr_SetString(PyExc_ImportError, msg);\
93            return ret;\
94        }\
95    } while(0)
96
97#define import_ufunc() \
98    do {\
99        UFUNC_NOFPE\
100        if (_import_umath() < 0) {\
101            PyErr_Print();\
102            PyErr_SetString(PyExc_ImportError,\
103                    "numpy.core.umath failed to import");\
104        }\
105    } while(0)
106
107#endif
108"""
109
110c_template = r"""
111/* These pointers will be stored in the C-object for use in other
112    extension modules
113*/
114
115void *PyUFunc_API[] = {
116%s
117};
118"""
119
120def generate_api(output_dir, force=False):
121    basename = 'ufunc_api'
122
123    h_file = os.path.join(output_dir, '__%s.h' % basename)
124    c_file = os.path.join(output_dir, '__%s.c' % basename)
125    d_file = os.path.join(output_dir, '%s.txt' % basename)
126    targets = (h_file, c_file, d_file)
127
128    sources = ['ufunc_api_order.txt']
129
130    if (not force and not genapi.should_rebuild(targets, sources + [__file__])):
131        return targets
132    else:
133        do_generate_api(targets, sources)
134
135    return targets
136
137def do_generate_api(targets, sources):
138    header_file = targets[0]
139    c_file = targets[1]
140    doc_file = targets[2]
141
142    ufunc_api_index = genapi.merge_api_dicts((
143            numpy_api.ufunc_funcs_api,
144            numpy_api.ufunc_types_api))
145    genapi.check_api_dict(ufunc_api_index)
146
147    ufunc_api_list = genapi.get_api_functions('UFUNC_API', numpy_api.ufunc_funcs_api)
148
149    # Create dict name -> *Api instance
150    ufunc_api_dict = {}
151    api_name = 'PyUFunc_API'
152    for f in ufunc_api_list:
153        name = f.name
154        index = ufunc_api_index[name][0]
155        annotations = ufunc_api_index[name][1:]
156        ufunc_api_dict[name] = FunctionApi(f.name, index, annotations,
157                                           f.return_type, f.args, api_name)
158
159    for name, val in numpy_api.ufunc_types_api.items():
160        index = val[0]
161        ufunc_api_dict[name] = TypeApi(name, index, 'PyTypeObject', api_name)
162
163    # set up object API
164    module_list = []
165    extension_list = []
166    init_list = []
167
168    for name, index in genapi.order_dict(ufunc_api_index):
169        api_item = ufunc_api_dict[name]
170        extension_list.append(api_item.define_from_array_api_string())
171        init_list.append(api_item.array_api_define())
172        module_list.append(api_item.internal_define())
173
174    # Write to header
175    s = h_template % ('\n'.join(module_list), '\n'.join(extension_list))
176    genapi.write_file(header_file, s)
177
178    # Write to c-code
179    s = c_template % ',\n'.join(init_list)
180    genapi.write_file(c_file, s)
181
182    # Write to documentation
183    s = '''
184=================
185NumPy Ufunc C-API
186=================
187'''
188    for func in ufunc_api_list:
189        s += func.to_ReST()
190        s += '\n\n'
191    genapi.write_file(doc_file, s)
192
193    return targets
194