8f43927575
examples/ is based on KasperskyLab/knp at 55470ed6 with local changes to the mnist-learn altai/blifat models, cuda-backend CMake and model_loader. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
301 lines
8 KiB
Python
301 lines
8 KiB
Python
#!/usr/bin/env python3
|
|
"""
|
|
@file pybind11_gen.py
|
|
@brief Binding generator for PyBind11.
|
|
|
|
@license Apache 2.0 License.
|
|
@copyright © 2024 AO Kaspersky Lab
|
|
@date 03.10.2023.
|
|
|
|
This file has been modified by AO Kaspersky Lab in 10-2023.
|
|
Scope of modification: Changes to the grammar processing function.
|
|
"""
|
|
|
|
#
|
|
# Apache 2 license
|
|
#
|
|
|
|
import argparse
|
|
import collections
|
|
import re
|
|
|
|
from CppHeaderParser_l import CppHeaderParser
|
|
|
|
op_re = re.compile(r'^.*operator\s*\w*\s*(.+)')
|
|
operators = {
|
|
'()': '__call__',
|
|
'<': '__lt__',
|
|
'<=': '__le__',
|
|
'>': '__gt__',
|
|
'>=': '__ge__',
|
|
'==': '__eq__',
|
|
'!=': '__ne__',
|
|
'[]': '__getitem__',
|
|
'bool': 'bool',
|
|
}
|
|
|
|
spaces_count = 4
|
|
|
|
# CppHeaderParser.debug = 1
|
|
# CppHeaderParser.debug_trace = 1
|
|
|
|
|
|
MethodHookData = collections.namedtuple('MethodHookData', ['method_name', 'in_params', 'ret_names', 'pre', 'post'])
|
|
|
|
|
|
def _process_enum(enum, class_name=None):
|
|
name = enum['name']
|
|
typename = name
|
|
parent = 'm'
|
|
|
|
if class_name:
|
|
typename = '{}::{}'.format(class_name, typename)
|
|
parent = class_name.lower()
|
|
|
|
ret = ['py::enum_<{}>({}, "{}")'.format(typename, parent, name)]
|
|
|
|
for v in enum['values']:
|
|
k = v['name']
|
|
ret.append('{}.value("{}", {}::{})'.format(' ' * spaces_count, k, typename, k))
|
|
|
|
ret[-1] = ret[-1] + ';'
|
|
return ret
|
|
|
|
|
|
def _process_fn(fn):
|
|
return [] # TODO
|
|
|
|
|
|
def _process_method(class_name, meth, hooks, overloaded=False):
|
|
# Skip destructor.
|
|
if meth.get('destructor', False):
|
|
return []
|
|
|
|
ret = []
|
|
method_name: str = meth['name']
|
|
parameters = meth['parameters']
|
|
|
|
# Fix types that are enums.
|
|
for p in parameters:
|
|
if p.get('enum'):
|
|
p['raw_type'] = p['enum']
|
|
p['type'] = p['enum']
|
|
|
|
# Constructor.
|
|
if method_name == class_name:
|
|
params = ', '.join(p['raw_type'] for p in parameters)
|
|
ret.append('{}.def(py::init<{}>())'.format(' ' * spaces_count, params))
|
|
else:
|
|
pre = []
|
|
post = []
|
|
ret_names = []
|
|
if meth['returns'] != 'void':
|
|
ret_names.append('__ret')
|
|
|
|
in_params = parameters[:]
|
|
|
|
modified = False
|
|
|
|
# Data that hooks can modify.
|
|
hook_data = MethodHookData(method_name, in_params, ret_names, pre, post)
|
|
|
|
for hook in hooks.get('method_hooks', []):
|
|
if hook(class_name, meth, hook_data):
|
|
modified = True
|
|
|
|
py_method_name = hook_data.method_name
|
|
|
|
m = op_re.match(method_name)
|
|
if m:
|
|
py_method_name = operators.get(m[1], method_name)
|
|
|
|
if modified:
|
|
in_args = ''
|
|
if in_params:
|
|
in_args = ', ' + ', '.join('%(type)s %(name)s' % p for p in in_params)
|
|
|
|
ret.append(
|
|
f'{" " * spaces_count}.def("%(py_method_name)s", '
|
|
f'[](%(class_name)s '
|
|
'&__inst%(in_args)s)\n'
|
|
f'{" " * 2 * spaces_count}'
|
|
'{' % locals()
|
|
)
|
|
|
|
if pre:
|
|
ret.append(' ' * 2 * spaces_count + '; '.join(pre) + ';')
|
|
|
|
meth_params = ', '.join(p['name'] for p in meth['parameters'])
|
|
|
|
fnret = 'auto __ret = '
|
|
|
|
if '__ret' not in ret_names:
|
|
fnret = ''
|
|
|
|
if not post:
|
|
if ret_names == ['__ret']:
|
|
ret_names = []
|
|
fnret = 'return '
|
|
|
|
ret.append(f'{" " * spaces_count * 2}%(fnret)s__inst.%(method_name)s(%(meth_params)s);' % locals())
|
|
|
|
if post:
|
|
ret.append(' ' * spaces_count + '; '.join(post) + ';')
|
|
|
|
if len(ret_names) == 0:
|
|
pass
|
|
elif len(ret_names) == 1:
|
|
ret.append(f'{" " * 2 * spaces_count}return %s;' % ret_names[0])
|
|
else:
|
|
ret.append(f'{" " * 2 * spaces_count}return std::make_tuple(%s);' % ', '.join(ret_names))
|
|
|
|
ret.append('%s})' % (' ' * spaces_count))
|
|
|
|
else:
|
|
overload = ''
|
|
if overloaded:
|
|
# Overloaded method.
|
|
params = ', '.join(p['raw_type'] for p in parameters)
|
|
overload = '({} ({}::*)({}))'.format(meth['returns'], class_name, params)
|
|
|
|
ret.append(
|
|
'{}.def("{}", {}&{}::{})'.format(' ' * spaces_count, py_method_name, overload, class_name, method_name)
|
|
)
|
|
|
|
return ret
|
|
|
|
|
|
def _process_class(cls, hooks):
|
|
class_name = cls['name']
|
|
varname = class_name.lower()
|
|
|
|
ret = ['py::class_<{}> {}(m, "{}");'.format(class_name, varname, class_name)]
|
|
|
|
# Collect methods first to determine if there are overloads.
|
|
# We're ignoring base classes here.
|
|
methods = cls['methods']['public']
|
|
if methods:
|
|
ret.append(varname)
|
|
|
|
# Collapse them to find overloads.
|
|
meths = collections.OrderedDict()
|
|
for meth in methods:
|
|
meths.setdefault(meth['name'], []).append(meth)
|
|
|
|
# Process it.
|
|
for ml in meths.values():
|
|
if len(ml) == 1:
|
|
ret += _process_method(class_name, ml[0], hooks)
|
|
else:
|
|
for mh in ml:
|
|
ret += _process_method(class_name, mh, hooks, True)
|
|
|
|
ret[-1] = ret[-1] + ';'
|
|
|
|
for e in cls['enums']['public']:
|
|
ret.append('')
|
|
ret += _process_enum(e, class_name=class_name)
|
|
|
|
return ret
|
|
|
|
|
|
def process_header(fname, hooks):
|
|
"""Return a list of lines."""
|
|
|
|
header = CppHeaderParser.CppHeader(fname)
|
|
output = []
|
|
|
|
for e in header.enums:
|
|
output += _process_enum(e)
|
|
output.append('')
|
|
|
|
# for fn in header.functions:
|
|
# output += _process_fn(fn)
|
|
# output.append('')
|
|
|
|
for cls in sorted(header.classes.values(), key=lambda c: c['line_number']):
|
|
output += _process_class(cls, hooks)
|
|
output.append('')
|
|
|
|
return output
|
|
|
|
|
|
#
|
|
# Hooks.
|
|
#
|
|
|
|
# Method hook parameters:
|
|
# class_name: name of the class
|
|
# method: method dictionary from cppheaderparser
|
|
# in_params: copy of method['parameters']
|
|
# ret_names: variables to return
|
|
# pre: statements to insert before function call
|
|
# post: statements to insert after function call
|
|
# .. returns True if method hook did something
|
|
|
|
|
|
def _reference_hook(class_name, method, hook_data):
|
|
parameters = method['parameters']
|
|
refs = [p for p in parameters if p['reference']]
|
|
if refs:
|
|
hook_data.in_params[:] = [p for p in hook_data.in_params if not p['reference']]
|
|
hook_data.pre.extend('%(raw_type)s %(name)s' % p for p in refs)
|
|
hook_data.ret_names.extend(p['name'] for p in refs)
|
|
return True
|
|
|
|
|
|
def _ctr_hook(class_name, method, hook_data):
|
|
if method['returns'] == 'CTR_Code':
|
|
hook_data.ret_names.remove('__ret')
|
|
hook_data.post.append('CheckCTRCode(__ret)')
|
|
return True
|
|
|
|
|
|
def process_module(module_name, headers, hooks, add_namespaces=None):
|
|
for header in headers:
|
|
print('#include <%s>' % header) # TODO: Not usually the actual path.
|
|
|
|
print()
|
|
print('#include <pybind11/pybind11.h>')
|
|
print()
|
|
print('namespace py = pybind11;')
|
|
if add_namespaces:
|
|
for ns in add_namespaces:
|
|
if not ns:
|
|
continue
|
|
|
|
print('Using namespace %s;' % ns)
|
|
print()
|
|
|
|
print()
|
|
print('PYBIND11_MODULE(%s, m)\n{' % module_name)
|
|
print()
|
|
print(
|
|
'%sm.doc() = R"pbdoc(\n%s%s\n%s)pbdoc";'
|
|
% (' ' * spaces_count, ' ' * spaces_count * 2, module_name, ' ' * spaces_count)
|
|
)
|
|
print()
|
|
|
|
for header in headers:
|
|
print(' ' * spaces_count + f'\n{" " * spaces_count}'.join(process_header(header, hooks)))
|
|
print()
|
|
|
|
print('}')
|
|
|
|
|
|
def main():
|
|
parser = argparse.ArgumentParser()
|
|
parser.add_argument('module_name')
|
|
parser.add_argument('-n', '--add-namespaces', default='', help='Additional namespaces')
|
|
parser.add_argument('headers', nargs='+')
|
|
|
|
args = parser.parse_args()
|
|
|
|
hooks = dict()
|
|
hooks['method_hooks'] = [_reference_hook, _ctr_hook]
|
|
|
|
process_module(args.module_name, args.headers, hooks, args.add_namespaces.split(','))
|
|
|
|
|
|
if __name__ == '__main__':
|
|
main()
|