From fee25b6440a36eeae64131e139d87e7d6ad0772e Mon Sep 17 00:00:00 2001 From: Boyan-MILANOV Date: Tue, 26 Dec 2023 16:56:32 +0100 Subject: [PATCH] Support injecting constant arguments in function call --- fickling/fickle.py | 8 ++++++++ 1 file changed, 8 insertions(+) diff --git a/fickling/fickle.py b/fickling/fickle.py index aab3346..dc8f75a 100644 --- a/fickling/fickle.py +++ b/fickling/fickle.py @@ -517,6 +517,7 @@ def insert_magic_int(self, magic: int, index: int = -1): def insert_function_call_on_unpickled_object( self, function_definition: str, + constant_args: List[Any], ): """Insert and call a function that takes the unpickled object as parameter. @@ -552,6 +553,10 @@ def insert_function_call_on_unpickled_object( self.insert(-1, Mark()) self.insert(-1, Get.create(2)) + # Add constant arguments + for arg in constant_args: + self.insert(-1, ConstantOpcode.new(arg)) + # Now the stack contains [func, mark, model]. # We need to add TUPLE which # packs the function arguments from the stack and then call REDUCE, which calls the injected @@ -1429,6 +1434,9 @@ class ShortBinBytes(DynamicLength, ConstantOpcode): priority = Unicode.priority + 1 length_bytes = 1 + def encode_body(self) -> bytes: + return self.arg + @classmethod def validate(cls, obj): if not isinstance(obj, bytes):