diff --git a/fickling/fickle.py b/fickling/fickle.py index aab3346..227d1db 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: Optional[List[Any]] = None, ): """Insert and call a function that takes the unpickled object as parameter. @@ -552,6 +553,11 @@ def insert_function_call_on_unpickled_object( self.insert(-1, Mark()) self.insert(-1, Get.create(2)) + # Add constant arguments + if constant_args: + 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 +1435,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):