fix + improve python type stubs

This commit is contained in:
Marius Guggenmos 2025-01-11 06:12:18 +01:00
parent cd7da95558
commit b7c065b551
5 changed files with 126 additions and 72 deletions

View file

@ -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

View file

@ -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

View file

@ -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));

View file

@ -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.
*/

View file

@ -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>