-
Notifications
You must be signed in to change notification settings - Fork 153
/
nb_lib.h
529 lines (371 loc) · 20.8 KB
/
nb_lib.h
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
452
453
454
455
456
457
458
459
460
461
462
463
464
465
466
467
468
469
470
471
472
473
474
475
476
477
478
479
480
481
482
483
484
485
486
487
488
489
490
491
492
493
494
495
496
497
498
499
500
501
502
503
504
505
506
507
508
509
510
511
512
513
514
515
516
517
518
519
520
521
522
523
524
525
526
527
528
529
/*
nanobind/nb_lib.h: Interface to libnanobind.so
Copyright (c) 2022 Wenzel Jakob
All rights reserved. Use of this source code is governed by a
BSD-style license that can be found in the LICENSE file.
*/
NAMESPACE_BEGIN(NB_NAMESPACE)
// Forward declarations for types in dlpack.h (1)
namespace dlpack { struct dltensor; struct dtype; }
NAMESPACE_BEGIN(detail)
// Forward declarations for types in dlpack.h (2)
struct ndarray_handle;
struct ndarray_req;
/**
* Helper class to clean temporaries created by function dispatch.
* The first element serves a special role: it stores the 'self'
* object of method calls (for rv_policy::reference_internal).
*/
struct NB_CORE cleanup_list {
public:
static constexpr uint32_t Small = 6;
cleanup_list(PyObject *self) :
m_size{1},
m_capacity{Small},
m_data{m_local} {
m_local[0] = self;
}
~cleanup_list() = default;
/// Append a single PyObject to the cleanup stack
NB_INLINE void append(PyObject *value) noexcept {
if (m_size >= m_capacity)
expand();
m_data[m_size++] = value;
}
NB_INLINE PyObject *self() const {
return m_local[0];
}
/// Decrease the reference count of all appended objects
void release() noexcept;
/// Does the list contain any entries? (besides the 'self' argument)
bool used() { return m_size != 1; }
/// Return the size of the cleanup stack
size_t size() const { return m_size; }
/// Subscript operator
PyObject *operator[](size_t index) const { return m_data[index]; }
protected:
/// Out of memory, expand..
void expand() noexcept;
protected:
uint32_t m_size;
uint32_t m_capacity;
PyObject **m_data;
PyObject *m_local[Small];
};
// ========================================================================
/// Raise a runtime error with the given message
#if defined(__GNUC__)
__attribute__((noreturn, __format__ (__printf__, 1, 2)))
#else
[[noreturn]]
#endif
NB_CORE void raise(const char *fmt, ...);
/// Raise a type error with the given message
#if defined(__GNUC__)
__attribute__((noreturn, __format__ (__printf__, 1, 2)))
#else
[[noreturn]]
#endif
NB_CORE void raise_type_error(const char *fmt, ...);
/// Abort the process with a fatal error
#if defined(__GNUC__)
__attribute__((noreturn, __format__ (__printf__, 1, 2)))
#else
[[noreturn]]
#endif
NB_CORE void fail(const char *fmt, ...) noexcept;
/// Raise nanobind::python_error after an error condition was found
[[noreturn]] NB_CORE void raise_python_error();
/// Raise nanobind::next_overload
NB_CORE void raise_next_overload_if_null(void *p);
/// Raise nanobind::cast_error
NB_CORE void raise_cast_error();
// ========================================================================
NB_CORE void init(const char *domain);
// ========================================================================
/// Convert a Python object into a Python unicode string
NB_CORE PyObject *str_from_obj(PyObject *o);
/// Convert an UTF8 null-terminated C string into a Python unicode string
NB_CORE PyObject *str_from_cstr(const char *c);
/// Convert an UTF8 C string + size into a Python unicode string
NB_CORE PyObject *str_from_cstr_and_size(const char *c, size_t n);
// ========================================================================
/// Convert a Python object into a Python byte string
NB_CORE PyObject *bytes_from_obj(PyObject *o);
/// Convert an UTF8 null-terminated C string into a Python byte string
NB_CORE PyObject *bytes_from_cstr(const char *c);
/// Convert an UTF8 C string + size into a Python byte string
NB_CORE PyObject *bytes_from_cstr_and_size(const char *c, size_t n);
// ========================================================================
/// Convert a Python object into a Python boolean object
NB_CORE PyObject *bool_from_obj(PyObject *o);
/// Convert a Python object into a Python integer object
NB_CORE PyObject *int_from_obj(PyObject *o);
/// Convert a Python object into a Python floating point object
NB_CORE PyObject *float_from_obj(PyObject *o);
// ========================================================================
/// Convert a Python object into a Python list
NB_CORE PyObject *list_from_obj(PyObject *o);
/// Convert a Python object into a Python tuple
NB_CORE PyObject *tuple_from_obj(PyObject *o);
// ========================================================================
/// Get an object attribute or raise an exception
NB_CORE PyObject *getattr(PyObject *obj, const char *key);
NB_CORE PyObject *getattr(PyObject *obj, PyObject *key);
/// Get an object attribute or return a default value (never raises)
NB_CORE PyObject *getattr(PyObject *obj, const char *key, PyObject *def) noexcept;
NB_CORE PyObject *getattr(PyObject *obj, PyObject *key, PyObject *def) noexcept;
/// Get an object attribute or raise an exception. Skip if 'out' is non-null
NB_CORE void getattr_or_raise(PyObject *obj, const char *key, PyObject **out);
NB_CORE void getattr_or_raise(PyObject *obj, PyObject *key, PyObject **out);
/// Set an object attribute or raise an exception
NB_CORE void setattr(PyObject *obj, const char *key, PyObject *value);
NB_CORE void setattr(PyObject *obj, PyObject *key, PyObject *value);
/// Delete an object attribute or raise an exception
NB_CORE void delattr(PyObject *obj, const char *key);
NB_CORE void delattr(PyObject *obj, PyObject *key);
// ========================================================================
/// Index into an object or raise an exception. Skip if 'out' is non-null
NB_CORE void getitem_or_raise(PyObject *obj, Py_ssize_t, PyObject **out);
NB_CORE void getitem_or_raise(PyObject *obj, const char *key, PyObject **out);
NB_CORE void getitem_or_raise(PyObject *obj, PyObject *key, PyObject **out);
/// Set an item or raise an exception
NB_CORE void setitem(PyObject *obj, Py_ssize_t, PyObject *value);
NB_CORE void setitem(PyObject *obj, const char *key, PyObject *value);
NB_CORE void setitem(PyObject *obj, PyObject *key, PyObject *value);
/// Delete an item or raise an exception
NB_CORE void delitem(PyObject *obj, Py_ssize_t);
NB_CORE void delitem(PyObject *obj, const char *key);
NB_CORE void delitem(PyObject *obj, PyObject *key);
// ========================================================================
/// Determine the length of a Python object
NB_CORE size_t obj_len(PyObject *o);
/// Try to roughly determine the length of a Python object
NB_CORE size_t obj_len_hint(PyObject *o) noexcept;
/// Obtain a string representation of a Python object
NB_CORE PyObject* obj_repr(PyObject *o);
/// Perform a comparison between Python objects and handle errors
NB_CORE bool obj_comp(PyObject *a, PyObject *b, int value);
/// Perform an unary operation on a Python object with error handling
NB_CORE PyObject *obj_op_1(PyObject *a, PyObject* (*op)(PyObject*));
/// Perform an unary operation on a Python object with error handling
NB_CORE PyObject *obj_op_2(PyObject *a, PyObject *b,
PyObject *(*op)(PyObject *, PyObject *));
// Perform a vector function call
NB_CORE PyObject *obj_vectorcall(PyObject *base, PyObject *const *args,
size_t nargsf, PyObject *kwnames,
bool method_call);
/// Create an iterator from 'o', raise an exception in case of errors
NB_CORE PyObject *obj_iter(PyObject *o);
/// Advance the iterator 'o', raise an exception in case of errors
NB_CORE PyObject *obj_iter_next(PyObject *o);
// ========================================================================
// Conversion validity check done by nb::make_tuple
NB_CORE void tuple_check(PyObject *tuple, size_t nargs);
// ========================================================================
// Append a single argument to a function call
NB_CORE void call_append_arg(PyObject *args, size_t &nargs, PyObject *value);
// Append a variable-length sequence of arguments to a function call
NB_CORE void call_append_args(PyObject *args, size_t &nargs, PyObject *value);
// Append a single keyword argument to a function call
NB_CORE void call_append_kwarg(PyObject *kwargs, const char *name, PyObject *value);
// Append a variable-length dictionary of keyword arguments to a function call
NB_CORE void call_append_kwargs(PyObject *kwargs, PyObject *value);
// ========================================================================
// If the given sequence has the size 'size', return a pointer to its contents.
// May produce a temporary.
NB_CORE PyObject **seq_get_with_size(PyObject *seq, size_t size,
PyObject **temp) noexcept;
// Like the above, but return the size instead of checking it.
NB_CORE PyObject **seq_get(PyObject *seq, size_t *size,
PyObject **temp) noexcept;
// ========================================================================
/// Create a new capsule object with a name
NB_CORE PyObject *capsule_new(const void *ptr, const char *name,
void (*cleanup)(void *) noexcept) noexcept;
// ========================================================================
/// Create a Python function object for the given function record
NB_CORE PyObject *nb_func_new(const void *data) noexcept;
// ========================================================================
/// Create a Python type object for the given type record
struct type_init_data;
NB_CORE PyObject *nb_type_new(const type_init_data *c) noexcept;
/// Extract a pointer to a C++ type underlying a Python object, if possible
NB_CORE bool nb_type_get(const std::type_info *t, PyObject *o, uint8_t flags,
cleanup_list *cleanup, void **out) noexcept;
/// Cast a C++ type instance into a Python object
NB_CORE PyObject *nb_type_put(const std::type_info *cpp_type, void *value,
rv_policy rvp, cleanup_list *cleanup,
bool *is_new = nullptr) noexcept;
// Special version of nb_type_put for polymorphic classes
NB_CORE PyObject *nb_type_put_p(const std::type_info *cpp_type,
const std::type_info *cpp_type_p, void *value,
rv_policy rvp, cleanup_list *cleanup,
bool *is_new = nullptr) noexcept;
// Special version of 'nb_type_put' for unique pointers and ownership transfer
NB_CORE PyObject *nb_type_put_unique(const std::type_info *cpp_type,
void *value, cleanup_list *cleanup,
bool cpp_delete) noexcept;
// Special version of 'nb_type_put_unique' for polymorphic classes
NB_CORE PyObject *nb_type_put_unique_p(const std::type_info *cpp_type,
const std::type_info *cpp_type_p,
void *value, cleanup_list *cleanup,
bool cpp_delete) noexcept;
/// Try to reliquish ownership from Python object to a unique_ptr
NB_CORE void nb_type_relinquish_ownership(PyObject *o, bool cpp_delete);
/// Get a pointer to a user-defined 'extra' value associated with the nb_type t.
NB_CORE void *nb_type_supplement(PyObject *t) noexcept;
/// Check if the given python object represents a nanobind type
NB_CORE bool nb_type_check(PyObject *t) noexcept;
/// Return the size of the type wrapped by the given nanobind type object
NB_CORE size_t nb_type_size(PyObject *t) noexcept;
/// Return the alignment of the type wrapped by the given nanobind type object
NB_CORE size_t nb_type_align(PyObject *t) noexcept;
/// Return a unicode string representing the long-form name of the given type
NB_CORE PyObject *nb_type_name(PyObject *t) noexcept;
/// Return a unicode string representing the long-form name of object's type
NB_CORE PyObject *nb_inst_name(PyObject *o) noexcept;
/// Return the C++ type_info wrapped by the given nanobind type object
NB_CORE const std::type_info *nb_type_info(PyObject *t) noexcept;
/// Get a pointer to the instance data of a nanobind instance (nb_inst)
NB_CORE void *nb_inst_ptr(PyObject *o) noexcept;
/// Check if a Python type object wraps an instance of a specific C++ type
NB_CORE bool nb_type_isinstance(PyObject *obj, const std::type_info *t) noexcept;
/// Search for the Python type object associated with a C++ type
NB_CORE PyObject *nb_type_lookup(const std::type_info *t) noexcept;
/// Allocate an instance of type 't'
NB_CORE PyObject *nb_inst_alloc(PyTypeObject *t);
/// Allocate an zero-initialized instance of type 't'
NB_CORE PyObject *nb_inst_alloc_zero(PyTypeObject *t);
/// Allocate an instance of type 't' referencing the existing 'ptr'
NB_CORE PyObject *nb_inst_reference(PyTypeObject *t, void *ptr,
PyObject *parent);
/// Allocate an instance of type 't' taking ownership of the existing 'ptr'
NB_CORE PyObject *nb_inst_take_ownership(PyTypeObject *t, void *ptr);
/// Call the destructor of the given python object
NB_CORE void nb_inst_destruct(PyObject *o) noexcept;
/// Zero-initialize a POD type and mark it as ready + to be destructed upon GC
NB_CORE void nb_inst_zero(PyObject *o) noexcept;
/// Copy-construct 'dst' from 'src', mark it as ready and to be destructed (must have the same nb_type)
NB_CORE void nb_inst_copy(PyObject *dst, const PyObject *src) noexcept;
/// Move-construct 'dst' from 'src', mark it as ready and to be destructed (must have the same nb_type)
NB_CORE void nb_inst_move(PyObject *dst, const PyObject *src) noexcept;
/// Destruct 'dst', copy-construct 'dst' from 'src', mark ready and retain 'destruct' status (must have the same nb_type)
NB_CORE void nb_inst_replace_copy(PyObject *dst, const PyObject *src) noexcept;
/// Destruct 'dst', move-construct 'dst' from 'src', mark ready and retain 'destruct' status (must have the same nb_type)
NB_CORE void nb_inst_replace_move(PyObject *dst, const PyObject *src) noexcept;
/// Check if a particular instance uses a Python-derived type
NB_CORE bool nb_inst_python_derived(PyObject *o) noexcept;
/// Overwrite the instance's ready/destruct flags
NB_CORE void nb_inst_set_state(PyObject *o, bool ready, bool destruct) noexcept;
/// Query the 'ready' and 'destruct' flags of an instance
NB_CORE std::pair<bool, bool> nb_inst_state(PyObject *o) noexcept;
// ========================================================================
// Create and install a Python property object
NB_CORE void property_install(PyObject *scope, const char *name,
PyObject *getter, PyObject *setter) noexcept;
NB_CORE void property_install_static(PyObject *scope, const char *name,
PyObject *getter,
PyObject *setter) noexcept;
// ========================================================================
NB_CORE PyObject *get_override(void *ptr, const std::type_info *type,
const char *name, bool pure);
// ========================================================================
// Ensure that 'patient' cannot be GCed while 'nurse' is alive
NB_CORE void keep_alive(PyObject *nurse, PyObject *patient);
// Keep 'payload' alive until 'nurse' is GCed
NB_CORE void keep_alive(PyObject *nurse, void *payload,
void (*deleter)(void *) noexcept) noexcept;
// ========================================================================
/// Indicate to nanobind that an implicit constructor can convert 'src' -> 'dst'
NB_CORE void implicitly_convertible(const std::type_info *src,
const std::type_info *dst) noexcept;
/// Register a callback to check if implicit conversion to 'dst' is possible
NB_CORE void implicitly_convertible(bool (*predicate)(PyTypeObject *,
PyObject *,
cleanup_list *),
const std::type_info *dst) noexcept;
// ========================================================================
/// Fill in slots for an enum type being built
NB_CORE void nb_enum_prepare(const type_init_data *t,
PyType_Slot *&slots, size_t max_slots) noexcept;
/// Add an entry to an enumeration
NB_CORE void nb_enum_put(PyObject *type, const char *name, const void *value,
const char *doc) noexcept;
/// Export enum entries to the parent scope
NB_CORE void nb_enum_export(PyObject *type);
// ========================================================================
/// Try to import a Python extension module, raises an exception upon failure
NB_CORE PyObject *module_import(const char *name);
/// Create a new extension module with the given name
NB_CORE PyObject *module_new(const char *name, PyModuleDef *def) noexcept;
/// Create a submodule of an existing module
NB_CORE PyObject *module_new_submodule(PyObject *base, const char *name,
const char *doc) noexcept;
// ========================================================================
// Try to import a reference-counted ndarray object via DLPack
NB_CORE ndarray_handle *ndarray_import(PyObject *o, const ndarray_req *req,
bool convert,
cleanup_list *cleanup) noexcept;
// Describe a local ndarray object using a DLPack capsule
NB_CORE ndarray_handle *ndarray_create(void *value, size_t ndim,
const size_t *shape, PyObject *owner,
const int64_t *strides,
dlpack::dtype *dtype, bool ro,
int32_t device, int32_t device_id);
/// Increase the reference count of the given ndarray object; returns a pointer
/// to the underlying DLTensor
NB_CORE dlpack::dltensor *ndarray_inc_ref(ndarray_handle *) noexcept;
/// Decrease the reference count of the given ndarray object
NB_CORE void ndarray_dec_ref(ndarray_handle *) noexcept;
/// Wrap a ndarray_handle* into a PyCapsule
NB_CORE PyObject *ndarray_wrap(ndarray_handle *, int framework,
rv_policy policy, cleanup_list *cleanup) noexcept;
/// Check if an object is a known ndarray type (NumPy, PyTorch, Tensorflow, JAX)
NB_CORE bool ndarray_check(PyObject *o) noexcept;
// ========================================================================
/// Print to stdout using Python
NB_CORE void print(PyObject *file, PyObject *str, PyObject *end);
// ========================================================================
typedef void (*exception_translator)(const std::exception_ptr &, void *);
NB_CORE void register_exception_translator(exception_translator translator,
void *payload);
NB_CORE PyObject *exception_new(PyObject *mod, const char *name,
PyObject *base);
// ========================================================================
NB_CORE bool load_i8 (PyObject *o, uint8_t flags, int8_t *out) noexcept;
NB_CORE bool load_u8 (PyObject *o, uint8_t flags, uint8_t *out) noexcept;
NB_CORE bool load_i16(PyObject *o, uint8_t flags, int16_t *out) noexcept;
NB_CORE bool load_u16(PyObject *o, uint8_t flags, uint16_t *out) noexcept;
NB_CORE bool load_i32(PyObject *o, uint8_t flags, int32_t *out) noexcept;
NB_CORE bool load_u32(PyObject *o, uint8_t flags, uint32_t *out) noexcept;
NB_CORE bool load_i64(PyObject *o, uint8_t flags, int64_t *out) noexcept;
NB_CORE bool load_u64(PyObject *o, uint8_t flags, uint64_t *out) noexcept;
NB_CORE bool load_f32(PyObject *o, uint8_t flags, float *out) noexcept;
NB_CORE bool load_f64(PyObject *o, uint8_t flags, double *out) noexcept;
// ========================================================================
/// Increase the reference count of 'o', and check that the GIL is held
NB_CORE void incref_checked(PyObject *o) noexcept;
/// Decrease the reference count of 'o', and check that the GIL is held
NB_CORE void decref_checked(PyObject *o) noexcept;
// ========================================================================
NB_CORE void set_leak_warnings(bool value) noexcept;
NB_CORE void set_implicit_cast_warnings(bool value) noexcept;
// ========================================================================
NB_CORE bool iterable_check(PyObject *o) noexcept;
// ========================================================================
NB_CORE void slice_compute(PyObject *slice, Py_ssize_t size,
Py_ssize_t &start, Py_ssize_t &stop,
Py_ssize_t &step, size_t &slice_length);
// ========================================================================
NB_CORE PyObject *repr_list(PyObject *o);
NB_CORE PyObject *repr_map(PyObject *o);
NB_CORE bool is_alive() noexcept;
#if NB_TYPE_GET_SLOT_IMPL
NB_CORE void *type_get_slot(PyTypeObject *t, int slot_id);
#endif
NAMESPACE_END(detail)
using detail::raise;
using detail::raise_type_error;
using detail::raise_python_error;
NAMESPACE_END(NB_NAMESPACE)