mirror of
https://github.com/JonathanSalwan/Triton
synced 2026-08-17 10:23:04 -04:00
fix + improve python type stubs
This commit is contained in:
parent
cd7da95558
commit
b7c065b551
5 changed files with 126 additions and 72 deletions
|
|
@ -1,49 +1,61 @@
|
|||
try:
|
||||
from typing import Dict, List
|
||||
from typing import Dict, List, Set
|
||||
except ImportError:
|
||||
pass
|
||||
|
||||
from collections import OrderedDict
|
||||
|
||||
template = '''
|
||||
def {fname}({arg_names}):
|
||||
# type: ({arg_types}) -> {ret_type}
|
||||
"""{doc_str}"""
|
||||
{code}
|
||||
pass
|
||||
'''
|
||||
|
||||
class Function():
|
||||
generate_code = True
|
||||
INDENT_STR = ' ' * 4
|
||||
DEBUG = False
|
||||
|
||||
|
||||
def __init__(self, name, args, ret_type, doc_str=''):
|
||||
# type: (str, Dict[str, str], str, str) -> None
|
||||
self.name = name
|
||||
self.args = OrderedDict(args)
|
||||
self.ret_type = ret_type
|
||||
self.doc_str = doc_str
|
||||
self.name = name
|
||||
self.arg_types = [{t} for t in args.values()]
|
||||
self.arg_names = [{n} for n in args.keys()]
|
||||
self.ret_type = ret_type
|
||||
self.doc_str = doc_str
|
||||
self.overloaded_args = []
|
||||
|
||||
def add_overload(self, other):
|
||||
assert(self.name == other.name)
|
||||
|
||||
def gen_code(self):
|
||||
# type: () -> str
|
||||
code = str()
|
||||
func_call = '{}return self.org.{name}({args})'.format(self.INDENT_STR * 2, name=self.name, args=', '.join(self.args.keys()))
|
||||
if self.DEBUG:
|
||||
code += "{}print({})\n".format(self.INDENT_STR * 2, repr(func_call.strip()))
|
||||
code += func_call
|
||||
return code
|
||||
# we have fewer args than the other overload and need to expand them
|
||||
if len(self.arg_names) < len(other.arg_names):
|
||||
for i in range(len(self.arg_names), len(other.arg_names)):
|
||||
self.arg_names.append(set())
|
||||
self.arg_types.append(set())
|
||||
|
||||
# create unions of the names and types
|
||||
for i in range(len(other.arg_names)):
|
||||
self.arg_names[i] |= other.arg_names[i]
|
||||
self.arg_types[i] |= other.arg_types[i]
|
||||
|
||||
def __str__(self):
|
||||
# type: () -> str
|
||||
|
||||
def handle_default_initializer(full_names):
|
||||
# type: (Set[str]) -> List[str]
|
||||
names = list(full_names)
|
||||
for i in range(len(names) - 1):
|
||||
# remove the default parameter for every arg except the last one
|
||||
if '=' in names[i]:
|
||||
names[i] = names[i][:names[i].find('=')]
|
||||
|
||||
return names
|
||||
|
||||
arg_names = ', '.join(['self'] + ['_'.join(handle_default_initializer(names)) for names in self.arg_names])
|
||||
arg_types = 'Self'
|
||||
if len(self.arg_types) != 1 or len(self.arg_types[0]) > 1 or list(self.arg_types[0])[0] != 'None':
|
||||
arg_types = ', '.join(['Self'] + [' | '.join(types) for types in self.arg_types])
|
||||
func_str = template.format(
|
||||
fname=self.name,
|
||||
ret_type=self.ret_type,
|
||||
arg_names=', '.join(['self'] + list(self.args.keys())),
|
||||
arg_types=', '.join('' if val == 'None' else val for val in self.args.values()),
|
||||
doc_str=self.doc_str,
|
||||
code='{}pass'.format(self.INDENT_STR * 2) if not self.generate_code else self.gen_code()
|
||||
)
|
||||
arg_names=arg_names,
|
||||
arg_types=arg_types,
|
||||
doc_str=self.doc_str)
|
||||
return func_str
|
||||
|
|
|
|||
|
|
@ -1,5 +1,5 @@
|
|||
try:
|
||||
from typing import Optional, List, Dict, Tuple
|
||||
from typing import Optional, List, Dict, Tuple, Self
|
||||
except ImportError:
|
||||
pass
|
||||
|
||||
|
|
@ -17,7 +17,7 @@ NAMESPACE_PREFIX = 'n'
|
|||
list_pattern = r'\[(.*?)(?:,(?: ?...)?)?\]'
|
||||
type_pattern = r'(?P<type>List\[.*?\]|[\w\.]+)'
|
||||
obj_doc_re = re.compile(r'-\s<b>(?P<sig>.*?)<\/b><br>\r?\n(?P<desc>.*?)\r?\n\r?\n', flags=re.DOTALL)
|
||||
name_doc_pattern = r'- \*\*{namespace}\.(?P<member>.*?)\*\*'
|
||||
name_doc_pattern = r'- \*\*(?P<member>.*?)\*\*'
|
||||
ref_re = re.compile(r'\\ref\spy_(.*?)_page')
|
||||
sig_re = re.compile(r'(?P<return>{}) (?P<name>\w+)\s?\((?P<args>.*?)\)'.format(type_pattern))
|
||||
list_re = re.compile(list_pattern)
|
||||
|
|
@ -107,7 +107,7 @@ def gen_module_for_object(classname, input_str):
|
|||
|
||||
# find functions
|
||||
matches = obj_doc_re.finditer(input_str)
|
||||
funcs = []
|
||||
funcs: Dict[str, Function] = {}
|
||||
if not matches:
|
||||
return ""
|
||||
|
||||
|
|
@ -117,60 +117,103 @@ def gen_module_for_object(classname, input_str):
|
|||
# print("Signature: {}\nDescription: {}\n".format(fsig, desc))
|
||||
func = gen_function(fsig, desc)
|
||||
if func:
|
||||
funcs.append(func)
|
||||
# overloaded function name
|
||||
if func.name in funcs:
|
||||
funcs[func.name].add_overload(func)
|
||||
funcs[func.name].doc_str += '\n\n' + desc
|
||||
# unique function name
|
||||
else:
|
||||
funcs[func.name] = func
|
||||
else:
|
||||
print('... in module {}'.format(classname))
|
||||
|
||||
# special case for TritonContext.registers since is it dynamically
|
||||
# set based on the architecture of the context
|
||||
if classname == 'TritonContext':
|
||||
additional_init = 'self.registers = None # type: Any'
|
||||
else:
|
||||
additional_init = 'pass'
|
||||
|
||||
# generate
|
||||
autogen_str = '''
|
||||
class {classname}:
|
||||
def __init__(self, *args, **kargs):
|
||||
self.org = triton.{classname}(*args, **kargs)
|
||||
{additional_init}
|
||||
|
||||
{functions}
|
||||
'''.format(classname=classname, functions='\n'.join([str(f) for f in funcs]))
|
||||
{functions}
|
||||
'''.format(classname=classname,
|
||||
functions='\n'.join([str(f) for f in funcs.values()]),
|
||||
additional_init=additional_init)
|
||||
|
||||
return autogen_str
|
||||
|
||||
|
||||
class Submodule:
|
||||
|
||||
def __init__(self, name, indentation_level):
|
||||
# type: (Self, str, int) -> None
|
||||
self.name = name
|
||||
self.indentation_level = indentation_level
|
||||
self.members = []
|
||||
|
||||
def __str__(self):
|
||||
# type: () -> str
|
||||
indent = ' ' * self.indentation_level
|
||||
s = indent + 'class {}(IntEnum):\n'.format(self.name)
|
||||
|
||||
# add one indentation level for the members
|
||||
indent += ' '
|
||||
if len(self.members) == 0:
|
||||
s += indent + 'pass'
|
||||
else:
|
||||
s += '\n'.join((indent + member for member in self.members))
|
||||
|
||||
return s
|
||||
|
||||
def add_member(self, member: str):
|
||||
self.members.append(member)
|
||||
|
||||
|
||||
def gen_module_for_namespace(classname, input_str):
|
||||
# type: (str, str) -> str
|
||||
global args
|
||||
|
||||
input_str = ref_re.sub(sub_ref, input_str)
|
||||
|
||||
# find functions
|
||||
pat = name_doc_pattern.format(namespace=classname)
|
||||
pat = name_doc_pattern
|
||||
matches = re.finditer(pat, input_str)
|
||||
members = []
|
||||
if not matches:
|
||||
return ""
|
||||
|
||||
submodules = set()
|
||||
submodules = {}
|
||||
counter = 1
|
||||
for match in matches:
|
||||
member = ' {member} = triton.{namespace}.{member}'.format(member = match.group('member'), namespace=classname)
|
||||
components = match.group('member').split('.')
|
||||
modules = components[:-1]
|
||||
member = components[-1]
|
||||
submodule_path = '.'.join(modules)
|
||||
|
||||
if(member == ' Z3 = triton.SOLVER.Z3' and not args.z3_enabled):
|
||||
continue
|
||||
elif(member == ' BITWUZLA = triton.SOLVER.BITWUZLA' and not args.bitwuzla_enabled):
|
||||
continue
|
||||
for i in range(len(modules)):
|
||||
submodule = '.'.join(modules[:i+1])
|
||||
if submodule not in submodules:
|
||||
submodules[submodule] = Submodule(modules[i], i)
|
||||
|
||||
members.append(member)
|
||||
submod = member.split('=')[0].split('.')[:-1]
|
||||
for x in submod:
|
||||
submodules.add(' class %s: pass' % (x.lstrip()))
|
||||
if classname == 'SOLVER':
|
||||
if member == 'Z3' and not args.z3_enabled:
|
||||
print('skipping Z3')
|
||||
continue
|
||||
elif member == 'BITWUZLA' and not args.bitwuzla_enabled:
|
||||
print('skipping BITWUZLA')
|
||||
continue
|
||||
|
||||
#print(submodules)
|
||||
if not members:
|
||||
print("warning: empty namespace {}".format(classname))
|
||||
members.append(' pass')
|
||||
# don't bother with a type for this since it's usually just enums
|
||||
member_str = '{member} = {counter}'.format(
|
||||
member=member, counter=counter)
|
||||
submodules[submodule_path].add_member(member_str)
|
||||
counter += 1
|
||||
|
||||
# generate
|
||||
autogen_str = '''
|
||||
class {classname}:
|
||||
{submodules}
|
||||
{members}
|
||||
'''.format(classname=classname, submodules='\n'.join(submodules), members='\n'.join(members))
|
||||
autogen_str = '\n'.join((str(s) for s in submodules.values()))
|
||||
|
||||
return autogen_str
|
||||
|
||||
|
|
@ -276,9 +319,12 @@ def get_namespaces(namespace_dir):
|
|||
def gen_init_file(modules):
|
||||
# type: (List[str]) -> str
|
||||
global args
|
||||
mod_str = """from typing import List, Union, Callable, Tuple
|
||||
mod_str = """from typing import List, Union, Callable, Tuple, Any
|
||||
from typing_extensions import Self
|
||||
from enum import IntEnum
|
||||
import triton
|
||||
{z3}
|
||||
|
||||
{modules}
|
||||
|
||||
raise ImportError
|
||||
|
|
|
|||
|
|
@ -2188,7 +2188,6 @@ According to the CPU architecture, the OPCODE namespace contains all kinds of op
|
|||
- **OPCODE.RV64.ADDW**<br>
|
||||
- **OPCODE.RV64.AND**<br>
|
||||
- **OPCODE.RV64.ANDI**<br>
|
||||
- **OPCODE.RV64.AND**<br>
|
||||
- **OPCODE.RV64.AUIPC**<br>
|
||||
- **OPCODE.RV64.BEQ**<br>
|
||||
- **OPCODE.RV64.BGE**<br>
|
||||
|
|
@ -2284,7 +2283,6 @@ According to the CPU architecture, the OPCODE namespace contains all kinds of op
|
|||
- **OPCODE.RV32.ADDI**<br>
|
||||
- **OPCODE.RV32.AND**<br>
|
||||
- **OPCODE.RV32.ANDI**<br>
|
||||
- **OPCODE.RV32.AND**<br>
|
||||
- **OPCODE.RV32.AUIPC**<br>
|
||||
- **OPCODE.RV32.BEQ**<br>
|
||||
- **OPCODE.RV32.BGE**<br>
|
||||
|
|
@ -4535,7 +4533,6 @@ namespace triton {
|
|||
xPyDict_SetItemString(riscv64OpcodesDict, "ADDW", PyLong_FromUint32(triton::arch::riscv::ID_INS_ADDW));
|
||||
xPyDict_SetItemString(riscv64OpcodesDict, "AND", PyLong_FromUint32(triton::arch::riscv::ID_INS_AND));
|
||||
xPyDict_SetItemString(riscv64OpcodesDict, "ANDI", PyLong_FromUint32(triton::arch::riscv::ID_INS_ANDI));
|
||||
xPyDict_SetItemString(riscv64OpcodesDict, "AND", PyLong_FromUint32(triton::arch::riscv::ID_INS_AND));
|
||||
xPyDict_SetItemString(riscv64OpcodesDict, "AUIPC", PyLong_FromUint32(triton::arch::riscv::ID_INS_AUIPC));
|
||||
xPyDict_SetItemString(riscv64OpcodesDict, "BEQ", PyLong_FromUint32(triton::arch::riscv::ID_INS_BEQ));
|
||||
xPyDict_SetItemString(riscv64OpcodesDict, "BGE", PyLong_FromUint32(triton::arch::riscv::ID_INS_BGE));
|
||||
|
|
@ -4636,7 +4633,6 @@ namespace triton {
|
|||
xPyDict_SetItemString(riscv32OpcodesDict, "ADDI", PyLong_FromUint32(triton::arch::riscv::ID_INS_ADDI));
|
||||
xPyDict_SetItemString(riscv32OpcodesDict, "AND", PyLong_FromUint32(triton::arch::riscv::ID_INS_AND));
|
||||
xPyDict_SetItemString(riscv32OpcodesDict, "ANDI", PyLong_FromUint32(triton::arch::riscv::ID_INS_ANDI));
|
||||
xPyDict_SetItemString(riscv32OpcodesDict, "AND", PyLong_FromUint32(triton::arch::riscv::ID_INS_AND));
|
||||
xPyDict_SetItemString(riscv32OpcodesDict, "AUIPC", PyLong_FromUint32(triton::arch::riscv::ID_INS_AUIPC));
|
||||
xPyDict_SetItemString(riscv32OpcodesDict, "BEQ", PyLong_FromUint32(triton::arch::riscv::ID_INS_BEQ));
|
||||
xPyDict_SetItemString(riscv32OpcodesDict, "BGE", PyLong_FromUint32(triton::arch::riscv::ID_INS_BGE));
|
||||
|
|
|
|||
|
|
@ -28,28 +28,28 @@ The STUBS namespace contains all stubs.
|
|||
\section STUBS_py_api Python API - Items of the STUBS namespace
|
||||
<hr>
|
||||
|
||||
- <b>STUBS.AARCH64.LIBC.code</b><br>
|
||||
- **STUBS.AARCH64.LIBC.code**<br>
|
||||
The libc stub on AArch64 architecture with the ARM64 ABI calling convention.
|
||||
|
||||
- <b>STUBS.AARCH64.LIBC.symbols</b><br>
|
||||
- **STUBS.AARCH64.LIBC.symbols**<br>
|
||||
The symbols map of the AArch64 libc stub.
|
||||
|
||||
- <b>STUBS.I386.SYSTEMV.LIBC.code</b><br>
|
||||
- **STUBS.I386.SYSTEMV.LIBC.code**<br>
|
||||
The libc stub on i386 architecture with the SystemV ABI calling convention.
|
||||
|
||||
- <b>STUBS.I386.SYSTEMV.LIBC.symbols</b><br>
|
||||
- **STUBS.I386.SYSTEMV.LIBC.symbols**<br>
|
||||
The symbols map of the i386-systemv libc stub.
|
||||
|
||||
- <b>STUBS.X8664.MS.LIBC.code</b><br>
|
||||
- **STUBS.X8664.MS.LIBC.code**<br>
|
||||
The libc stub on x8664 architecture with the MS ABI calling convention.
|
||||
|
||||
- <b>STUBS.X8664.MS.LIBC.symbols</b><br>
|
||||
- **STUBS.X8664.MS.LIBC.symbols**<br>
|
||||
The symbols map of the x8664-ms libc stub.
|
||||
|
||||
- <b>STUBS.X8664.SYSTEMV.LIBC.code</b><br>
|
||||
- **STUBS.X8664.SYSTEMV.LIBC.code**<br>
|
||||
The libc stub on x8664 architecture with the SystemV ABI calling convention.
|
||||
|
||||
- <b>STUBS.X8664.SYSTEMV.LIBC.symbols</b><br>
|
||||
- **STUBS.X8664.SYSTEMV.LIBC.symbols**<br>
|
||||
The symbols map of the x8664-systemv libc stub.
|
||||
|
||||
*/
|
||||
|
|
|
|||
|
|
@ -295,10 +295,10 @@ Lifts a symbolic expression and all its references to Python format. If `icommen
|
|||
- <b>string liftToSMT(\ref py_SymbolicExpression_page expr, bool assert_=False, bool icomment=False)</b><br>
|
||||
Lifts a symbolic expression and all its references to SMT format. If `assert_` is true, then (assert <expr>). If `icomment` is true, then print instructions assembly in expression comments.
|
||||
|
||||
- <b>\ref py_SymbolicExpression_page newSymbolicExpression(\ref py_AstNode_page node, string comment)</b><br>
|
||||
- <b>\ref py_SymbolicExpression_page newSymbolicExpression(\ref py_AstNode_page node, string comment="")</b><br>
|
||||
Returns a new symbolic expression. Note that if there are simplification passes recorded, simplifications will be applied.
|
||||
|
||||
- <b>\ref py_SymbolicVariable_page newSymbolicVariable(integer varSize, string alias)</b><br>
|
||||
- <b>\ref py_SymbolicVariable_page newSymbolicVariable(integer varSize, string alias="")</b><br>
|
||||
Returns a new symbolic variable.
|
||||
|
||||
- <b>void popPathConstraint(void)</b><br>
|
||||
|
|
@ -380,13 +380,13 @@ Performs a dead store elimination simplification on a given block. If `padding`
|
|||
- <b>dict sliceExpressions(\ref py_SymbolicExpression_page expr)</b><br>
|
||||
Slices expressions from a given one (backward slicing) and returns all symbolic expressions as a dictionary of {integer SymExprId : \ref py_SymbolicExpression_page expr}.
|
||||
|
||||
- <b>\ref py_SymbolicVariable_page symbolizeExpression(integer symExprId, integer symVarSize, string symVarAlias)</b><br>
|
||||
- <b>\ref py_SymbolicVariable_page symbolizeExpression(integer symExprId, integer symVarSize, string symVarAlias="")</b><br>
|
||||
Converts a symbolic expression to a symbolic variable. `symVarSize` must be in bits. This function returns the new symbolic variable created.
|
||||
|
||||
- <b>\ref py_SymbolicVariable_page symbolizeMemory(\ref py_MemoryAccess_page mem, string symVarAlias)</b><br>
|
||||
- <b>\ref py_SymbolicVariable_page symbolizeMemory(\ref py_MemoryAccess_page mem, string symVarAlias="")</b><br>
|
||||
Converts a symbolic memory expression to a symbolic variable. This function returns the new symbolic variable created.
|
||||
|
||||
- <b>\ref py_SymbolicVariable_page symbolizeRegister(\ref py_Register_page reg, string symVarAlias)</b><br>
|
||||
- <b>\ref py_SymbolicVariable_page symbolizeRegister(\ref py_Register_page reg, string symVarAlias="")</b><br>
|
||||
Converts a symbolic register expression to a symbolic variable. This function returns the new symbolic variable created.
|
||||
|
||||
- <b>\ref py_AstNode_page synthesize(\ref py_AstNode_page node, bool constant=True, bool subexpr=True, bool opaque=False)</b><br>
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue