diff --git a/.gitignore b/.gitignore index 9eb08bb4..e9bc5499 100644 --- a/.gitignore +++ b/.gitignore @@ -20,4 +20,6 @@ mise.toml dist/ decompiled_*/ decompiled_*.py +ROADMAP.md +.opencode-cache/ diff --git a/pylingual/decompiler.py b/pylingual/decompiler.py index 34d045b9..2fbaad21 100644 --- a/pylingual/decompiler.py +++ b/pylingual/decompiler.py @@ -44,6 +44,7 @@ from pylingual.models import CacheTranslator, load_models from pylingual.utils.generate_bytecode import CompileError, PyenvError, compile_version from pylingual.masking.model_disasm import create_global_masker, restore_masked_source_text +from pylingual.preprocessor import Preprocessor from pylingual.editable_bytecode import PYCFile from pylingual.segmentation.segmentation_search_strategies import get_top_k_predictions, m_deep_top_k, naive_confidence_priority, filter_subwords from pylingual.segmentation.sliding_window import merge, sliding_window @@ -122,6 +123,8 @@ def __call__(self): with tempfile.TemporaryDirectory() as tmp: self.tmp = Path(tmp) + self.preprocess_bytecode() + self.check_timeout() self.mask_bytecode() self.check_timeout() self.run_segmentation() @@ -208,6 +211,13 @@ def purge_comp_errors(self): self.source_context.purged_cfts = [] return [] + def preprocess_bytecode(self): + logger.info(f"Preprocessing bytecode for {self.name}...") + try: + Preprocessor().preprocess(self.pyc) + except Exception as e: + logger.warning(f"Preprocessing failed for {self.name}, continuing without container recovery: {e}") + def mask_bytecode(self): logger.info(f"Masking bytecode for {self.name}...") try: diff --git a/pylingual/editable_bytecode/EditableBytecode.py b/pylingual/editable_bytecode/EditableBytecode.py index cd561870..ed160a70 100644 --- a/pylingual/editable_bytecode/EditableBytecode.py +++ b/pylingual/editable_bytecode/EditableBytecode.py @@ -18,6 +18,20 @@ from typing import Callable +def _encode_name_arg(opname: str, name_index: int, original_arg: int | None, version: PythonVersion) -> int: + shifts = { + "LOAD_GLOBAL": 1 if version >= (3, 11) else 0, + "LOAD_ATTR": 1 if version >= (3, 12) else 0, + "LOAD_SUPER_ATTR": 2 if version >= (3, 12) else 0, + "IMPORT_NAME": 2 if version >= (3, 15) else 0, + } + shift = shifts.get(opname, 0) + if not shift: + return name_index + flags = int(original_arg or 0) & ((1 << shift) - 1) + return (name_index << shift) | flags + + class EditableBytecode: """ An editable representation of Pythonic bytecode. Many values in it, including its instructions, may be incorrect @@ -389,6 +403,7 @@ def _bake_jumps(self, add_extended=False): instruction.argval = instruction.target.offset else: instruction.arg = instruction.argval = instruction.target.offset + instruction.refresh_argrepr() if add_extended: # count the number of necessary preceeding EXTENDED_ARG instructions n_extendeds_needed = (instruction.arg > 0xFF) + (instruction.arg > 0xFFFF) + (instruction.arg > 0xFFFFFF) @@ -820,7 +835,8 @@ def insert_insts(self, insert_dict: dict[int, list[Inst]]) -> int: elif inst.optype == "name": if inst.argval not in self.co_names: self.co_names.append(inst.argval) - inst.arg = self.co_names.index(inst.argval) + name_index = self.co_names.index(inst.argval) + inst.arg = _encode_name_arg(inst.opname, name_index, inst.arg, self.version) elif inst.optype == "free" or inst.optype == "local": if inst.argval not in self.co_varnames: self.co_varnames.append(inst.argval) @@ -850,6 +866,18 @@ def insert_insts(self, insert_dict: dict[int, list[Inst]]) -> int: return len(to_insert) + def insert_instruction(self, index: int, instruction: Inst) -> int: + """Inserts one instruction at the specified index.""" + if index < len(self.instructions): + displaced = self.instructions[index] + if instruction.starts_line is None: + instruction.starts_line = displaced.starts_line + displaced.starts_line = None + for inst in self.instructions: + if inst.is_jump and inst.target is displaced: + inst._target = instruction + return self.insert_insts({index: [instruction]}) + def new_instruction(self, *args, **kwargs): """Creates a new instruction for use with this EditableBytecode object. This function does NOT automatically insert the instruction.""" return Inst(self, *args, **kwargs) diff --git a/pylingual/editable_bytecode/Instruction.py b/pylingual/editable_bytecode/Instruction.py index 354a6ba7..10905d42 100644 --- a/pylingual/editable_bytecode/Instruction.py +++ b/pylingual/editable_bytecode/Instruction.py @@ -133,6 +133,12 @@ def target(self, value): raise AttributeError("Only jump instructions have target attributes, not " + repr(self)) self._target = value + def refresh_argrepr(self): + """Refresh argrepr to reflect the current state of the instruction, e.g. after offset changes.""" + if self.is_jump and self._target is not None: + self.argval = self._target.offset + self.argrepr = "to " + repr(self._target.offset) + def add_reqs(self, *reqs): for r in reqs: if isinstance(r, tuple): diff --git a/pylingual/masking/global_masker.py b/pylingual/masking/global_masker.py index 82e33b16..0c959156 100644 --- a/pylingual/masking/global_masker.py +++ b/pylingual/masking/global_masker.py @@ -11,19 +11,24 @@ class TypeSensitiveDict(MutableMapping): def __init__(self, *args, **kwargs): self.store = dict() + self.original_keys = dict() self.update(dict(*args, **kwargs)) # use the free update to set keys def __getitem__(self, key): return self.store[self._key_transform(key)] def __setitem__(self, key, value): - self.store[self._key_transform(key)] = value + transformed = self._key_transform(key) + self.store[transformed] = value + self.original_keys[transformed] = key def __delitem__(self, key): - del self.store[self._key_transform(key)] + transformed = self._key_transform(key) + del self.store[transformed] + del self.original_keys[transformed] def __iter__(self): - return iter(self.store) + return iter(self.original_keys.values()) def __len__(self): return len(self.store) @@ -32,10 +37,10 @@ def __contains__(self, key): return self._key_transform(key) in self.store def keys(self) -> list: - return [self._key_restore(key) for key in self.store.keys()] + return list(self.original_keys.values()) def items(self): - return ((self._key_restore(key), value) for key, value in self.store.items()) + return ((self.original_keys[key], value) for key, value in self.store.items()) def values(self): return self.store.values() @@ -45,11 +50,21 @@ def _key_transform(self, key): return (key.value.decode("utf-8"), str) if type(key) == LongTypeForPython3: return (key.value, int) + if type(key) in (list, tuple): + return (tuple(self._key_transform(value) for value in key), type(key)) + if type(key) in (set, frozenset): + return (frozenset(self._key_transform(value) for value in key), type(key)) + if type(key) is dict: + return ( + frozenset((self._key_transform(k), self._key_transform(v)) for k, v in key.items()), + dict, + ) + try: + hash(key) + except TypeError: + return (repr(key), type(key)) return (key, type(key)) - def _key_restore(self, key): - return key[0] - #### Main Masker class Masker: @@ -184,16 +199,20 @@ def get_model_view(self, inst: Inst) -> str: elif isinstance(inst.argval, str) and inst.argval in self.global_tab: # have to do this check incase string is varname of type annotation view = f"{inst.opname} , {repr(self.mask(inst.argval))}" + elif type(inst.argval) in (list, tuple, frozenset, set, dict) and getattr(inst, "preprocessed_container", False): + # Container constants recovered by the preprocessor are masked atomically. + # TODO: Providing type information via type(inst.argval).__name__() would + # give the model useful context, but the current model was not trained on that format + # and produces garbled output (e.g. splitting "dict" across a mask boundary). Reverting + # to a plain mask token matches the training data. Revisit this after model retraining. + view = f"{inst.opname} , {self.mask(inst.argval)}" elif type(inst.argval) in (list, tuple, frozenset, set): - # do recursive in-place replacement of list elems if they are strs or bytestrs - def replace_list(consts): - """recursive replacement of elements in arbitrary list-like objects""" for idx, const in enumerate(consts): if isinstance(const, str): consts[idx] = repr(self.mask(inst.bytecode.resolve_namespace(const))) elif const is None: - continue # don't mask None + continue elif type(const) in (list, tuple, frozenset): consts[idx] = type(const)(replace_list(list(const))) elif type(const) is slice: @@ -202,19 +221,12 @@ def replace_list(consts): consts[idx] = self.mask(const) return consts - consts = list(deepcopy(inst.argval)) - consts = replace_list(consts) - + consts = replace_list(list(deepcopy(inst.argval))) if inst.bytecode.version < (3, 11): - # Format keyword argument list - # We left pad the list of kwargs so the model doesnt have to "look ahead" next_insts = inst.next_instructions - next_inst = next_insts[0] if next_insts != [] else None + next_inst = next_insts[0] if next_insts else None if next_inst is not None and inst.opname == "LOAD_CONST" and next_inst.opname == "CALL_FUNCTION_KW": consts = [""] * (next_inst.argval - len(consts)) + consts - - # cast back to original type and print repr - # demote quotes one layer arg_repr = repr(type(inst.argval)(consts)).replace("'", "").replace('"', "'") view = f"{inst.opname} , {arg_repr}" elif type(inst.argval) is slice: diff --git a/pylingual/masking/model_disasm.py b/pylingual/masking/model_disasm.py index 2fc25367..6d7b7f45 100644 --- a/pylingual/masking/model_disasm.py +++ b/pylingual/masking/model_disasm.py @@ -36,6 +36,7 @@ def create_global_masker(bytecode: EditableBytecode) -> Masker: for bc in bytecode.iter_bytecodes(): bc_co = bc.to_code(no_lnotab=True) + preprocessed = [inst.argval for inst in bc.instructions if getattr(inst, "preprocessed_container", False)] # create consts consts = list(deepcopy(bc_co.co_consts)) @@ -54,13 +55,13 @@ def create_global_masker(bytecode: EditableBytecode) -> Masker: # Don't needlessly increment the global_idx if const in global_tab: continue - if type(const) in (list, tuple, frozenset, set): + if type(const) in (list, tuple, frozenset, set) and not any(type(const) is type(value) and const == value for value in preprocessed): consts.extend(const) elif type(const) is slice: # decompose slice constants for 3.14+ consts.extend([const.start, const.stop, const.step]) else: - global_tab.update({bc.resolve_namespace(const): f""}) + global_tab[bc.resolve_namespace(const)] = f"" global_idx += 1 # create names @@ -69,25 +70,25 @@ def create_global_masker(bytecode: EditableBytecode) -> Masker: for n in name: if n in global_tab: continue - global_tab.update({bc.resolve_namespace(n): f""}) + global_tab[bc.resolve_namespace(n)] = f"" global_idx += 1 else: if name in global_tab: continue - global_tab.update({bc.resolve_namespace(name): f""}) + global_tab[bc.resolve_namespace(name)] = f"" global_idx += 1 for free in bc_co.co_freevars: if free in global_tab: continue - global_tab.update({free: f""}) + global_tab[free] = f"" global_idx += 1 if bc.version >= (3, 11): for cell in bc_co.co_cellvars: if cell in global_tab: continue - global_tab.update({cell: f""}) + global_tab[cell] = f"" global_idx += 1 for local in bc_co.co_varnames: @@ -95,15 +96,15 @@ def create_global_masker(bytecode: EditableBytecode) -> Masker: for local_item in local: if local_item in global_tab: continue - global_tab.update({bc.resolve_namespace(local_item): f""}) + global_tab[bc.resolve_namespace(local_item)] = f"" global_idx += 1 else: if local in global_tab: continue - global_tab.update({bc.resolve_namespace(local): f""}) + global_tab[bc.resolve_namespace(local)] = f"" global_idx += 1 - global_tab.update({bc_co.co_name: f""}) + global_tab[bc_co.co_name] = f"" global_idx += 1 return global_masker @@ -150,10 +151,12 @@ def format_source_replacement(mask_value: str) -> str: re_rel_pattern = re.compile(r"^(\s*)(import|from)\s*(\d+)(.*)", re.MULTILINE) -def unmask(source_line: str, replacements: dict, re_pattern: Pattern[str]): +def unmask(source_line: str, replacements: dict, container_masks: set[str], re_pattern: Pattern[str]): def m(match): s = match.span() r = replacements[match.group()] + if match.group() in container_masks: + return r if s[0] == 0 or s[1] >= len(match.string) or match.string[s[0] - 1] not in "\"'{}" and match.string[s[1]] not in "\"'{}": return r return use_escape_sequences(r) @@ -178,9 +181,19 @@ def fix_jump_targets(disasm: str) -> str: def restore_masked_source_text(lines: list[str], masker: Masker) -> list[str]: """Creates a large regex of all the tokens and their respective values Replaces everything in file text in one pass.""" - replacements = {re.escape(v): format_source_replacement(k) for k, v in masker.global_tab.items()} + replacements = {} + container_masks = set() + for value, mask in masker.global_tab.items(): + replacement = format_source_replacement(value) + if type(value) in (list, tuple, frozenset, set, dict): + # The model sometimes quotes an atomic container mask as though it were a string. + # Consume those wrapper quotes rather than escaping the container's contents. + replacements[re.escape(f"'{mask}'")] = replacement + replacements[re.escape(f'"{mask}"')] = replacement + container_masks.update((f"'{mask}'", f'"{mask}"', mask)) + replacements[re.escape(mask)] = replacement re_pattern = re.compile("|".join(replacements.keys())) - return [unmask(x, replacements, re_pattern) for x in lines] + return [unmask(x, replacements, container_masks, re_pattern) for x in lines] # replace mask values to start at 0 and count up diff --git a/pylingual/preprocessor/__init__.py b/pylingual/preprocessor/__init__.py new file mode 100644 index 00000000..f0d1e867 --- /dev/null +++ b/pylingual/preprocessor/__init__.py @@ -0,0 +1 @@ +from .preprocessor import Preprocessor, Tracer diff --git a/pylingual/preprocessor/container_recovery/__init__.py b/pylingual/preprocessor/container_recovery/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/pylingual/preprocessor/container_recovery/recovery.py b/pylingual/preprocessor/container_recovery/recovery.py new file mode 100644 index 00000000..aa2508fe --- /dev/null +++ b/pylingual/preprocessor/container_recovery/recovery.py @@ -0,0 +1,33 @@ +from __future__ import annotations + +from collections import defaultdict + +from .segment import Recovery, Segment + + +_strategies: dict[int, list] = defaultdict(list) + + +def register_recovery_strategy(priority: int): + def decorator(func): + _strategies[priority].append(func) + return func + return decorator + + +def get_strategies() -> list: + result = [] + for priority in sorted(_strategies.keys()): + result.extend(_strategies[priority]) + return result + + +def recover(seg: Segment, indent: int = 0) -> Recovery: + for strategy in get_strategies(): + result = strategy(seg, indent) + if result is not None: + return result + return Recovery(None, False) + + +from . import strategies # noqa: E402,F401 — trigger strategy registration \ No newline at end of file diff --git a/pylingual/preprocessor/container_recovery/segment.py b/pylingual/preprocessor/container_recovery/segment.py new file mode 100644 index 00000000..01da890c --- /dev/null +++ b/pylingual/preprocessor/container_recovery/segment.py @@ -0,0 +1,49 @@ +from __future__ import annotations + +from dataclasses import dataclass + + +INDENT = " " +SIDE_OFFSET = " " * 14 + + +@dataclass +class Recovery: + value: object + complete: bool + + +@dataclass +class Segment: + tag: str + ordered_children: list[Segment | tuple[int, object]] + base_stack_depth: int | None = None + start_offset: int | None = None + end_offset: int | None = None + + def compute_offsets(self) -> None: + offsets = [] + for child in self.ordered_children: + if isinstance(child, Segment): + child.compute_offsets() + if child.start_offset is not None: + offsets.append(child.start_offset) + if child.end_offset is not None: + offsets.append(child.end_offset) + else: + instr = child[1] + if hasattr(instr, "offset"): + offsets.append(instr.offset) + if offsets: + self.start_offset = min(offsets) + self.end_offset = max(offsets) + + def print(self, indent=0): + print(INDENT * indent, self.tag, sep="") + for child in self.ordered_children: + if isinstance(child, type(self)): + child.print(indent + 1) + else: + colpad = len(f"{SIDE_OFFSET}{child[0]} {child[1].opname}") + print(f"{SIDE_OFFSET}{child[0]} {child[1].opname}" + (" " * (20 - colpad) + f"{self.base_stack_depth}" if self.base_stack_depth else "")) + print() \ No newline at end of file diff --git a/pylingual/preprocessor/container_recovery/stack_analysis.py b/pylingual/preprocessor/container_recovery/stack_analysis.py new file mode 100644 index 00000000..8a56e62f --- /dev/null +++ b/pylingual/preprocessor/container_recovery/stack_analysis.py @@ -0,0 +1,271 @@ +from __future__ import annotations + +from xdis.cross_dis import xstack_effect + +from .segment import Segment + + +def analyze_stack(bytecode) -> list[tuple[int, object]]: + opc = bytecode.opcode + depth = 0 + results = [] + for instr in bytecode.instructions: + arg = instr.arg if instr.arg is not None else 0 + effect = xstack_effect(instr.opcode, opc, arg) + if effect is None: + effect = opc.oppush[instr.opcode] - opc.oppop[instr.opcode] + depth += effect + results.append((depth, instr)) + return results + + +def parse_list_recursive(remaining: list[tuple[int, object]]) -> Segment: + (post_base_depth, build_instr) = remaining.pop() + base_depth = remaining[-1][0] if remaining else post_base_depth + size = build_instr.arg + end_depth = base_depth - size + cur_elem_id = size + output_segs = [Segment("BUILD", [(post_base_depth, build_instr)])] + + if size == 0: + return Segment("LIST", output_segs[::-1]) + + cur_seg = [] + + while remaining: + if base_depth == end_depth or cur_elem_id <= 0: + break + (stack_depth, instr) = remaining.pop() + if stack_depth < base_depth: + output_segs.append(Segment(f"ELEM {cur_elem_id}", parse_bytecode_recursive(cur_seg[::-1]), base_depth)) + cur_elem_id -= 1 + base_depth -= 1 + cur_seg = [] + cur_seg.append((stack_depth, instr)) + + if cur_seg and cur_elem_id > 0: + output_segs.append(Segment(f"ELEM {cur_elem_id}", parse_bytecode_recursive(cur_seg[::-1]), base_depth)) + elif cur_seg: + remaining.extend(cur_seg) + + return Segment("LIST", output_segs[::-1]) + + +def parse_set_recursive(remaining: list[tuple[int, object]]) -> Segment: + (post_base_depth, build_instr) = remaining.pop() + base_depth = remaining[-1][0] if remaining else post_base_depth + size = build_instr.arg + end_depth = base_depth - size + cur_elem_id = size + output_segs = [Segment("BUILD", [(post_base_depth, build_instr)])] + + if size == 0: + return Segment("SET", output_segs[::-1]) + + cur_seg = [] + + while remaining: + if base_depth == end_depth or cur_elem_id <= 0: + break + (stack_depth, instr) = remaining.pop() + if stack_depth < base_depth: + output_segs.append(Segment(f"ELEM {cur_elem_id}", parse_bytecode_recursive(cur_seg[::-1]), base_depth)) + cur_elem_id -= 1 + base_depth -= 1 + cur_seg = [] + cur_seg.append((stack_depth, instr)) + + if cur_seg and cur_elem_id > 0: + output_segs.append(Segment(f"ELEM {cur_elem_id}", parse_bytecode_recursive(cur_seg[::-1]), base_depth)) + elif cur_seg: + remaining.extend(cur_seg) + + return Segment("SET", output_segs[::-1]) + + +def parse_tuple_recursive(remaining: list[tuple[int, object]]) -> Segment: + (post_base_depth, build_instr) = remaining.pop() + base_depth = remaining[-1][0] if remaining else post_base_depth + size = build_instr.arg + end_depth = base_depth - size + cur_elem_id = size + output_segs = [Segment("BUILD", [(post_base_depth, build_instr)])] + + if size == 0: + return Segment("TUPLE", output_segs[::-1]) + + cur_seg = [] + + while remaining: + if base_depth == end_depth or cur_elem_id <= 0: + break + (stack_depth, instr) = remaining.pop() + if stack_depth < base_depth: + output_segs.append(Segment(f"ELEM {cur_elem_id}", parse_bytecode_recursive(cur_seg[::-1]), base_depth)) + cur_elem_id -= 1 + base_depth -= 1 + cur_seg = [] + cur_seg.append((stack_depth, instr)) + + if cur_seg and cur_elem_id > 0: + output_segs.append(Segment(f"ELEM {cur_elem_id}", parse_bytecode_recursive(cur_seg[::-1]), base_depth)) + elif cur_seg: + remaining.extend(cur_seg) + + return Segment("TUPLE", output_segs[::-1]) + + +def parse_dict_recursive(remaining: list[tuple[int, object]]) -> Segment: + (post_base_depth, build_instr) = remaining.pop() + base_depth = remaining[-1][0] if remaining else post_base_depth + size = build_instr.arg + end_depth = base_depth - size * 2 + cur_elem_id = size * 2 + output_segs = [Segment("BUILD", [(post_base_depth, build_instr)])] + + if size == 0: + return Segment("DICT", output_segs[::-1]) + + cur_seg = [] + + is_key = False + + while remaining: + if base_depth == end_depth or cur_elem_id <= 0: + break + (stack_depth, instr) = remaining.pop() + if stack_depth < base_depth: + output_segs.append(Segment(f"KEY {(cur_elem_id + 1) // 2}" if is_key else f"VALUE {(cur_elem_id + 1) // 2}", parse_bytecode_recursive(cur_seg[::-1]), base_depth)) + cur_elem_id -= 1 + base_depth -= 1 + cur_seg = [] + is_key = not is_key + cur_seg.append((stack_depth, instr)) + + if cur_seg and cur_elem_id > 0: + output_segs.append(Segment(f"KEY {(cur_elem_id + 1) // 2}" if is_key else f"VALUE {(cur_elem_id + 1) // 2}", parse_bytecode_recursive(cur_seg[::-1]), base_depth)) + elif cur_seg: + remaining.extend(cur_seg) + + return Segment("DICT", output_segs[::-1]) + + +def parse_const_key_map_recursive(remaining: list[tuple[int, object]]) -> Segment: + (post_base_depth, build_instr) = remaining.pop() + size = build_instr.arg + output_segs = [Segment("BUILD", [(post_base_depth, build_instr)])] + + if size == 0: + return Segment("DICT", output_segs[::-1]) + + (key_depth, key_instr) = remaining.pop() + output_segs.append(Segment("KEY_TUPLE", [(key_depth, key_instr)])) + + base_depth = remaining[-1][0] if remaining else post_base_depth + end_depth = base_depth - size + cur_elem_id = size + cur_seg = [] + + while remaining: + if base_depth == end_depth or cur_elem_id <= 0: + break + (stack_depth, instr) = remaining.pop() + if stack_depth < base_depth: + output_segs.append(Segment(f"VALUE {cur_elem_id}", parse_bytecode_recursive(cur_seg[::-1]), base_depth)) + cur_elem_id -= 1 + base_depth -= 1 + cur_seg = [] + cur_seg.append((stack_depth, instr)) + + if cur_seg and cur_elem_id > 0: + output_segs.append(Segment(f"VALUE {cur_elem_id}", parse_bytecode_recursive(cur_seg[::-1]), base_depth)) + elif cur_seg: + remaining.extend(cur_seg) + + return Segment("DICT", output_segs[::-1]) + + +def parse_bytecode_recursive(remaining: list[tuple[int, object]]) -> list[Segment]: + output_segs = [] + cur_seg = [] + + while remaining: + (stack_depth, instr) = remaining[-1] + if instr.opname == "BUILD_MAP": + if cur_seg: output_segs.append(Segment("UNKNOWN", cur_seg[::-1])) + cur_seg = [] + dict_seg = parse_dict_recursive(remaining) + output_segs.append(dict_seg) + elif instr.opname == "BUILD_CONST_KEY_MAP": + if cur_seg: output_segs.append(Segment("UNKNOWN", cur_seg[::-1])) + cur_seg = [] + dict_seg = parse_const_key_map_recursive(remaining) + output_segs.append(dict_seg) + elif instr.opname in ("LIST_EXTEND", "LIST_APPEND"): + if cur_seg: output_segs.append(Segment("UNKNOWN", cur_seg[::-1])) + cur_seg = [] + trigger_opname = instr.opname + remaining.pop() + raw = [(stack_depth, instr)] + while remaining: + (_, i) = remaining[-1] + remaining.pop() + raw.append((_, i)) + if i.opname == "BUILD_LIST": + break + reversed_raw = raw[::-1] + children = [Segment("BUILD", [reversed_raw[0]])] + if trigger_opname == "LIST_APPEND": + elem_id = 1 + i = 1 + while i + 1 < len(reversed_raw): + children.append(Segment(f"ELEM {elem_id}", [reversed_raw[i], reversed_raw[i + 1]])) + elem_id += 1 + i += 2 + if i < len(reversed_raw): + children.append(Segment("UNKNOWN", reversed_raw[i:])) + else: + children.append(Segment("EXTEND", reversed_raw[1:])) + output_segs.append(Segment("LIST", children)) + elif instr.opname == "SET_UPDATE": + if cur_seg: output_segs.append(Segment("UNKNOWN", cur_seg[::-1])) + cur_seg = [] + remaining.pop() + raw = [(stack_depth, instr)] + while remaining: + (_, i) = remaining[-1] + remaining.pop() + raw.append((_, i)) + if i.opname == "BUILD_SET": + break + reversed_raw = raw[::-1] + children = [Segment("BUILD", [reversed_raw[0]]), Segment("EXTEND", reversed_raw[1:])] + output_segs.append(Segment("SET", children)) + elif instr.opname == "BUILD_LIST": + if cur_seg: output_segs.append(Segment("UNKNOWN", cur_seg[::-1])) + cur_seg = [] + list_seg = parse_list_recursive(remaining) + output_segs.append(list_seg) + elif instr.opname == "BUILD_SET": + if cur_seg: output_segs.append(Segment("UNKNOWN", cur_seg[::-1])) + cur_seg = [] + set_seg = parse_set_recursive(remaining) + output_segs.append(set_seg) + elif instr.opname == "BUILD_TUPLE": + if cur_seg: output_segs.append(Segment("UNKNOWN", cur_seg[::-1])) + cur_seg = [] + tuple_seg = parse_tuple_recursive(remaining) + output_segs.append(tuple_seg) + elif instr.opname == "LOAD_CONST" and isinstance(instr.argval, (tuple, list, dict, set, frozenset)): + remaining.pop() + if cur_seg: output_segs.append(Segment("UNKNOWN", cur_seg[::-1])) + cur_seg = [] + output_segs.append(Segment("LOAD_CONST_CONTAINER", [(stack_depth, instr)])) + else: + remaining.pop() + cur_seg.append((stack_depth, instr)) + + if cur_seg: + output_segs.append(Segment("UNKNOWN", cur_seg[::-1])) + + return output_segs[::-1] \ No newline at end of file diff --git a/pylingual/preprocessor/container_recovery/strategies/__init__.py b/pylingual/preprocessor/container_recovery/strategies/__init__.py new file mode 100644 index 00000000..a7958180 --- /dev/null +++ b/pylingual/preprocessor/container_recovery/strategies/__init__.py @@ -0,0 +1,4 @@ +from pathlib import Path + +__all__ = [x.stem for x in Path(__file__).parent.glob("*.py") if x.stem != "__init__"] +from . import * \ No newline at end of file diff --git a/pylingual/preprocessor/container_recovery/strategies/const.py b/pylingual/preprocessor/container_recovery/strategies/const.py new file mode 100644 index 00000000..e5d25e3a --- /dev/null +++ b/pylingual/preprocessor/container_recovery/strategies/const.py @@ -0,0 +1,50 @@ +from __future__ import annotations + +from ..recovery import register_recovery_strategy, recover +from ..segment import Recovery, Segment +from ..utils import _non_empty_children, _resolve_common_constant, _single_instr + + +@register_recovery_strategy(7) +def recover_const_container(seg: Segment, indent: int) -> Recovery | None: + if seg.tag != "LOAD_CONST_CONTAINER": + return None + instr = _single_instr(seg) + if instr is None: + return None + val = instr[1].argval + if isinstance(val, frozenset): + val = set(val) + return Recovery(val, True) + + +@register_recovery_strategy(8) +def recover_load_const(seg: Segment, indent: int) -> Recovery | None: + instr = _single_instr(seg) + if instr is None or instr[1].opname != "LOAD_CONST": + return None + return Recovery(instr[1].argval, True) + + +@register_recovery_strategy(9) +def recover_small_int(seg: Segment, indent: int) -> Recovery | None: + instr = _single_instr(seg) + if instr is None or instr[1].opname != "LOAD_SMALL_INT": + return None + return Recovery(instr[1].argval, True) + + +@register_recovery_strategy(10) +def recover_common_constant(seg: Segment, indent: int) -> Recovery | None: + instr = _single_instr(seg) + if instr is None or instr[1].opname != "LOAD_COMMON_CONSTANT": + return None + return Recovery(_resolve_common_constant(instr[1]), True) + + +@register_recovery_strategy(11) +def recover_delegate(seg: Segment, indent: int) -> Recovery | None: + children = _non_empty_children(seg) + if len(children) == 1 and isinstance(children[0], Segment): + return recover(children[0], indent) + return None diff --git a/pylingual/preprocessor/container_recovery/strategies/dict.py b/pylingual/preprocessor/container_recovery/strategies/dict.py new file mode 100644 index 00000000..3a157a5d --- /dev/null +++ b/pylingual/preprocessor/container_recovery/strategies/dict.py @@ -0,0 +1,65 @@ +from __future__ import annotations + +from ..recovery import register_recovery_strategy, recover +from ..segment import Recovery, Segment +from ..utils import _non_empty_children, _get_build_map_size + + +@register_recovery_strategy(0) +def recover_dict(seg: Segment, indent: int) -> Recovery | None: + size = _get_build_map_size(seg) + if size is None: + return None + keys_values = [c for c in seg.ordered_children if not (isinstance(c, Segment) and c.tag == "BUILD")] + if not keys_values: + return Recovery({}, True) + result = {} + complete = True + for i in range(0, len(keys_values), 2): + key_r = recover(keys_values[i], indent + 1) + val_r = recover(keys_values[i + 1], indent + 1) + if not key_r.complete: + complete = False + if not val_r.complete: + complete = False + # Duplicate keys affect source construction even though the runtime dict overwrites them. + if key_r.value in result: + complete = False + result[key_r.value] = val_r.value + return Recovery(result, complete) + + +@register_recovery_strategy(0) +def recover_const_key_map(seg: Segment, indent: int) -> Recovery | None: + build_instr = None + for child in reversed(seg.ordered_children): + if isinstance(child, Segment) and child.tag == "BUILD": + if len(child.ordered_children) == 1 and isinstance(child.ordered_children[0], tuple): + instr = child.ordered_children[0][1] + if instr.opname == "BUILD_CONST_KEY_MAP": + build_instr = instr + break + if build_instr is None: + return None + + key_tuple_seg = None + values = [] + for child in seg.ordered_children: + if isinstance(child, Segment): + if child.tag == "KEY_TUPLE": + key_tuple_seg = child + elif child.tag.startswith("VALUE"): + values.append(child) + if key_tuple_seg is None: + return None + + keys = key_tuple_seg.ordered_children[0][1].argval + + value_recoveries = [recover(v, indent + 1) for v in values] + complete = all(r.complete for r in value_recoveries) + + result = {keys[i]: value_recoveries[i].value for i in range(min(len(keys), len(value_recoveries)))} + # Duplicate keys affect source construction even though the runtime dict overwrites them. + if len(result) != min(len(keys), len(value_recoveries)): + complete = False + return Recovery(result, complete) diff --git a/pylingual/preprocessor/container_recovery/strategies/list.py b/pylingual/preprocessor/container_recovery/strategies/list.py new file mode 100644 index 00000000..7eefc4a9 --- /dev/null +++ b/pylingual/preprocessor/container_recovery/strategies/list.py @@ -0,0 +1,79 @@ +from __future__ import annotations + +from ..recovery import register_recovery_strategy, recover +from ..segment import Recovery, Segment +from ..utils import _non_empty_children, _recover_load_instr, _get_build_list_size, _has_list_append_pattern + + +@register_recovery_strategy(1) +def recover_build_list(seg: Segment, indent: int) -> Recovery | None: + size = _get_build_list_size(seg) + if size is None: + return None + if _has_list_append_pattern(seg): + return None + elems = [c for c in seg.ordered_children if not (isinstance(c, Segment) and c.tag == "BUILD")] + if not elems: + return Recovery([], True) + items = [] + complete = True + for elem in elems: + r = recover(elem, indent + 1) + if not r.complete: + complete = False + items.append(r.value) + return Recovery(items, complete) + + +@register_recovery_strategy(4) +def recover_list_extend(seg: Segment, indent: int) -> Recovery | None: + children = _non_empty_children(seg) + if len(children) != 2: + return None + build_child, extend_child = children + if not (isinstance(build_child, Segment) and build_child.tag == "BUILD"): + return None + if not (isinstance(extend_child, Segment) and extend_child.tag == "EXTEND"): + return None + build_instrs = build_child.ordered_children + if len(build_instrs) != 1 or not isinstance(build_instrs[0], tuple): + return None + if build_instrs[0][1].opname != "BUILD_LIST": + return None + extend_instrs = extend_child.ordered_children + if len(extend_instrs) != 2 or not all(isinstance(i, tuple) for i in extend_instrs): + return None + names = [i[1].opname for i in extend_instrs] + if names != ["LOAD_CONST", "LIST_EXTEND"]: + return None + val = extend_instrs[0][1].argval + return Recovery(list(val) if isinstance(val, tuple) else val, True) + + +@register_recovery_strategy(6) +def recover_list_append(seg: Segment, indent: int) -> Recovery | None: + children = _non_empty_children(seg) + if len(children) < 2: + return None + build_child = children[0] + if not (isinstance(build_child, Segment) and build_child.tag == "BUILD"): + return None + build_instrs = build_child.ordered_children + if len(build_instrs) != 1 or not isinstance(build_instrs[0], tuple): + return None + if build_instrs[0][1].opname != "BUILD_LIST": + return None + elem_children = children[1:] + if not all(isinstance(c, Segment) and c.tag.startswith("ELEM") for c in elem_children): + return None + items = [] + for elem in elem_children: + inner = elem.ordered_children + if len(inner) < 1 or not isinstance(inner[0], tuple): + return None + load_instr = inner[0][1] + r = _recover_load_instr(load_instr) + if r is None: + return Recovery(None, False) + items.append(r.value) + return Recovery(items, True) \ No newline at end of file diff --git a/pylingual/preprocessor/container_recovery/strategies/set.py b/pylingual/preprocessor/container_recovery/strategies/set.py new file mode 100644 index 00000000..c5d2cb9a --- /dev/null +++ b/pylingual/preprocessor/container_recovery/strategies/set.py @@ -0,0 +1,50 @@ +from __future__ import annotations + +from ..recovery import register_recovery_strategy, recover +from ..segment import Recovery, Segment +from ..utils import _non_empty_children, _get_build_set_size + + +@register_recovery_strategy(2) +def recover_build_set(seg: Segment, indent: int) -> Recovery | None: + size = _get_build_set_size(seg) + if size is None or size == 0: + return None + elems = [c for c in seg.ordered_children if not (isinstance(c, Segment) and c.tag == "BUILD")] + if not elems: + return Recovery(set(), True) + items = [] + complete = True + for elem in elems: + r = recover(elem, indent + 1) + if not r.complete: + complete = False + items.append(r.value) + return Recovery(set(items), complete) + + +@register_recovery_strategy(5) +def recover_set_update(seg: Segment, indent: int) -> Recovery | None: + children = _non_empty_children(seg) + if len(children) != 2: + return None + build_child, extend_child = children + if not (isinstance(build_child, Segment) and build_child.tag == "BUILD"): + return None + if not (isinstance(extend_child, Segment) and extend_child.tag == "EXTEND"): + return None + build_instrs = build_child.ordered_children + if len(build_instrs) != 1 or not isinstance(build_instrs[0], tuple): + return None + if build_instrs[0][1].opname != "BUILD_SET": + return None + extend_instrs = extend_child.ordered_children + if len(extend_instrs) != 2 or not all(isinstance(i, tuple) for i in extend_instrs): + return None + names = [i[1].opname for i in extend_instrs] + if names != ["LOAD_CONST", "SET_UPDATE"]: + return None + val = extend_instrs[0][1].argval + if isinstance(val, (tuple, frozenset, set)): + return Recovery(set(val), True) + return Recovery(val, True) \ No newline at end of file diff --git a/pylingual/preprocessor/container_recovery/strategies/tuple.py b/pylingual/preprocessor/container_recovery/strategies/tuple.py new file mode 100644 index 00000000..3ef2bb8e --- /dev/null +++ b/pylingual/preprocessor/container_recovery/strategies/tuple.py @@ -0,0 +1,23 @@ +from __future__ import annotations + +from ..recovery import register_recovery_strategy, recover +from ..segment import Recovery, Segment +from ..utils import _get_build_tuple_size + + +@register_recovery_strategy(3) +def recover_build_tuple(seg: Segment, indent: int) -> Recovery | None: + size = _get_build_tuple_size(seg) + if size is None or size == 0: + return None + elems = [c for c in seg.ordered_children if not (isinstance(c, Segment) and c.tag == "BUILD")] + if not elems: + return Recovery((), True) + items = [] + complete = True + for elem in elems: + r = recover(elem, indent + 1) + if not r.complete: + complete = False + items.append(r.value) + return Recovery(tuple(items), complete) \ No newline at end of file diff --git a/pylingual/preprocessor/container_recovery/utils.py b/pylingual/preprocessor/container_recovery/utils.py new file mode 100644 index 00000000..1866f16e --- /dev/null +++ b/pylingual/preprocessor/container_recovery/utils.py @@ -0,0 +1,102 @@ +from __future__ import annotations + +from .segment import Recovery, Segment + + +_COMMON_CONSTANTS = { + "AssertionError": AssertionError, + "NotImplementedError": NotImplementedError, + "tuple": tuple, + "all": all, + "any": any, + "list": list, + "set": set, + "None": None, + '""': "", + "True": True, + "False": False, + "-1": -1, + "frozenset": frozenset, + "()": (), +} + + +def _resolve_common_constant(instr): + return instr.argval if instr.argval != instr.arg else _COMMON_CONSTANTS[instr.argrepr] + + +def _get_build_map_size(seg: Segment) -> int | None: + for child in reversed(seg.ordered_children): + if isinstance(child, Segment) and child.tag == "BUILD": + if len(child.ordered_children) == 1 and isinstance(child.ordered_children[0], tuple): + instr = child.ordered_children[0][1] + if instr.opname == "BUILD_MAP": + return instr.arg + return None + + +def _get_build_list_size(seg: Segment) -> int | None: + for child in reversed(seg.ordered_children): + if isinstance(child, Segment) and child.tag == "BUILD": + if len(child.ordered_children) == 1 and isinstance(child.ordered_children[0], tuple): + instr = child.ordered_children[0][1] + if instr.opname == "BUILD_LIST": + return instr.arg + return None + + +def _get_build_set_size(seg: Segment) -> int | None: + for child in reversed(seg.ordered_children): + if isinstance(child, Segment) and child.tag == "BUILD": + if len(child.ordered_children) == 1 and isinstance(child.ordered_children[0], tuple): + instr = child.ordered_children[0][1] + if instr.opname == "BUILD_SET": + return instr.arg + return None + + +def _get_build_tuple_size(seg: Segment) -> int | None: + for child in reversed(seg.ordered_children): + if isinstance(child, Segment) and child.tag == "BUILD": + if len(child.ordered_children) == 1 and isinstance(child.ordered_children[0], tuple): + instr = child.ordered_children[0][1] + if instr.opname == "BUILD_TUPLE": + return instr.arg + return None + + +def _has_list_append_pattern(seg: Segment) -> bool: + for c in seg.ordered_children: + if isinstance(c, Segment) and c.tag.startswith("ELEM"): + for instr in c.ordered_children: + if isinstance(instr, tuple) and instr[1].opname == "LIST_APPEND": + return True + if isinstance(c, Segment) and c.tag == "EXTEND": + for instr in c.ordered_children: + if isinstance(instr, tuple) and instr[1].opname == "LIST_EXTEND": + return True + return False + + +def _non_empty_children(seg: Segment) -> list: + return [c for c in seg.ordered_children + if not (isinstance(c, Segment) and c.tag == "UNKNOWN" and len(c.ordered_children) == 0)] + + +def _single_instr(seg: Segment) -> tuple | None: + children = _non_empty_children(seg) + if len(children) == 1 and isinstance(children[0], tuple): + return children[0] + if len(children) == 1 and isinstance(children[0], Segment) and children[0].tag == "UNKNOWN" and len(children[0].ordered_children) == 1 and isinstance(children[0].ordered_children[0], tuple): + return children[0].ordered_children[0] + return None + + +def _recover_load_instr(instr) -> Recovery | None: + if instr.opname == "LOAD_SMALL_INT": + return Recovery(instr.argval, True) + if instr.opname == "LOAD_CONST": + return Recovery(instr.argval, True) + if instr.opname == "LOAD_COMMON_CONSTANT": + return Recovery(_resolve_common_constant(instr), True) + return None diff --git a/pylingual/preprocessor/preprocessor.py b/pylingual/preprocessor/preprocessor.py new file mode 100644 index 00000000..acf25cb6 --- /dev/null +++ b/pylingual/preprocessor/preprocessor.py @@ -0,0 +1,126 @@ +from __future__ import annotations + +from xdis.cross_dis import instruction_size, xstack_effect + +from pylingual.editable_bytecode import EditableBytecode, Inst +from .container_recovery.recovery import recover +from .container_recovery.segment import Segment +from .container_recovery.stack_analysis import analyze_stack, parse_bytecode_recursive + +_CONTAINER_TAGS = {"LIST", "SET", "TUPLE", "DICT"} + + +def _make_function_stack_effect(arg, version) -> int: + return -int(arg or 0).bit_count() - (version < (3, 11)) + + +class Tracer: + """Trace an object placed on the stack before an instruction sequence.""" + + def __init__(self, instructions: list[Inst]): + self.instructions = instructions + + def trace(self) -> tuple[int, int] | None: + """Return the consuming instruction index and its zero-based TOS argument.""" + depth = 1 + for index, inst in enumerate(self.instructions): + effect = xstack_effect(inst.opcode, inst.bytecode.opcode, inst.arg or 0) + if effect is None and inst.opname == "MAKE_FUNCTION": + effect = _make_function_stack_effect(inst.arg, inst.bytecode.version) + elif effect is None: + effect = inst.bytecode.opcode.oppush[inst.opcode] - inst.bytecode.opcode.oppop[inst.opcode] + + push = 1 if inst.opname.startswith("BUILD_") else inst.bytecode.opcode.oppush[inst.opcode] + pop = push - effect + if push >= 0 and pop >= depth: + return index, depth - 1 + depth += effect + if depth <= 0: + break + return None + + +def _preserve_container(consumer: Inst, argnum: int) -> bool: + if consumer.opname in ("MAKE_FUNCTION", "SET_FUNCTION_ATTRIBUTE"): + return True + if consumer.opname == "DICT_MERGE" and argnum in (0, 1): + return True + if consumer.opname == "CALL_FUNCTION_EX" and argnum == 0: + return consumer.arg is None or isinstance(consumer.arg, int) and bool(consumer.arg & 1) + return False + + +class Preprocessor: + """Rewrites container construction bytecode into single LOAD_CONST instructions. + + For each code object, analyzes the stack to identify container construction + patterns (lists, dicts, sets, tuples, frozensets), recovers their values via + strategy-based recovery, and replaces the bytecode span with a single + LOAD_CONST referencing the materialized value in co_consts. + """ + + def preprocess(self, bc: EditableBytecode) -> None: + """Process all code objects (top-level + nested).""" + for code_obj in bc.iter_bytecodes(): + self._process_code_object(code_obj) + + def _process_code_object(self, bc: EditableBytecode) -> None: + depths = analyze_stack(bc) + segments = parse_bytecode_recursive(depths) + for seg in reversed(segments): + self._process_segment(bc, seg) + bc._bake_jumps() + + def _process_segment(self, bc: EditableBytecode, seg: Segment) -> None: + if seg.tag in _CONTAINER_TAGS: + seg.compute_offsets() + if seg.start_offset is not None and seg.end_offset is not None: + recovery = recover(seg) + if recovery.complete: + end_index = bc.instructions.index(bc.get_by_offset(seg.end_offset)) + following = bc.instructions[end_index + 1:] + consumption = Tracer(following).trace() + if consumption is None or not _preserve_container(following[consumption[0]], consumption[1]): + self._collapse_segment(bc, seg, recovery.value) + return + for child in reversed(seg.ordered_children): + if isinstance(child, Segment): + self._process_segment(bc, child) + + def _collapse_segment(self, bc: EditableBytecode, seg: Segment, value) -> None: + if len(value) == 0: + return + + const_index = None + for i, existing in enumerate(bc.co_consts): + if value is existing or (type(value) is type(existing) and value == existing): + const_index = i + break + if const_index is None: + bc.co_consts.append(value) + const_index = len(bc.co_consts) - 1 + + start_inst = bc.get_by_offset(seg.start_offset) + end_inst = bc.get_by_offset(seg.end_offset) + start_idx = bc.instructions.index(start_inst) + end_idx = bc.instructions.index(end_inst) + + opcode = bc.opcode.opmap["LOAD_CONST"] + new_inst = bc.new_instruction( + "LOAD_CONST", + opcode, + "const", + instruction_size(opcode, bc.opcode), + const_index, + value, + repr(value), + True, + -1, + None, + False, + False, + ) + new_inst.preprocessed_container = True + + bc.remove_instructions(bc.instructions[start_idx : end_idx + 1]) + bc.insert_instruction(start_idx, new_inst) diff --git a/pyproject.toml b/pyproject.toml index 753ce6e3..40e80bfc 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -36,6 +36,9 @@ dependencies = [ "tensorboardx>=2.6.4", ] +[project.optional-dependencies] +test = ["pytest"] + [project.urls] homepage = "https://pylingual.io" diff --git a/tests/preprocessing/fixtures/complex_keys.py b/tests/preprocessing/fixtures/complex_keys.py new file mode 100644 index 00000000..337018ea --- /dev/null +++ b/tests/preprocessing/fixtures/complex_keys.py @@ -0,0 +1,4 @@ +a = {1 + 2: 3} +b = {len("hi"): "hello"} +c = {True: False, (1, 2): "tuple_key"} +print(a, b, c) \ No newline at end of file diff --git a/tests/preprocessing/fixtures/complex_nested.py b/tests/preprocessing/fixtures/complex_nested.py new file mode 100644 index 00000000..4a6e185a --- /dev/null +++ b/tests/preprocessing/fixtures/complex_nested.py @@ -0,0 +1,6 @@ +# Complex nested structures +a = {1: {2: 3}, 4: {5, 6}} +b = [{(1, 2): 3}, {"a": (4, 5)}] +c = ({1, 2}, [3, 4], {5: 6}) +e = [{"x": {1, 2}}, {"y": (3, 4)}] +print(a, b, c, e) diff --git a/tests/preprocessing/fixtures/complex_nested_mixed.py b/tests/preprocessing/fixtures/complex_nested_mixed.py new file mode 100644 index 00000000..659e93ed --- /dev/null +++ b/tests/preprocessing/fixtures/complex_nested_mixed.py @@ -0,0 +1,41 @@ +# Complex nested structures with mixed container types +# Tuple of dicts containing tuples and lists +tuple_of_dicts = ( + {1: (2, 3), 'a': [4, 5]}, + {6: [7, 8], 'b': (9, 10)}, +) + +# Dict of sets containing tuples +dict_of_sets = { + 'x': {(1, 2), (3, 4)}, + 'y': {(5, 6), (7, 8)}, +} + +# List of tuples containing dicts +list_of_tuples = [ + ({'a': 1}, {'b': 2}), + ({'c': 3}, {'d': 4}), +] + +# Tuple of lists containing sets +tuple_of_lists = ( + [{1, 2}, {3, 4}], + [{5, 6}, {7, 8}], +) + +# Nested: dict containing list containing dict containing set +deeply_nested = { + 'level1': [ + {'level2': {1, 2, 3}}, + {'level3': {4, 5, 6}}, + ], +} + +# Mixed with empty containers +mixed_empties = { + 'empty_dict': {}, + 'empty_list': [], + 'nonempty': [1, 2], +} + +print("Done") \ No newline at end of file diff --git a/tests/preprocessing/fixtures/dict_call.py b/tests/preprocessing/fixtures/dict_call.py new file mode 100644 index 00000000..c08838f7 --- /dev/null +++ b/tests/preprocessing/fixtures/dict_call.py @@ -0,0 +1,2 @@ +a = dict(x=1, y=2) +print(a) \ No newline at end of file diff --git a/tests/preprocessing/fixtures/dict_complex_keys.py b/tests/preprocessing/fixtures/dict_complex_keys.py new file mode 100644 index 00000000..69ffb5fa --- /dev/null +++ b/tests/preprocessing/fixtures/dict_complex_keys.py @@ -0,0 +1,6 @@ +# Dicts with complex keys +a = {(1, 2): "tuple key"} +b = {frozenset({1, 2}): "frozenset key"} +c = {True: "bool key", None: "none key", 1: "int key"} +d = {(1, (2, 3)): {"nested": [4, 5]}} +print(a, b, c, d) diff --git a/tests/preprocessing/fixtures/dict_with_containers.py b/tests/preprocessing/fixtures/dict_with_containers.py new file mode 100644 index 00000000..56423043 --- /dev/null +++ b/tests/preprocessing/fixtures/dict_with_containers.py @@ -0,0 +1,2 @@ +d = {"a": (1, 2), "b": {3, 4}} +print(d) \ No newline at end of file diff --git a/tests/preprocessing/fixtures/edge_cases.py b/tests/preprocessing/fixtures/edge_cases.py new file mode 100644 index 00000000..347c4a89 --- /dev/null +++ b/tests/preprocessing/fixtures/edge_cases.py @@ -0,0 +1,11 @@ +# Edge cases +empty_dict = {} +empty_list = [] +empty_tuple = () +# Note: empty set is set() call, not literal +single_in_list = [1] +single_in_tuple = (1,) +single_in_dict = {"x": 1} +nested_empty = {"a": [], "b": {}, "c": ()} +deep_nesting = [[[{"x": {1, (2, 3)}}]]] +print(empty_dict, empty_list, empty_tuple, single_in_list, single_in_tuple, single_in_dict, nested_empty, deep_nesting) diff --git a/tests/preprocessing/fixtures/empty_containers.py b/tests/preprocessing/fixtures/empty_containers.py new file mode 100644 index 00000000..15e3c10d --- /dev/null +++ b/tests/preprocessing/fixtures/empty_containers.py @@ -0,0 +1,15 @@ +# Empty container tests +empty_dict = {} +empty_list = [] +empty_tuple = () +# Note: empty set is set() call, not literal + +# Nested empty containers +nested_empty = { + 'a': {}, + 'b': [], + 'c': (), + 'd': {'x': {}, 'y': []}, +} + +print(empty_dict, empty_list, empty_tuple) \ No newline at end of file diff --git a/tests/preprocessing/fixtures/empty_dict.py b/tests/preprocessing/fixtures/empty_dict.py new file mode 100644 index 00000000..e08cd9f0 --- /dev/null +++ b/tests/preprocessing/fixtures/empty_dict.py @@ -0,0 +1,2 @@ +a = {} +print(a) \ No newline at end of file diff --git a/tests/preprocessing/fixtures/empty_set_tuple.py b/tests/preprocessing/fixtures/empty_set_tuple.py new file mode 100644 index 00000000..55e7c6f4 --- /dev/null +++ b/tests/preprocessing/fixtures/empty_set_tuple.py @@ -0,0 +1,3 @@ +x = () +y = set() +print(x, y) \ No newline at end of file diff --git a/tests/preprocessing/fixtures/list_of_dicts.py b/tests/preprocessing/fixtures/list_of_dicts.py new file mode 100644 index 00000000..dbde1a8e --- /dev/null +++ b/tests/preprocessing/fixtures/list_of_dicts.py @@ -0,0 +1,2 @@ +x = [{"a": 1}, {"b": 2, "c": 3}] +print(x) \ No newline at end of file diff --git a/tests/preprocessing/fixtures/long_list.py b/tests/preprocessing/fixtures/long_list.py new file mode 100644 index 00000000..c1c15c3f --- /dev/null +++ b/tests/preprocessing/fixtures/long_list.py @@ -0,0 +1,2 @@ +x = [0, 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] +print(x) \ No newline at end of file diff --git a/tests/preprocessing/fixtures/mixed_list.py b/tests/preprocessing/fixtures/mixed_list.py new file mode 100644 index 00000000..6a31bb8e --- /dev/null +++ b/tests/preprocessing/fixtures/mixed_list.py @@ -0,0 +1,6 @@ +def foo(x): + return x + 1 + +x = [1, 2, 3, foo(3), 4, 5, 6] + +print(x) diff --git a/tests/preprocessing/fixtures/mixed_list2.py b/tests/preprocessing/fixtures/mixed_list2.py new file mode 100644 index 00000000..8672a89f --- /dev/null +++ b/tests/preprocessing/fixtures/mixed_list2.py @@ -0,0 +1,5 @@ +def foo(x): + return x + 1 + +x = [1, 2, foo(3), 4, 5, 6, foo(7), 8, 9] +print(x) \ No newline at end of file diff --git a/tests/preprocessing/fixtures/mixed_types.py b/tests/preprocessing/fixtures/mixed_types.py new file mode 100644 index 00000000..a0a4e9fc --- /dev/null +++ b/tests/preprocessing/fixtures/mixed_types.py @@ -0,0 +1,2 @@ +x = [1, "hello", True, None, 3.14, (1, 2), b"bytes"] +print(x) \ No newline at end of file diff --git a/tests/preprocessing/fixtures/multiple_dicts.py b/tests/preprocessing/fixtures/multiple_dicts.py new file mode 100644 index 00000000..c456a799 --- /dev/null +++ b/tests/preprocessing/fixtures/multiple_dicts.py @@ -0,0 +1,3 @@ +a = {"x": 1} +b = {"y": 2, "z": 3} +print(a, b) \ No newline at end of file diff --git a/tests/preprocessing/fixtures/nested_containers.py b/tests/preprocessing/fixtures/nested_containers.py new file mode 100644 index 00000000..31e63515 --- /dev/null +++ b/tests/preprocessing/fixtures/nested_containers.py @@ -0,0 +1,2 @@ +d = {"tuple": (1, (2, 3)), "list": [1, [2, 3]], "set": {1, 2}} +print(d) \ No newline at end of file diff --git a/tests/preprocessing/fixtures/nested_lists.py b/tests/preprocessing/fixtures/nested_lists.py new file mode 100644 index 00000000..c2ff756f --- /dev/null +++ b/tests/preprocessing/fixtures/nested_lists.py @@ -0,0 +1,2 @@ +x = [[1, 2], [3, 4], [5, 6]] +print(x) \ No newline at end of file diff --git a/tests/preprocessing/fixtures/set_and_tuple.py b/tests/preprocessing/fixtures/set_and_tuple.py new file mode 100644 index 00000000..542e5268 --- /dev/null +++ b/tests/preprocessing/fixtures/set_and_tuple.py @@ -0,0 +1,3 @@ +s = {1, 2, 3} +t = (4, 5, 6) +print(s, t) \ No newline at end of file diff --git a/tests/preprocessing/fixtures/set_basic.py b/tests/preprocessing/fixtures/set_basic.py new file mode 100644 index 00000000..fcbba70f --- /dev/null +++ b/tests/preprocessing/fixtures/set_basic.py @@ -0,0 +1,2 @@ +s = {1, 2, 3} +print(s) \ No newline at end of file diff --git a/tests/preprocessing/fixtures/set_complex.py b/tests/preprocessing/fixtures/set_complex.py new file mode 100644 index 00000000..a38ba6c1 --- /dev/null +++ b/tests/preprocessing/fixtures/set_complex.py @@ -0,0 +1,5 @@ +# Sets with various element types +a = {(1, 2), (3, 4)} +b = {(1, 2, 3), (4, 5, 6)} +c = {frozenset({1, 2}), frozenset({3, 4})} +print(a, b, c) diff --git a/tests/preprocessing/fixtures/set_mixed.py b/tests/preprocessing/fixtures/set_mixed.py new file mode 100644 index 00000000..42a5b40f --- /dev/null +++ b/tests/preprocessing/fixtures/set_mixed.py @@ -0,0 +1,2 @@ +s = {1, "hello", True, None, 3.14} +print(s) \ No newline at end of file diff --git a/tests/preprocessing/fixtures/set_of_tuples.py b/tests/preprocessing/fixtures/set_of_tuples.py new file mode 100644 index 00000000..6fbfb56e --- /dev/null +++ b/tests/preprocessing/fixtures/set_of_tuples.py @@ -0,0 +1,2 @@ +s = {(1, 2), (3, 4)} +print(s) \ No newline at end of file diff --git a/tests/preprocessing/fixtures/single_dict.py b/tests/preprocessing/fixtures/single_dict.py new file mode 100644 index 00000000..2fc5e291 --- /dev/null +++ b/tests/preprocessing/fixtures/single_dict.py @@ -0,0 +1,2 @@ +x = {1: 2} +print(x) diff --git a/tests/preprocessing/fixtures/single_element.py b/tests/preprocessing/fixtures/single_element.py new file mode 100644 index 00000000..c2d1d6cb --- /dev/null +++ b/tests/preprocessing/fixtures/single_element.py @@ -0,0 +1,2 @@ +a = {"x": 1} +print(a) \ No newline at end of file diff --git a/tests/preprocessing/fixtures/single_element_mixed.py b/tests/preprocessing/fixtures/single_element_mixed.py new file mode 100644 index 00000000..b6d2aec2 --- /dev/null +++ b/tests/preprocessing/fixtures/single_element_mixed.py @@ -0,0 +1,10 @@ +# Single element containers with nested types +single_dict_in_list = [{'x': 1}] +single_list_in_dict = {'a': [1]} +single_set_in_tuple = ({1},) +single_tuple_in_set = {(1,)} + +# Single element with nested container +single_nested = {'outer': {'inner': [1]}} + +print(single_dict_in_list, single_list_in_dict) \ No newline at end of file diff --git a/tests/preprocessing/fixtures/single_element_set_tuple.py b/tests/preprocessing/fixtures/single_element_set_tuple.py new file mode 100644 index 00000000..9c732b7b --- /dev/null +++ b/tests/preprocessing/fixtures/single_element_set_tuple.py @@ -0,0 +1,3 @@ +x = (1,) +y = {1} +print(x, y) \ No newline at end of file diff --git a/tests/preprocessing/fixtures/tuple_basic.py b/tests/preprocessing/fixtures/tuple_basic.py new file mode 100644 index 00000000..69c0a743 --- /dev/null +++ b/tests/preprocessing/fixtures/tuple_basic.py @@ -0,0 +1,2 @@ +t = (1, 2, 3) +print(t) \ No newline at end of file diff --git a/tests/preprocessing/fixtures/tuple_mixed.py b/tests/preprocessing/fixtures/tuple_mixed.py new file mode 100644 index 00000000..071f2f42 --- /dev/null +++ b/tests/preprocessing/fixtures/tuple_mixed.py @@ -0,0 +1,2 @@ +t = (1, "hello", True, None, 3.14) +print(t) \ No newline at end of file diff --git a/tests/preprocessing/fixtures/tuple_with_nesting.py b/tests/preprocessing/fixtures/tuple_with_nesting.py new file mode 100644 index 00000000..7c7db08f --- /dev/null +++ b/tests/preprocessing/fixtures/tuple_with_nesting.py @@ -0,0 +1,2 @@ +t = (1, [2, 3], {"a": 1}) +print(t) \ No newline at end of file diff --git a/tests/preprocessing/test_edge_cases.py b/tests/preprocessing/test_edge_cases.py new file mode 100644 index 00000000..2cbec801 --- /dev/null +++ b/tests/preprocessing/test_edge_cases.py @@ -0,0 +1,159 @@ +from __future__ import annotations + +import sys +from types import SimpleNamespace + +import pytest +from xdis import get_opcode + +from pylingual.editable_bytecode import EditableBytecode +from pylingual.preprocessor import Preprocessor +from pylingual.preprocessor.container_recovery.utils import _resolve_common_constant + + +def _preprocess(source: str) -> EditableBytecode: + code = compile(source, "", "exec") + opc = get_opcode(sys.version_info[:2], False) + bc = EditableBytecode(code, opc, sys.version_info[:3]) + Preprocessor().preprocess(bc) + return bc + + +def _co_const_values(bc: EditableBytecode) -> list: + return [c for c in bc.co_consts if isinstance(c, (list, tuple, set, frozenset, dict))] + + +def test_empty_containers_not_added_by_preprocessor(): + code = compile("a = {}\nb = []\nc = ()\n", "", "exec") + opc = get_opcode(sys.version_info[:2], False) + bc = EditableBytecode(code, opc, sys.version_info[:3]) + before = [c for c in bc.co_consts if isinstance(c, (list, tuple, set, frozenset, dict))] + Preprocessor().preprocess(bc) + after = [c for c in bc.co_consts if isinstance(c, (list, tuple, set, frozenset, dict))] + assert after == before, f"Preprocessor added empty containers: before={before!r}, after={after!r}" + + +def test_falsy_non_empty_list_in_co_consts(): + bc = _preprocess("a = [0]\n") + values = _co_const_values(bc) + assert [0] in values, f"[0] not in co_consts values {values!r}" + + +def test_falsy_non_empty_dict_in_co_consts(): + bc = _preprocess("a = {0: 0}\n") + values = _co_const_values(bc) + assert {0: 0} in values, f"{{0: 0}} not in co_consts values {values!r}" + + +def test_falsy_non_empty_tuple_in_co_consts(): + bc = _preprocess("a = (0,)\n") + values = _co_const_values(bc) + assert (0,) in values, f"(0,) not in co_consts values {values!r}" + + +def test_falsy_non_empty_set_in_co_consts(): + bc = _preprocess("a = {0}\n") + values = _co_const_values(bc) + assert any(type(v) is set and v == {0} for v in values), f"{{0}} not in co_consts values {values!r}" + + +def test_set_recovered_from_frozenset(): + bc = _preprocess("a = {1, 2, 3}\n") + values = _co_const_values(bc) + sets = [v for v in values if type(v) is set] + assert any(v == {1, 2, 3} for v in sets), f"set {{1,2,3}} not in {values!r}" + + +def test_nested_containers_in_co_consts(): + bc = _preprocess("a = [[1, 2], [3, 4]]\n") + values = _co_const_values(bc) + assert [[1, 2], [3, 4]] in values, f"[[1,2],[3,4]] not in co_consts values {values!r}" + + +def test_deeply_nested_containers(): + bc = _preprocess("a = {'x': [1, {'y': (2, 3)}]}\n") + values = _co_const_values(bc) + assert {'x': [1, {'y': (2, 3)}]} in values, f"nested dict not in co_consts values {values!r}" + + +def test_multiple_containers_same_type(): + bc = _preprocess("a = [1, 2]\nb = [3, 4]\nc = [5, 6]\n") + values = _co_const_values(bc) + assert [1, 2] in values, f"[1,2] not in {values!r}" + assert [3, 4] in values, f"[3,4] not in {values!r}" + assert [5, 6] in values, f"[5,6] not in {values!r}" + + +def test_dedup_identical_containers(): + bc = _preprocess("a = [1, 2, 3]\nb = [1, 2, 3]\n") + values = _co_const_values(bc) + count = sum(1 for v in values if type(v) is list and v == [1, 2, 3]) + assert count == 1, f"Expected 1 deduplicated [1,2,3], got {count} in {values!r}" + + +def test_no_dedup_different_type_same_value(): + bc = _preprocess("a = (1, 2)\nb = [1, 2]\n") + values = _co_const_values(bc) + tuples = [v for v in values if type(v) is tuple and v == (1, 2)] + lists = [v for v in values if type(v) is list and v == [1, 2]] + assert len(tuples) == 1, f"Expected 1 tuple (1,2), got {tuples!r}" + assert len(lists) == 1, f"Expected 1 list [1,2], got {lists!r}" + + +def test_preprocessor_rebases_exception_table_offsets(): + code = compile( + "def f():\n" + " try:\n" + " value = {'a': 1, 'b': 2, 'c': 3}\n" + " finally:\n" + " value = None\n", + "", + "exec", + ) + root = EditableBytecode(code, get_opcode(sys.version_info[:2], False), sys.version_info[:3]) + bc = root.bytecode_lookup["f"] + targets = [bc.get_by_offset(entry.target) for entry in bc.named_exception_table] + + Preprocessor().preprocess(root) + + assert [entry.target for entry in bc.named_exception_table] == [target.offset for target in targets] + assert all( + entry.start in bc.offsets and entry.end in bc.offsets and entry.target in bc.offsets + for entry in bc.named_exception_table + ) + + +@pytest.mark.parametrize( + ("arg", "argrepr", "expected"), + [(0, "AssertionError", AssertionError), (7, "None", None), (9, "True", True), (10, "False", False), (11, "-1", -1)], +) +def test_resolve_raw_common_constant(arg, argrepr, expected): + instr = SimpleNamespace(arg=arg, argval=arg, argrepr=argrepr) + + assert _resolve_common_constant(instr) is expected + + +def test_preserve_xdis_resolved_common_constant(): + instr = SimpleNamespace(arg=9, argval=True, argrepr="True") + + assert _resolve_common_constant(instr) is True + + +def test_preprocessor_preserves_container_line_start(): + root = _preprocess( + "def f(flag):\n" + " if flag:\n" + " value = 1\n" + " else:\n" + " value = 2\n" + " result = [[0, 0], [0, 0]]\n" + " return result\n" + ) + bc = root.bytecode_lookup["f"] + folded = next(inst for inst in bc.instructions if getattr(inst, "preprocessed_container", False)) + store = bc.instructions[bc.instructions.index(folded) + 1] + + assert folded.starts_line == 6 + assert store.opname == "STORE_FAST" + assert store.starts_line is None + assert any(inst.is_jump and inst.target is folded for inst in bc.instructions) diff --git a/tests/preprocessing/test_function_defaults.py b/tests/preprocessing/test_function_defaults.py new file mode 100644 index 00000000..3eacb2e3 --- /dev/null +++ b/tests/preprocessing/test_function_defaults.py @@ -0,0 +1,144 @@ +from __future__ import annotations + +import sys +from types import SimpleNamespace + +from xdis import get_opcode + +from pylingual.editable_bytecode import EditableBytecode +from pylingual.masking.global_masker import Masker, TypeSensitiveDict +from pylingual.masking.model_disasm import create_global_masker +from pylingual.preprocessor import Preprocessor, Tracer +from pylingual.preprocessor.preprocessor import _make_function_stack_effect + + +def _preprocess(source: str) -> EditableBytecode: + code = compile(source, "", "exec") + opc = get_opcode(sys.version_info[:2], False) + bytecode = EditableBytecode(code, opc, sys.version_info[:3]) + Preprocessor().preprocess(bytecode) + return bytecode + + +def test_positional_default_tuple_shape_remains_visible_to_model(): + bytecode = _preprocess( + "def f(first=None, second=None, third=None, fourth=None):\n" + " pass\n" + ) + masker = create_global_masker(bytecode) + defaults = next( + inst + for inst in bytecode.instructions + if inst.opname == "LOAD_CONST" and inst.argval == (None, None, None, None) + ) + + assert masker.get_model_view(defaults) == "LOAD_CONST , (None, None, None, None)" + + +def test_keyword_only_default_mapping_remains_visible_to_model(): + bytecode = _preprocess( + "def f(*, retries=3, timeout=30):\n" + " pass\n" + ) + + assert any(inst.opname in ("BUILD_CONST_KEY_MAP", "BUILD_MAP") for inst in bytecode.instructions) + assert not any(inst.opname == "LOAD_CONST" and isinstance(inst.argval, dict) for inst in bytecode.instructions) + + masker = create_global_masker(bytecode) + keys = next(( + inst + for inst in bytecode.instructions + if inst.opname == "LOAD_CONST" and inst.argval == ("retries", "timeout") + ), None) + if keys is not None: + assert masker.get_model_view(keys).startswith("LOAD_CONST , (") + + +def test_inner_container_still_folds_before_unrelated_make_function(): + bytecode = _preprocess("value = ([1, 2], lambda: 0)\n") + + assert any( + inst.opname == "LOAD_CONST" + and inst.argval == [1, 2] + and getattr(inst, "preprocessed_container", False) + for inst in bytecode.instructions + ) + + +def test_call_function_kw_pads_positional_arguments(): + table = TypeSensitiveDict() + table["opts"] = "" + table["typ"] = "" + bytecode = SimpleNamespace( + version=(3, 6), + named_exception_table=None, + opcode=SimpleNamespace(MAKE_FUNCTION=-1), + resolve_namespace=lambda value: value, + ) + inst = SimpleNamespace( + argval=("opts", "typ"), + bytecode=bytecode, + has_arg=True, + is_jump_target=False, + next_instructions=[SimpleNamespace(opname="CALL_FUNCTION_KW", argval=4)], + opcode=100, + opname="LOAD_CONST", + optype="const", + ) + + assert Masker(table).get_model_view(inst) == ( + "LOAD_CONST , (, , '', '')" + ) + + +def test_pre311_make_function_effect_includes_qualname(): + assert _make_function_stack_effect(0b1110, (3, 6)) == -4 + assert _make_function_stack_effect(0b1110, (3, 12)) == -3 + + +def test_explicit_keyword_map_before_kwargs_remains_visible_to_model(): + bytecode = _preprocess( + "def f(foo, value, options):\n" + " return foo(value, conditional=True, **options)\n" + ) + function = next(bc for bc in bytecode.iter_bytecodes() if bc.codeobj.co_name == "f") + + assert any(inst.opname == "BUILD_MAP" and inst.arg == 1 for inst in function.instructions) + assert not any( + inst.opname == "LOAD_CONST" + and inst.argval == {"conditional": True} + and getattr(inst, "preprocessed_container", False) + for inst in function.instructions + ) + + +def test_dict_used_as_keyword_value_can_still_fold(): + bytecode = _preprocess( + "def f(foo, options):\n" + " return foo(value={'nested': 2}, **options)\n" + ) + function = next(bc for bc in bytecode.iter_bytecodes() if bc.codeobj.co_name == "f") + + assert any( + inst.opname == "LOAD_CONST" + and inst.argval == {"nested": 2} + and getattr(inst, "preprocessed_container", False) + for inst in function.instructions + ) + + +def test_tracer_reports_consumer_argument_from_top_of_stack(): + bytecode = _preprocess( + "def f(foo, value, options):\n" + " return foo(value, conditional=True, **options)\n" + ) + function = next(bc for bc in bytecode.iter_bytecodes() if bc.codeobj.co_name == "f") + keyword_map = next(inst for inst in function.instructions if inst.opname == "BUILD_MAP" and inst.arg == 1) + following = function.instructions[function.instructions.index(keyword_map) + 1:] + + consumption = Tracer(following).trace() + + assert consumption is not None + index, argnum = consumption + assert following[index].opname == "DICT_MERGE" + assert argnum == 1 diff --git a/tests/preprocessing/test_preprocessor.py b/tests/preprocessing/test_preprocessor.py new file mode 100644 index 00000000..b5254958 --- /dev/null +++ b/tests/preprocessing/test_preprocessor.py @@ -0,0 +1,136 @@ +from __future__ import annotations + +import ast +import sys +from pathlib import Path + +import pytest +from xdis import get_opcode + +from pylingual.editable_bytecode import EditableBytecode +from pylingual.preprocessor import Preprocessor + +FIXTURES_DIR = Path(__file__).parent / "fixtures" + + +def _extract_container_literals(source: str) -> list: + tree = ast.parse(source) + containers = [] + for node in ast.walk(tree): + if isinstance(node, (ast.List, ast.Set, ast.Tuple, ast.Dict)): + if isinstance(node, ast.Dict): + keys = [] + duplicate = False + for key_node in node.keys: + if key_node is None: + continue + try: + key = ast.literal_eval(key_node) + except (ValueError, SyntaxError): + continue + if key in keys: + duplicate = True + break + keys.append(key) + if duplicate: + continue + try: + val = ast.literal_eval(node) + except (ValueError, SyntaxError): + continue + if val is None: + continue + if not hasattr(val, "__len__"): + continue + if len(val) == 0: + continue + containers.append(val) + return containers + + +def _value_in_container(value, container) -> bool: + if isinstance(container, dict): + items = list(container.keys()) + list(container.values()) + elif isinstance(container, (list, tuple, set, frozenset)): + items = list(container) + else: + return False + for item in items: + if _values_match(value, item): + return True + if isinstance(item, (list, tuple, set, frozenset, dict)): + if _value_in_container(value, item): + return True + return False + + +def _values_match(expected, actual) -> bool: + if isinstance(expected, frozenset): + if type(actual) is set and actual == set(expected): + return True + if type(actual) is frozenset and actual == expected: + return True + return False + if type(expected) is not type(actual): + return False + return expected == actual + + +def _value_in_co_consts(value, co_consts) -> bool: + for const in co_consts: + if _values_match(value, const): + return True + if isinstance(const, (list, tuple, set, frozenset, dict)): + if _value_in_container(value, const): + return True + return False + + +def _preprocess_source(source: str) -> EditableBytecode: + code = compile(source, "", "exec") + opc = get_opcode(sys.version_info[:2], False) + bc = EditableBytecode(code, opc, sys.version_info[:3]) + Preprocessor().preprocess(bc) + return bc + + +@pytest.mark.parametrize( + "fixture", + sorted(FIXTURES_DIR.glob("*.py")), + ids=lambda p: p.stem, +) +def test_container_values_in_co_consts(fixture: Path) -> None: + source = fixture.read_text() + containers = _extract_container_literals(source) + if not containers: + pytest.skip("No non-empty container literals in fixture") + + bc = _preprocess_source(source) + + missing = [] + for value in containers: + if not _value_in_co_consts(value, bc.co_consts): + missing.append(value) + + assert not missing, ( + f"{len(missing)} container value(s) not found in co_consts:\n" + f" missing: {missing!r}\n" + f" co_consts: {bc.co_consts!r}" + ) + + +@pytest.mark.parametrize( + "source", + [ + "value = {'key': 1, 'key': 2}", + "first = object()\nsecond = object()\nvalue = {'key': first, 'key': second}", + ], +) +def test_duplicate_dict_keys_are_not_folded(source: str) -> None: + bc = _preprocess_source(source) + + assert any(inst.opname in ("BUILD_MAP", "BUILD_CONST_KEY_MAP") and inst.arg == 2 for inst in bc.instructions) + assert not any( + getattr(inst, "preprocessed_container", False) and inst.argval == {"key": 2} + for inst in bc.instructions + ) diff --git a/tests/test_editable_bytecode.py b/tests/test_editable_bytecode.py new file mode 100644 index 00000000..c156f028 --- /dev/null +++ b/tests/test_editable_bytecode.py @@ -0,0 +1,41 @@ +import copy +import sys + +import pytest +from xdis import get_opcode + +from pylingual.editable_bytecode import EditableBytecode +from pylingual.editable_bytecode.EditableBytecode import _encode_name_arg +from pylingual.utils.version import PythonVersion + + +def _bytecode(source: str) -> EditableBytecode: + code = compile(source, "", "exec") + return EditableBytecode(code, get_opcode(sys.version_info[:2], False), sys.version_info[:3]) + + +def test_insert_insts_reencodes_load_global_name_index(): + destination = _bytecode("existing = 0") + source = _bytecode("def f():\n return target\n").child_bytecodes[0] + load_global = copy.copy(next(inst for inst in source.instructions if inst.opname == "LOAD_GLOBAL")) + original_flags = load_global.arg & 1 + + destination.insert_insts({0: [load_global]}) + + name_index = destination.co_names.index("target") + assert load_global.arg == name_index << 1 | original_flags + + +@pytest.mark.parametrize( + ("opname", "name_index", "original_arg", "version", "expected"), + [ + ("LOAD_GLOBAL", 3, 2, (3, 14), 6), + ("LOAD_GLOBAL", 3, 3, (3, 14), 7), + ("LOAD_ATTR", 4, 1, (3, 14), 9), + ("LOAD_SUPER_ATTR", 5, 3, (3, 14), 23), + ("IMPORT_NAME", 6, 2, (3, 15), 26), + ("STORE_ATTR", 7, 1, (3, 14), 7), + ], +) +def test_encode_name_arg_preserves_packed_flags(opname, name_index, original_arg, version, expected): + assert _encode_name_arg(opname, name_index, original_arg, PythonVersion(version)) == expected diff --git a/tests/test_global_masker.py b/tests/test_global_masker.py new file mode 100644 index 00000000..83dd5d85 --- /dev/null +++ b/tests/test_global_masker.py @@ -0,0 +1,51 @@ +from pylingual.masking.global_masker import TypeSensitiveDict + + +def test_equal_sets_share_a_type_sensitive_key_regardless_of_order(): + values = [ + "restore_selector_state", + "restore_tree_expansions", + "create_nodes", + "restore_runtime_state", + "build_tree.descendants_init", + "build_tree.root_lookup", + "create_state_resolver", + "build_tree.tree_construct", + "create_state_handles", + "link_nodes", + "build_tree.populate_descendants", + "restore_node_runtime_state", + "build_tree.group_nodes_by_depth", + "create_runtime", + "validate_payload", + "build_tree", + "assemble_runtime_dependencies", + ] + original = set(values) + reordered = set(reversed(values)) + table = TypeSensitiveDict() + + table[original] = "" + + assert table[reordered] == "" + assert reordered in table + + +def test_equal_nested_containers_share_a_type_sensitive_key(): + original = {"phases": {"load", "build"}, "counts": [1, 2]} + reordered = {"counts": [1, 2], "phases": {"build", "load"}} + table = TypeSensitiveDict() + + table[original] = "" + + assert table[reordered] == "" + + +def test_container_keys_retain_the_original_object_for_iteration(): + original = {"load", "build"} + table = TypeSensitiveDict() + + table[original] = "" + + assert table.keys() == [original] + assert list(table.items()) == [(original, "")] diff --git a/tests/test_model_disasm.py b/tests/test_model_disasm.py new file mode 100644 index 00000000..c57a0363 --- /dev/null +++ b/tests/test_model_disasm.py @@ -0,0 +1,29 @@ +import pytest + +from pylingual.masking.global_masker import Masker +from pylingual.masking.model_disasm import restore_masked_source_text + + +@pytest.mark.parametrize("mask", ["", "''", '""']) +def test_container_masks_restore_without_wrapper_quotes(mask): + masker = Masker() + masker.global_tab[{"size": (32, 32), "crop": True}] = "" + + assert restore_masked_source_text([f"value = {mask}"], masker) == ["value = {'size': (32, 32), 'crop': True}"] + + +def test_quoted_string_mask_keeps_literal_quotes(): + masker = Masker() + masker.global_tab["it's"] = "" + + assert restore_masked_source_text(["value = ''"], masker) == ["value = 'it\\'s'"] + + +def test_container_mask_before_closing_brace_is_not_escaped(): + masker = Masker() + masker.global_tab[{"size": (32, 32), "crop": True}] = "" + + source = "value = {'first': , 'last': }" + expected = "value = {'first': {'size': (32, 32), 'crop': True}, 'last': {'size': (32, 32), 'crop': True}}" + + assert restore_masked_source_text([source], masker) == [expected] diff --git a/uv.lock b/uv.lock index 4ce907ad..2520b21b 100644 --- a/uv.lock +++ b/uv.lock @@ -742,6 +742,15 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/1e/5e/d4e9f1a599fb8e573b7b87160658329fbf28d19eac2718f51fc3def3aa5a/idna-3.18-py3-none-any.whl", hash = "sha256:7f952cbe720b688055e3f87de14f5c3e5fdaa8bc3928985c4077ca689de849a2", size = 65455, upload-time = "2026-06-02T14:34:06.319Z" }, ] +[[package]] +name = "iniconfig" +version = "2.3.0" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/72/34/14ca021ce8e5dfedc35312d08ba8bf51fdd999c576889fc2c24cb97f4f10/iniconfig-2.3.0.tar.gz", hash = "sha256:c76315c77db068650d49c5b56314774a7804df16fee4402c1f19d6d15d8c4730", size = 20503, upload-time = "2025-10-18T21:55:43.219Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/cb/b1/3846dd7f199d53cb17f49cba7e651e9ce294d8497c8c150530ed11865bb8/iniconfig-2.3.0-py3-none-any.whl", hash = "sha256:f631c04d2c48c52b84d0d0549c99ff3859c98df65b3101406327ecc7d53fbf12", size = 7484, upload-time = "2025-10-18T21:55:41.639Z" }, +] + [[package]] name = "jinja2" version = "3.1.6" @@ -1473,6 +1482,15 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/ff/6e/cf826fae916b8658848d7b9f38d88da6396895c676e8086fc0988073aaf8/pillow-12.2.0-cp314-cp314t-win_arm64.whl", hash = "sha256:aa88ccfe4e32d362816319ed727a004423aab09c5cea43c01a4b435643fa34eb", size = 2556579, upload-time = "2026-04-01T14:45:52.529Z" }, ] +[[package]] +name = "pluggy" +version = "1.6.0" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/f9/e2/3e91f31a7d2b083fe6ef3fa267035b518369d9511ffab804f839851d2779/pluggy-1.6.0.tar.gz", hash = "sha256:7dcc130b76258d33b90f61b658791dede3486c3e6bfb003ee5c9bfb396dd22f3", size = 69412, upload-time = "2025-05-15T12:30:07.975Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/54/20/4d324d65cc6d9205fabedc306948156824eb9f0ee1633355a8f7ec5c66bf/pluggy-1.6.0-py3-none-any.whl", hash = "sha256:e920276dd6813095e9377c0bc5566d94c932c33b27a3e3945d8389c374dd4746", size = 20538, upload-time = "2025-05-15T12:30:06.134Z" }, +] + [[package]] name = "propcache" version = "0.5.2" @@ -1699,6 +1717,11 @@ dependencies = [ { name = "xdis" }, ] +[package.optional-dependencies] +test = [ + { name = "pytest" }, +] + [package.metadata] requires-dist = [ { name = "asttokens" }, @@ -1710,6 +1733,7 @@ requires-dist = [ { name = "networkx" }, { name = "numpy" }, { name = "pydot" }, + { name = "pytest", marker = "extra == 'test'" }, { name = "requests" }, { name = "rich" }, { name = "seqeval" }, @@ -1721,6 +1745,7 @@ requires-dist = [ { name = "transformers", extras = ["torch"] }, { name = "xdis", git = "https://github.com/rocky/python-xdis" }, ] +provides-extras = ["test"] [[package]] name = "pyparsing" @@ -1731,6 +1756,22 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/10/bd/c038d7cc38edc1aa5bf91ab8068b63d4308c66c4c8bb3cbba7dfbc049f9c/pyparsing-3.3.2-py3-none-any.whl", hash = "sha256:850ba148bd908d7e2411587e247a1e4f0327839c40e2e5e6d05a007ecc69911d", size = 122781, upload-time = "2026-01-21T03:57:55.912Z" }, ] +[[package]] +name = "pytest" +version = "9.1.1" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "colorama", marker = "sys_platform == 'win32'" }, + { name = "iniconfig" }, + { name = "packaging" }, + { name = "pluggy" }, + { name = "pygments" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/e4/47/b9efed96c114afcfa3c9d3fe98a76a1d14c74a9e266d397cf6eb64be5e01/pytest-9.1.1.tar.gz", hash = "sha256:1088fbde8f2b49d95a549a195707afa7a76a3ce9bcadc26b6d71f0ffda5fe313", size = 1636369, upload-time = "2026-06-19T10:58:32.857Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/24/25/1de2678b631f5a49215c6c96fff41ba892b0a34df68d6d80292b1b48aa7f/pytest-9.1.1-py3-none-any.whl", hash = "sha256:37a86b45efb9a47a61a36449063e8e18d0cab3161329fc099eb21783169c4f0c", size = 386536, upload-time = "2026-06-19T10:58:31.347Z" }, +] + [[package]] name = "python-dateutil" version = "2.9.0.post0"