diff --git a/README.md b/README.md index 7d2c9c9..62056a7 100644 --- a/README.md +++ b/README.md @@ -4,6 +4,12 @@ Python source obfuscator. Give it a `.py` file, get back a single self-contained Each build is unique by default. Pass `--seed N` if you want a reproducible one. +Version 0.8 adds a normalization layer that lowers f-strings and simple +sequence `match` statements before the rest of the pipeline. This lets literal +fragments pass through string protection and lets Abyss cover functions that +previously fell back because of generator expressions or simple match/case +syntax. + Version 0.7 hardens Abyss with a second sealed-packet layer: protected VM instructions, constants, globals, locals, and entry offsets are individually masked inside the encrypted Abyss asset, and the runtime decodes one tagged @@ -27,6 +33,20 @@ keep-name files, dry runs, stats JSON, HTML reports, manifest verification, output budget gates, a version flag, and CI metadata that matches the supported Python versions. +## 0.8 Improvements + +- F-string literals are lowered into ordinary formatting and joining calls + before string encryption, so fragments such as separators, prefixes, and + format specs no longer bypass literal protection. +- Simple sequence `match` cases are lowered into guarded `if` chains before + obfuscation, allowing the normal transforms and Abyss to process their + literal text and branch structure. +- Abyss now supports simple single-generator generator expressions and list + comprehensions by lowering them into VM-managed loops. +- The public disrobe sample now protects `rotate_name`, `compact_checksum`, + `describe_amounts`, and `build_report` under broad `--abyss` with no skipped + functions. + ## 0.7 Improvements - Abyss assets now use a second internal sealing layer after the outer asset @@ -177,6 +197,8 @@ python -m patchwork INPUT [-o OUTPUT] [options] --no-anti-debug turn off runtime anti-debug probes --abyss virtualize eligible functions into encrypted VM assets --abyss-functions NAMES virtualize only these functions, comma-separated or repeatable + --no-lower-fstrings leave f-strings as JoinedStr nodes + --no-lower-match leave match/case statements intact --audit-only analyze the input and exit without writing output --audit-json PATH write static audit metadata as JSON --manifest PATH write build manifest with hashes/options/audit data @@ -254,6 +276,8 @@ obf = Obfuscator( anti_debug=True, abyss=False, abyss_functions=[], + lower_fstrings=True, + lower_match=True, layers=3, stage2_layers=3, keep={'public_api_name'}, diff --git a/patchwork/__init__.py b/patchwork/__init__.py index fe037f5..58a14ae 100644 --- a/patchwork/__init__.py +++ b/patchwork/__init__.py @@ -1,7 +1,7 @@ from .core import Obfuscator, obfuscate, obfuscate_file from .audit import analyze_source, build_manifest, build_stats, verify_manifest -__version__ = '0.7.0' +__version__ = '0.8.0' __all__ = [ 'Obfuscator', diff --git a/patchwork/abyss.py b/patchwork/abyss.py index d0042fb..df1c0d1 100644 --- a/patchwork/abyss.py +++ b/patchwork/abyss.py @@ -222,6 +222,16 @@ class _ScopeCollector(ast.NodeVisitor): self.local_names.update(_target_names(node.target)) self.visit(node.value) + def visit_GeneratorExp(self, node: ast.GeneratorExp) -> None: + for generator in node.generators: + self.local_names.update(_target_names(generator.target)) + self.visit(generator.iter) + for condition in generator.ifs: + self.visit(condition) + self.visit(node.elt) + + visit_ListComp = visit_GeneratorExp + def finalize(self) -> tuple[set[str], set[str]]: return self.local_names - self.global_names, self.global_names @@ -234,6 +244,7 @@ class AbyssCompiler: self.global_names: set[str] = set() self.external_names: set[str] = set() self.loop_stack: list[tuple[_Label, _Label, int]] = [] + self._temp_counter = 0 def compile(self, node: ast.FunctionDef) -> dict[str, Any]: self._reject_function_shape(node) @@ -274,6 +285,12 @@ class AbyssCompiler: self.constants.append(value) return idx + def _temp(self) -> str: + self._temp_counter += 1 + name = f"_pw_ab_tmp_{self._temp_counter}" + self.local_names.add(name) + return name + def _stmt(self, node: ast.stmt) -> None: if isinstance(node, ast.Return): if node.value is None: @@ -458,6 +475,9 @@ class AbyssCompiler: if isinstance(node, ast.List): self._sequence(node.elts, "BUILD_LIST") return + if isinstance(node, (ast.GeneratorExp, ast.ListComp)): + self._comprehension(node.elt, node.generators) + return if isinstance(node, ast.Tuple): self._sequence(node.elts, "BUILD_TUPLE") return @@ -494,6 +514,34 @@ class AbyssCompiler: return raise UnsupportedAbyssNode(f"unsupported expression {type(node).__name__}") + def _comprehension(self, elt: ast.expr, generators: list[ast.comprehension]) -> None: + if len(generators) != 1: + raise UnsupportedAbyssNode("only single-generator comprehensions are supported") + generator = generators[0] + if generator.is_async: + raise UnsupportedAbyssNode("async comprehensions are not supported") + result_name = self._temp() + start_label = self.emitter.label() + end_label = self.emitter.label() + self.emitter.emit("BUILD_LIST", 0) + self.emitter.emit("STORE", result_name) + self._expr(generator.iter) + self.emitter.emit("GET_ITER") + self.emitter.mark(start_label) + self.emitter.emit("FOR_ITER", end_label) + self._store_target(generator.target) + for condition in generator.ifs: + self._expr(condition) + self.emitter.emit("JUMP_IF_FALSE", start_label) + self.emitter.emit("LOAD", result_name) + self.emitter.emit("GET_ATTR", "append") + self._expr(elt) + self.emitter.emit("CALL", 1, []) + self.emitter.emit("POP") + self.emitter.emit("JUMP", start_label) + self.emitter.mark(end_label) + self.emitter.emit("LOAD", result_name) + def _bool_op(self, node: ast.BoolOp) -> None: if not node.values: raise UnsupportedAbyssNode("empty boolean operation") diff --git a/patchwork/cli.py b/patchwork/cli.py index 2cf01b8..f941fa5 100644 --- a/patchwork/cli.py +++ b/patchwork/cli.py @@ -35,6 +35,8 @@ def _parser() -> argparse.ArgumentParser: p.add_argument("--no-anti-debug", action="store_false", dest="anti_debug", default=None) p.add_argument("--abyss", action="store_true", default=None, help="virtualize eligible functions into encrypted Abyss VM assets") p.add_argument("--abyss-functions", action="append", default=None, metavar="NAME[,NAME]", help="only virtualize the named function(s); repeatable or comma-separated") + p.add_argument("--no-lower-fstrings", action="store_false", dest="lower_fstrings", default=None) + p.add_argument("--no-lower-match", action="store_false", dest="lower_match", default=None) p.add_argument("--audit-only", action="store_true", help="analyze the input and exit without writing obfuscated output") p.add_argument("--audit-json", metavar="PATH", help="write static audit metadata as JSON") p.add_argument("--manifest", metavar="PATH", help="write build manifest with hashes, options, and audit metadata") @@ -73,6 +75,8 @@ def _effective_config(args: argparse.Namespace) -> dict[str, object]: "lazy_funcs", "anti_debug", "abyss", + "lower_fstrings", + "lower_match", "max_output_bytes", "max_ratio", "max_review_indicators", @@ -103,6 +107,8 @@ def _obfuscator_options(config: dict[str, object]) -> dict[str, object]: "anti_debug": config["anti_debug"], "abyss": config["abyss"], "abyss_functions": config["abyss_functions"], + "lower_fstrings": config["lower_fstrings"], + "lower_match": config["lower_match"], "layers": config["layers"], "stage2_layers": config["stage2_layers"], "keep": sorted(config["keep"]), diff --git a/patchwork/config.py b/patchwork/config.py index 1401c14..8f24e3b 100644 --- a/patchwork/config.py +++ b/patchwork/config.py @@ -19,6 +19,8 @@ DEFAULT_CONFIG: dict[str, Any] = { "anti_debug": True, "abyss": False, "abyss_functions": [], + "lower_fstrings": True, + "lower_match": True, "keep": [], "max_output_bytes": None, "max_ratio": None, @@ -35,6 +37,8 @@ BOOLEAN_KEYS = { "lazy_funcs", "anti_debug", "abyss", + "lower_fstrings", + "lower_match", } INTEGER_KEYS = {"seed", "layers", "stage2_layers", "max_output_bytes", "max_review_indicators"} diff --git a/patchwork/core.py b/patchwork/core.py index 1e2da93..17e419b 100644 --- a/patchwork/core.py +++ b/patchwork/core.py @@ -5,7 +5,7 @@ from .abyss import ASSETS_NAME as _ABYSS_ASSETS_NAME, DISPATCH_NAME as _ABYSS_DI from .lazy import encrypt_user_functions from .loader import build_loader from .packer import pack -from .transforms import IdentifierRenamer, NumberObfuscator, OpaquePredicateInjector, StringEncryptor, build_decrypt_helper, collect_skip_names +from .transforms import FStringLowerer, IdentifierRenamer, MatchLowerer, NumberObfuscator, OpaquePredicateInjector, StringEncryptor, build_decrypt_helper, collect_skip_names from .transforms.junk import JunkBranchInjector from .transforms.mba import MBATransformer from .transforms.opaque import build_seed_stmt @@ -35,9 +35,26 @@ def _insert_runtime_stmts(tree: ast.Module, stmts: list[ast.stmt]) -> None: insert_at += 1 tree.body[insert_at:insert_at] = stmts +def _normalize_future_imports(tree: ast.Module) -> None: + if not tree.body: + return + doc: list[ast.stmt] = [] + rest = tree.body + if ( + rest + and isinstance(rest[0], ast.Expr) + and isinstance(rest[0].value, ast.Constant) + and isinstance(rest[0].value.value, str) + ): + doc = [rest[0]] + rest = rest[1:] + future = [stmt for stmt in rest if isinstance(stmt, ast.ImportFrom) and stmt.module == '__future__'] + other = [stmt for stmt in rest if not (isinstance(stmt, ast.ImportFrom) and stmt.module == '__future__')] + tree.body = [*doc, *future, *other] + class Obfuscator: - def __init__(self, *, seed: int | None=None, rename: bool=True, encrypt_strings: bool=True, obfuscate_numbers: bool=True, opaque_predicates: bool=True, mba: bool=True, junk_branches: bool=True, lazy_funcs: bool=True, anti_debug: bool=True, layers: int=3, stage2_layers: int=3, keep: set[str] | None=None, abyss: bool=False, abyss_functions: list[str] | set[str] | tuple[str, ...] | None=None): + def __init__(self, *, seed: int | None=None, rename: bool=True, encrypt_strings: bool=True, obfuscate_numbers: bool=True, opaque_predicates: bool=True, mba: bool=True, junk_branches: bool=True, lazy_funcs: bool=True, anti_debug: bool=True, layers: int=3, stage2_layers: int=3, keep: set[str] | None=None, abyss: bool=False, abyss_functions: list[str] | set[str] | tuple[str, ...] | None=None, lower_fstrings: bool=True, lower_match: bool=True): self.rng, self.seed = make_rng(seed) self.rename = rename self.encrypt_strings = encrypt_strings @@ -52,11 +69,17 @@ class Obfuscator: self.keep_extra = set(keep or ()) self.abyss = bool(abyss) self.abyss_functions = set(abyss_functions or ()) + self.lower_fstrings = bool(lower_fstrings) + self.lower_match = bool(lower_match) self.abyss_stats: dict[str, object] = {"protected": [], "skipped": []} def obfuscate(self, source: str) -> str: tree = ast.parse(source) skip = collect_skip_names(tree) + if self.lower_fstrings: + tree = FStringLowerer().visit(tree) + if self.lower_match: + tree = MatchLowerer().visit(tree) abyss_runtime_stmts: list[ast.stmt] = [] abyss_keep: set[str] = set() if self.abyss or self.abyss_functions: @@ -98,6 +121,7 @@ class Obfuscator: tree = renamer.visit(tree) if abyss_runtime_stmts: _insert_runtime_stmts(tree, abyss_runtime_stmts) + _normalize_future_imports(tree) ast.fix_missing_locations(tree) user_code = compile(tree, '', 'exec') if self.lazy_funcs: diff --git a/patchwork/transforms/__init__.py b/patchwork/transforms/__init__.py index 44aaf91..33ef18a 100644 --- a/patchwork/transforms/__init__.py +++ b/patchwork/transforms/__init__.py @@ -1,5 +1,6 @@ -from .identifiers import IdentifierRenamer, collect_skip_names -from .strings import StringEncryptor, build_decrypt_helper -from .numbers import NumberObfuscator -from .opaque import OpaquePredicateInjector -__all__ = ['IdentifierRenamer', 'collect_skip_names', 'StringEncryptor', 'build_decrypt_helper', 'NumberObfuscator', 'OpaquePredicateInjector'] +from .identifiers import IdentifierRenamer, collect_skip_names +from .strings import StringEncryptor, build_decrypt_helper +from .numbers import NumberObfuscator +from .opaque import OpaquePredicateInjector +from .normalize import FStringLowerer, MatchLowerer +__all__ = ['IdentifierRenamer', 'collect_skip_names', 'StringEncryptor', 'build_decrypt_helper', 'NumberObfuscator', 'OpaquePredicateInjector', 'FStringLowerer', 'MatchLowerer'] diff --git a/patchwork/transforms/normalize.py b/patchwork/transforms/normalize.py new file mode 100644 index 0000000..252f944 --- /dev/null +++ b/patchwork/transforms/normalize.py @@ -0,0 +1,185 @@ +from __future__ import annotations + +import ast + + +class FStringLowerer(ast.NodeTransformer): + def visit_JoinedStr(self, node: ast.JoinedStr) -> ast.AST: + parts: list[ast.expr] = [] + for value in node.values: + if isinstance(value, ast.Constant) and isinstance(value.value, str): + if value.value: + parts.append(ast.Constant(value=value.value)) + elif isinstance(value, ast.FormattedValue): + parts.append(self._format_value(value)) + else: + return self.generic_visit(node) + if not parts: + return ast.Constant(value="") + return ast.Call( + func=ast.Attribute(value=ast.Constant(value=""), attr="join", ctx=ast.Load()), + args=[ast.List(elts=parts, ctx=ast.Load())], + keywords=[], + ) + + def _format_value(self, node: ast.FormattedValue) -> ast.expr: + value = self.visit(node.value) + if node.conversion == 115: + value = ast.Call(func=ast.Name(id="str", ctx=ast.Load()), args=[value], keywords=[]) + elif node.conversion == 114: + value = ast.Call(func=ast.Name(id="repr", ctx=ast.Load()), args=[value], keywords=[]) + elif node.conversion == 97: + value = ast.Call(func=ast.Name(id="ascii", ctx=ast.Load()), args=[value], keywords=[]) + elif node.conversion != -1: + return ast.FormattedValue(value=value, conversion=node.conversion, format_spec=node.format_spec) + if node.format_spec is None: + spec: ast.expr = ast.Constant(value="") + else: + spec_node = self.visit(node.format_spec) + if not isinstance(spec_node, ast.expr): + return ast.FormattedValue(value=value, conversion=-1, format_spec=node.format_spec) + spec = spec_node + return ast.Call(func=ast.Name(id="format", ctx=ast.Load()), args=[value, spec], keywords=[]) + + +class MatchLowerer(ast.NodeTransformer): + def __init__(self) -> None: + self._counter = 0 + + def visit_Match(self, node: ast.Match) -> ast.AST: + self.generic_visit(node) + if any(case.guard is not None for case in node.cases): + return node + subject_name = self._name() + subject_store = ast.Assign(targets=[ast.Name(id=subject_name, ctx=ast.Store())], value=node.subject) + chain = self._case_chain(subject_name, node.cases) + if chain is None: + return node + return [subject_store, chain] + + def _name(self) -> str: + self._counter += 1 + return f"_pw_match_{self._counter}" + + def _case_chain(self, subject_name: str, cases: list[ast.match_case]) -> ast.stmt | None: + next_stmt: ast.stmt | None = None + for case in reversed(cases): + lowered = self._case_if(subject_name, case, [] if next_stmt is None else [next_stmt]) + if lowered is None: + return None + next_stmt = lowered + return next_stmt + + def _case_if(self, subject_name: str, case: ast.match_case, orelse: list[ast.stmt]) -> ast.stmt | None: + result = self._pattern(subject_name, case.pattern) + if result is None: + return None + test, binders = result + if case.guard is not None: + test = ast.BoolOp(op=ast.And(), values=[test, case.guard]) + return ast.If(test=test, body=[*binders, *case.body], orelse=orelse) + + def _pattern(self, subject_name: str, pattern: ast.pattern) -> tuple[ast.expr, list[ast.stmt]] | None: + subject = ast.Name(id=subject_name, ctx=ast.Load()) + if isinstance(pattern, ast.MatchAs) and pattern.pattern is None: + if pattern.name is None: + return ast.Constant(value=True), [] + return ast.Constant(value=True), [self._bind(pattern.name, subject)] + if isinstance(pattern, ast.MatchValue): + return ast.Compare(left=subject, ops=[ast.Eq()], comparators=[pattern.value]), [] + if isinstance(pattern, ast.MatchSingleton): + return ast.Compare(left=subject, ops=[ast.Is()], comparators=[ast.Constant(value=pattern.value)]), [] + if isinstance(pattern, ast.MatchSequence): + return self._sequence_pattern(subject_name, pattern) + return None + + def _sequence_pattern(self, subject_name: str, pattern: ast.MatchSequence) -> tuple[ast.expr, list[ast.stmt]] | None: + patterns = pattern.patterns + star_index = next((index for index, item in enumerate(patterns) if isinstance(item, ast.MatchStar)), None) + tests: list[ast.expr] = [self._sequence_test(subject_name)] + binders: list[ast.stmt] = [] + if star_index is None: + tests.append(self._len_compare(subject_name, ast.Eq(), len(patterns))) + else: + tests.append(self._len_compare(subject_name, ast.GtE(), len(patterns) - 1)) + for index, item in enumerate(patterns): + if isinstance(item, ast.MatchAs): + value = self._index(subject_name, index, star_index, len(patterns)) + if item.pattern is not None: + return None + if item.name is not None: + binders.append(self._bind(item.name, value)) + elif isinstance(item, ast.MatchStar): + if item.name is not None: + binders.append(self._bind(item.name, self._star_slice(subject_name, index, len(patterns)))) + else: + value = self._index(subject_name, index, star_index, len(patterns)) + sub_name = self._name() + binders.append(ast.Assign(targets=[ast.Name(id=sub_name, ctx=ast.Store())], value=value)) + sub_result = self._pattern(sub_name, item) + if sub_result is None: + return None + sub_test, sub_binders = sub_result + tests.append(sub_test) + binders.extend(sub_binders) + return ast.BoolOp(op=ast.And(), values=tests), binders + + def _len_compare(self, subject_name: str, op: ast.cmpop, length: int) -> ast.expr: + return ast.Compare( + left=ast.Call(func=ast.Name(id="len", ctx=ast.Load()), args=[ast.Name(id=subject_name, ctx=ast.Load())], keywords=[]), + ops=[op], + comparators=[ast.Constant(value=length)], + ) + + def _sequence_test(self, subject_name: str) -> ast.expr: + subject = ast.Name(id=subject_name, ctx=ast.Load()) + sequence_type = ast.Attribute( + value=ast.Attribute( + value=ast.Call(func=ast.Name(id="__import__", ctx=ast.Load()), args=[ast.Constant(value="collections.abc")], keywords=[]), + attr="abc", + ctx=ast.Load(), + ), + attr="Sequence", + ctx=ast.Load(), + ) + return ast.BoolOp( + op=ast.And(), + values=[ + ast.Call(func=ast.Name(id="isinstance", ctx=ast.Load()), args=[subject, sequence_type], keywords=[]), + ast.UnaryOp( + op=ast.Not(), + operand=ast.Call( + func=ast.Name(id="isinstance", ctx=ast.Load()), + args=[ + subject, + ast.Tuple( + elts=[ + ast.Name(id="str", ctx=ast.Load()), + ast.Name(id="bytes", ctx=ast.Load()), + ast.Name(id="bytearray", ctx=ast.Load()), + ], + ctx=ast.Load(), + ), + ], + keywords=[], + ), + ), + ], + ) + + def _index(self, subject_name: str, index: int, star_index: int | None, total: int) -> ast.expr: + if star_index is not None and index > star_index: + index = index - total + return ast.Subscript(value=ast.Name(id=subject_name, ctx=ast.Load()), slice=ast.Constant(value=index), ctx=ast.Load()) + + def _star_slice(self, subject_name: str, index: int, total: int) -> ast.expr: + stop = index - total + 1 + upper: ast.expr | None = None if stop == 0 else ast.Constant(value=stop) + return ast.Subscript( + value=ast.Call(func=ast.Name(id="list", ctx=ast.Load()), args=[ast.Name(id=subject_name, ctx=ast.Load())], keywords=[]), + slice=ast.Slice(lower=ast.Constant(value=index), upper=upper, step=None), + ctx=ast.Load(), + ) + + def _bind(self, name: str, value: ast.expr) -> ast.stmt: + return ast.Assign(targets=[ast.Name(id=name, ctx=ast.Store())], value=value) diff --git a/pyproject.toml b/pyproject.toml index 7dd7915..7419899 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta" [project] name = "patchwork" -version = "0.7.0" +version = "0.8.0" description = "Python source transformation tool with reproducible builds, static audit metadata, and manifest generation" requires-python = ">=3.10" license = {text = "MIT"} diff --git a/tests/test_abyss.py b/tests/test_abyss.py index b47ae8a..f4699e0 100644 --- a/tests/test_abyss.py +++ b/tests/test_abyss.py @@ -153,6 +153,74 @@ print(secret(9, 5)) self.assertEqual(0, result.returncode, result.stderr) self.assertEqual(run_file(src_path), run_file(out_path)) + def test_public_sample_all_private_functions_are_abyss_supported(self) -> None: + source = """ +from __future__ import annotations +from dataclasses import dataclass +SALT = b"patchwork-disrobe-public-sample-v1" +ALPHABET = "abcdefghijklmnopqrstuvwxyz" +@dataclass(frozen=True) +class Ticket: + owner: str + amounts: tuple[int, ...] +def rotate_name(text: str, offset: int) -> str: + rotated: list[str] = [] + for char in text.lower(): + if char in ALPHABET: + index = (ALPHABET.index(char) + offset) % len(ALPHABET) + rotated.append(ALPHABET[index]) + else: + rotated.append(char) + return "".join(rotated) +def compact_checksum(ticket: Ticket) -> int: + state = 0x2A17 + payload = f"{ticket.owner}:{','.join(str(item) for item in ticket.amounts)}".encode() + for position, byte in enumerate(SALT + payload, start=1): + state ^= byte * position + state = ((state << 5) | (state >> 11)) & 0xFFFF + state = (state + 0x31D + len(ticket.amounts)) & 0xFFFF + return state +def describe_amounts(amounts: tuple[int, ...]) -> str: + match amounts: + case (single,): + return f"single={single}" + case (first, second): + return f"pair={first + second}" + case (first, *middle, last): + return f"range={first}->{last};mid={sum(middle)}" + case _: + return "empty" +def build_report() -> str: + tickets = ( + Ticket("Ada", (3, 5, 8, 13)), + Ticket("Grace", (21, 34)), + Ticket("Linus", (55,)), + Ticket("Katherine", (89, 144, 233)), + ) + lines: list[str] = [] + for index, ticket in enumerate(tickets, start=1): + alias = rotate_name(ticket.owner, index + 2) + check = compact_checksum(ticket) + lines.append(f"{index:02d}|{alias}|{describe_amounts(ticket.amounts)}|{check:04x}") + return "\\n".join(lines) +if __name__ == "__main__": + print(build_report()) +""" + obfuscator = Obfuscator(seed=8080, abyss=True) + obfuscated = obfuscator.obfuscate(source) + self.assertEqual(["rotate_name", "compact_checksum", "describe_amounts", "build_report"], obfuscator.abyss_stats["protected"]) + self.assertEqual([], obfuscator.abyss_stats["skipped"]) + self.assertNotIn("patchwork-disrobe-public-sample-v1", obfuscated) + self.assertNotIn("single=", obfuscated) + + with tempfile.TemporaryDirectory() as raw: + directory = Path(raw) + src_path = directory / "public_sample.py" + obf_path = directory / "public_sample_obf.py" + src_path.write_text(source, encoding="utf-8") + obf_path.write_text(obfuscated, encoding="utf-8") + self.assertEqual(run_file(src_path), run_file(obf_path)) + if __name__ == "__main__": unittest.main() diff --git a/tests/test_audit.py b/tests/test_audit.py index 024a5a7..4aab909 100644 --- a/tests/test_audit.py +++ b/tests/test_audit.py @@ -68,7 +68,7 @@ class AuditTests(unittest.TestCase): result = self.run_cli("--version") self.assertEqual(0, result.returncode) - self.assertIn("patchwork 0.7.0", result.stdout) + self.assertIn("patchwork 0.8.0", result.stdout) def test_config_dump_and_keep_file_merge_options(self) -> None: with tempfile.TemporaryDirectory() as raw: diff --git a/tests/test_transforms.py b/tests/test_transforms.py index f59f579..ab32115 100644 --- a/tests/test_transforms.py +++ b/tests/test_transforms.py @@ -12,6 +12,7 @@ ROOT = Path(__file__).resolve().parent.parent sys.path.insert(0, str(ROOT)) from patchwork.transforms.identifiers import IdentifierRenamer +from patchwork.transforms.normalize import MatchLowerer from patchwork.transforms.numbers import NumberObfuscator from patchwork.transforms.strings import StringEncryptor, build_decrypt_helper @@ -21,6 +22,12 @@ def eval_ast_expr(node: ast.AST) -> object: ast.fix_missing_locations(expr) return eval(compile(expr, "", "eval"), {}) +def run_module(source: str) -> str: + stdout = io.StringIO() + with contextlib.redirect_stdout(stdout): + exec(compile(source, "", "exec"), {}) + return stdout.getvalue() + class TransformTests(unittest.TestCase): def test_new_number_transforms_preserve_values(self) -> None: @@ -70,6 +77,24 @@ class TransformTests(unittest.TestCase): exec(compile(renamed, "", "exec"), {}) self.assertIn("deque([1, 2, 3])", stdout.getvalue()) + def test_match_lowerer_preserves_general_sequence_patterns(self) -> None: + source = """ +from collections import deque +def classify(value): + match value: + case (first, *middle, last): + return f"{type(value).__name__}:{first}:{middle}:{last}" + case _: + return f"{type(value).__name__}:no" +print(classify(range(4))) +print(classify("abcd")) +print(classify(deque([1, 2, 3]))) +""" + tree = MatchLowerer().visit(ast.parse(source)) + ast.fix_missing_locations(tree) + + self.assertEqual(run_module(source), run_module(ast.unparse(tree))) + if __name__ == "__main__": unittest.main()