21697093787

Committed 04 Feb 2026 09:56PM UTC coverage: 86.962% (-0.004%) from 86.966%

Build # 21697093787

Build Type

push

github

Committed by

web-flow

Commit Message

improve system information sent in session and container_info (#13680)

Run Details

10 of 17 new or added lines in 2 files covered. (58.82%)

222 existing lines in 17 files now uncovered.

70560 of 81139 relevant lines covered (86.96%)

0.87 hits per line

Source File
Press 'n' to go to next uncovered line, 'b' for previous

91.07

/localstack-core/localstack/state/pickle.py

"""
A small wrapper around dill that integrates with our state API, and allows registering custom serializer methods for
class hierarchies.

For your convenience, you can simply call ``dumps`` or ``loads`` as you would pickle or dill::

    from localstack.state import pickle
    foo = pickle.loads(pickle.dumps(Foo()))


You can register custom state serializers and deserializers to dill's dispatch table, but can also apply them to the
entire subclass hierarchy::

    @register(PriorityQueue, subclasses=True)
    def my_queue_pickler(pickler, obj):
        pickler.save_reduce(_recreate, (type(obj), obj.queue,), obj=obj)

    def _recreate(obj_type, obj_queue):
        # this method will be called when the object is de-serialized. you won't be able to reach it with the
        # debugger though, it's saved into the pickle! Make sure it's outside the actual reduce hook, otherwise a new
        # function is created every time for every serialized object of that type.

        q = obj_type()
        q.queue = obj_queue
        return q

To learn more about this mechanism, read https://docs.python.org/3/library/copyreg.html and
https://dill.readthedocs.io/en/latest/index.html?highlight=register#dill.Pickler.dispatch.
"""

import inspect
from collections.abc import Callable
from typing import Any, BinaryIO

import dill
from dill._dill import MetaCatchingDict

from .core import Decoder, Encoder

PythonPickler = Any
"""Type placeholder for pickle._Pickler (which has for instance the save_reduce method)"""


def register(cls: type = None, subclasses: bool = False):
    """
    Decorator to register a custom type or type tree into the dill pickling dispatcher table.

    :param cls: the type
    :param subclasses: whether to dispatch all subclasses to this function as well
    :return:
    """

    def _wrapper(fn: Any | Callable[[PythonPickler, Any], None]):
        if inspect.isclass(fn) and issubclass(fn, ObjectStateReducer):
            if cls is not None:
                raise ValueError("superfluous cls attribute for registering classes")
            obj = fn.create()
            add_dispatch_entry(obj.cls, obj._pickle, subclasses)
        elif callable(fn):
            add_dispatch_entry(cls, fn, subclasses=subclasses)
        else:
            raise ValueError(f"cannot register {fn}")

        return fn

    return _wrapper


def reducer(cls: type, restore: Callable = None, subclasses: bool = False):
    """
    Convenience decorator to simplify the following pattern::

        def _create_something(attr1, attr2):
            return Something(attr1, attr2)

        @register(Something)
        def pickle_something(pickler, obj):
            attr1 = obj.attr1
            attr2 = obj.attr2
            return pickler.save_reduce(_create_something, (attr1, attr2), obj=obj)

    into::

        def _create_something(attr1, attr2):
            return Something(attr1, attr2)

        @reducer(Something, _create_something)
        def pickle_something(pickler, obj):
            return obj.attr1, obj.attr2

    in some cases, if your constructor matches the arguments you return, into::

        @reducer(Something)
        def pickle_something(pickler, obj):
            return obj.attr1, obj.attr2

    Note that this option creates larger pickles than the previous option, since this option also needs to store the
    ``Something`` class into the pickle.

    :param cls:
    :param restore:
    :param subclasses:
    :return:
    """

    def _wrapper(fn):
        def _reducer(pickler, obj):
            return pickler.save_reduce(restore or cls, fn(obj), obj=obj)

        add_dispatch_entry(cls, _reducer, subclasses)
        return fn

    return _wrapper


def add_dispatch_entry(
    cls: type, fn: Callable[[PythonPickler, Any], None], subclasses: bool = False
):
    Pickler.dispatch_overwrite[cls] = fn
    if subclasses:
        Pickler.match_subclasses_of.add(cls)


def remove_dispatch_entry(cls: type):
    try:
        del Pickler.dispatch_overwrite[cls]
    except KeyError:
        pass

    try:
        Pickler.match_subclasses_of.remove(cls)
    except KeyError:
        pass


def dumps(obj: Any) -> bytes:
    """
    Pickle an object into bytes using a ``Encoder``.

    :param obj: the object to pickle
    :return: the pickled object
    """
    return get_default_encoder().encodes(obj)


def dump(obj: Any, file: BinaryIO):
    """
    Pickle an object into a buffer using a ``Encoder``.

    :param obj: the object to pickle
    :param file: the IO buffer
    """
    return get_default_encoder().encode(obj, file)


def loads(data: bytes) -> Any:
    """
    Unpickle am object from bytes using a ``Decoder``.

    :param data: the pickled object
    :return: the unpickled object
    """
    return get_default_decoder().decodes(data)


def load(file: BinaryIO) -> Any:
    """
    Unpickle am object from a buffer using a ``Decoder``.

    :param file: the buffer containing the pickled object
    :return: the unpickled object
    """
    return get_default_decoder().decode(file)


class _SuperclassMatchingTypeDict(MetaCatchingDict):
    """
    A special dictionary where keys are types, and keys are also optionally matched on their subclasses. Types where
    subclass matching should happen can be registered through the ``dispatch_subclasses_of`` property. Example::

        d = _SuperclassMatchingTypeDict()
        d[dict] = "a dict"
        d[defaultdict] # raises key error
        d.match_subclasses_of.add(dict)
        d[defaultdict] # returns "a dict"

    """

    def __init__(self, seq=None, match_subclasses_of: set[type] = None):
        if seq is not None:
            super().__init__(seq)
        else:
            super().__init__()

        self.match_subclasses_of = match_subclasses_of or set()

    def __missing__(self, key):
        for c in key.__mro__[1:]:
            # traverse the superclasses upwards until a dispatcher is found
            if c not in self.match_subclasses_of:
                continue

            if fn := super().get(c):
                return fn

        return super().__missing__(key)


class Pickler(dill.Pickler):
    """
    Custom dill pickler that considers dispatchers and subclass dispatchers registered via ``register``.
    """

    match_subclasses_of: set[type] = set()
    dispatch_overwrite: dict[type, Callable] = {}

    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)

        # create the dispatch table (inherit the dill dispatchers)
        dispatch = _SuperclassMatchingTypeDict(dill.Pickler.dispatch.copy())
        dispatch.update(Pickler.dispatch_overwrite.copy())  # makes sure ours take precedence
        dispatch.match_subclasses_of.update(Pickler.match_subclasses_of.copy())
        self.dispatch = dispatch


class PickleEncoder(Encoder):
    """
    An Encoder that use a dill pickling under the hood, and by default uses the custom ``Pickler`` that can be
    extended with custom serializers.
    """

    pickler_class: type[dill.Pickler]

    def __init__(self, pickler_class: type[dill.Pickler] = None):
        self.pickler_class = pickler_class or Pickler

    def encode(self, obj: Any, file: BinaryIO, py_type: type = None) -> Any:
        return self.pickler_class(file).dump(obj)


class PickleDecoder(Decoder):
    """
    A Decoder that use a dill pickling under the hood, and by default uses the custom ``Unpickler`` that can be
    extended with custom serializers.
    """

    unpickler_class: type[dill.Unpickler]

    def __init__(self, unpickler_class: type[dill.Unpickler] = None):
        self.unpickler_class = unpickler_class or dill.Unpickler

    def decode(self, file: BinaryIO, py_type=None) -> Any:
        return self.unpickler_class(file).load()


def get_default_encoder() -> Encoder:
    from .codecs import get_default_encoder

    return get_default_encoder()


def get_default_decoder() -> Decoder:
    from .codecs import get_default_decoder

    return get_default_decoder()


class ObjectStateReducer[T]:
    """
    A generalization of the following pattern::

        def _create_something(cls: Type, state: dict):
            obj = cls.__new__(self.cls)

            # do stuff on the state (perhaps re-create some attributes)
            state["this_one_doesnt_serialize"] = restore(state["this_one_serialized"])

            obj.__dict__.update(state)
            return obj

        @register(Something)
        def pickle_something(pickler, obj):
            state = obj.__dict__.copy()
            state.pop("this_one_doesnt_serialize")
            return pickler.save_reduce(_create_something, (state,), obj=obj)


    With the ObjectStateReducer, this can now be expressed as:

        @register()
        class SomethingPickler(ObjectStatePickler):
            cls = Something

            def prepare(state: dict):
                state.pop("this_one_doesnt_serialize")

            def restore(state: dict):
                state["this_one_doesnt_serialize"] = restore(state["this_one_serialized"])
    """

    cls: T

    @classmethod
    def create(cls):
        return cls()

    def register(self, subclasses=False):
        """
        Registers this ObjectStateReducer's reducer function. See ``pickle.register``.
        """
        add_dispatch_entry(self.cls, self._pickle, subclasses=subclasses)

    def _pickle(self, pickler, obj: T):
        state = self.get_state(obj)
        self.prepare(obj, state)
        return pickler.save_reduce(self._unpickle, (state,), obj=obj)

    def _unpickle(self, state: dict) -> dict:
        obj = self.cls.__new__(self.cls)
        self.restore(obj, state)
        self.set_state(obj, state)
        return obj

    def get_state(self, obj: T) -> Any:
        """
        Return the objects state. Can be overwritten by subclasses to return custom state.

        :param obj: the object
        :return: the unprepared state
        """
        return obj.__dict__.copy()

    def set_state(self, obj: T, state: Any):
        """
        Set the state of the object. Can be overwritten by subclasses to set custom state.

        :param obj: the object
        :param state: the restored object state.
        """
        obj.__dict__.update(state)

    def prepare(self, obj: T, state: Any):
        """
        Can be overwritten by subclasses to prepare the object state for pickling.

        :param obj: the object
        :param state: the object state to serialize
        """
        pass

    def restore(self, obj: T, state: Any):
        """
        Can be overwritten by subclasses to modify the object state to restore any previously removed attributes.

        :param obj: the object
        :param state: the object's state to restore
        """
        pass

1	"""
2	A small wrapper around dill that integrates with our state API, and allows registering custom serializer methods for
3	class hierarchies.
4
5	For your convenience, you can simply call ``dumps`` or ``loads`` as you would pickle or dill::
6
7	from localstack.state import pickle
8	foo = pickle.loads(pickle.dumps(Foo()))
9
10
11	You can register custom state serializers and deserializers to dill's dispatch table, but can also apply them to the
12	entire subclass hierarchy::
13
14	@register(PriorityQueue, subclasses=True)
15	def my_queue_pickler(pickler, obj):
16	pickler.save_reduce(_recreate, (type(obj), obj.queue,), obj=obj)
17
18	def _recreate(obj_type, obj_queue):
19	# this method will be called when the object is de-serialized. you won't be able to reach it with the
20	# debugger though, it's saved into the pickle! Make sure it's outside the actual reduce hook, otherwise a new
21	# function is created every time for every serialized object of that type.
22
23	q = obj_type()
24	q.queue = obj_queue
25	return q
26
27	To learn more about this mechanism, read https://docs.python.org/3/library/copyreg.html and
28	https://dill.readthedocs.io/en/latest/index.html?highlight=register#dill.Pickler.dispatch.
29	"""
30
31	import inspect	1✔
32	from collections.abc import Callable	1✔
33	from typing import Any, BinaryIO	1✔
34
35	import dill	1✔
36	from dill._dill import MetaCatchingDict	1✔
37
38	from .core import Decoder, Encoder	1✔
39
40	PythonPickler = Any	1✔
41	"""Type placeholder for pickle._Pickler (which has for instance the save_reduce method)"""	1✔
42
43
44	def register(cls: type = None, subclasses: bool = False):	1✔
45	"""
46	Decorator to register a custom type or type tree into the dill pickling dispatcher table.
47
48	:param cls: the type
49	:param subclasses: whether to dispatch all subclasses to this function as well
50	:return:
51	"""
52
53	def _wrapper(fn: Any \| Callable[[PythonPickler, Any], None]):	1✔
54	if inspect.isclass(fn) and issubclass(fn, ObjectStateReducer):	1✔
55	if cls is not None:	1✔
UNCOV 56	raise ValueError("superfluous cls attribute for registering classes")	×
57	obj = fn.create()	1✔
58	add_dispatch_entry(obj.cls, obj._pickle, subclasses)	1✔
UNCOV 59	elif callable(fn):	×
UNCOV 60	add_dispatch_entry(cls, fn, subclasses=subclasses)	×
61	else:
62	raise ValueError(f"cannot register {fn}")	×
63
64	return fn	1✔
65
66	return _wrapper	1✔
67
68
69	def reducer(cls: type, restore: Callable = None, subclasses: bool = False):	1✔
70	"""
71	Convenience decorator to simplify the following pattern::
72
73	def _create_something(attr1, attr2):
74	return Something(attr1, attr2)
75
76	@register(Something)
77	def pickle_something(pickler, obj):
78	attr1 = obj.attr1
79	attr2 = obj.attr2
80	return pickler.save_reduce(_create_something, (attr1, attr2), obj=obj)
81
82	into::
83
84	def _create_something(attr1, attr2):
85	return Something(attr1, attr2)
86
87	@reducer(Something, _create_something)
88	def pickle_something(pickler, obj):
89	return obj.attr1, obj.attr2
90
91	in some cases, if your constructor matches the arguments you return, into::
92
93	@reducer(Something)
94	def pickle_something(pickler, obj):
95	return obj.attr1, obj.attr2
96
97	Note that this option creates larger pickles than the previous option, since this option also needs to store the
98	``Something`` class into the pickle.
99
100	:param cls:
101	:param restore:
102	:param subclasses:
103	:return:
104	"""
105
106	def _wrapper(fn):	1✔
107	def _reducer(pickler, obj):	1✔
108	return pickler.save_reduce(restore or cls, fn(obj), obj=obj)	1✔
109
110	add_dispatch_entry(cls, _reducer, subclasses)	1✔
111	return fn	1✔
112
113	return _wrapper	1✔
114
115
116	def add_dispatch_entry(	1✔
117	cls: type, fn: Callable[[PythonPickler, Any], None], subclasses: bool = False
118	):
119	Pickler.dispatch_overwrite[cls] = fn	1✔
120	if subclasses:	1✔
121	Pickler.match_subclasses_of.add(cls)	1✔
122
123
124	def remove_dispatch_entry(cls: type):	1✔
125	try:	1✔
126	del Pickler.dispatch_overwrite[cls]	1✔
127	except KeyError:	1✔
128	pass	1✔
129
130	try:	1✔
131	Pickler.match_subclasses_of.remove(cls)	1✔
132	except KeyError:	1✔
133	pass	1✔
134
135
136	def dumps(obj: Any) -> bytes:	1✔
137	"""
138	Pickle an object into bytes using a ``Encoder``.
139
140	:param obj: the object to pickle
141	:return: the pickled object
142	"""
143	return get_default_encoder().encodes(obj)	1✔
144
145
146	def dump(obj: Any, file: BinaryIO):	1✔
147	"""
148	Pickle an object into a buffer using a ``Encoder``.
149
150	:param obj: the object to pickle
151	:param file: the IO buffer
152	"""
UNCOV 153	return get_default_encoder().encode(obj, file)	×
154
155
156	def loads(data: bytes) -> Any:	1✔
157	"""
158	Unpickle am object from bytes using a ``Decoder``.
159
160	:param data: the pickled object
161	:return: the unpickled object
162	"""
163	return get_default_decoder().decodes(data)	1✔
164
165
166	def load(file: BinaryIO) -> Any:	1✔
167	"""
168	Unpickle am object from a buffer using a ``Decoder``.
169
170	:param file: the buffer containing the pickled object
171	:return: the unpickled object
172	"""
UNCOV 173	return get_default_decoder().decode(file)	×
174
175
176	class _SuperclassMatchingTypeDict(MetaCatchingDict):	1✔
177	"""
178	A special dictionary where keys are types, and keys are also optionally matched on their subclasses. Types where
179	subclass matching should happen can be registered through the ``dispatch_subclasses_of`` property. Example::
180
181	d = _SuperclassMatchingTypeDict()
182	d[dict] = "a dict"
183	d[defaultdict] # raises key error
184	d.match_subclasses_of.add(dict)
185	d[defaultdict] # returns "a dict"
186
187	"""
188
189	def __init__(self, seq=None, match_subclasses_of: set[type] = None):	1✔
190	if seq is not None:	1✔
191	super().__init__(seq)	1✔
192	else:
UNCOV 193	super().__init__()	×
194
195	self.match_subclasses_of = match_subclasses_of or set()	1✔
196
197	def __missing__(self, key):	1✔
198	for c in key.__mro__[1:]:	1✔
199	# traverse the superclasses upwards until a dispatcher is found
200	if c not in self.match_subclasses_of:	1✔
201	continue	1✔
202
203	if fn := super().get(c):	1✔
204	return fn	1✔
205
206	return super().__missing__(key)	1✔
207
208
209	class Pickler(dill.Pickler):	1✔
210	"""
211	Custom dill pickler that considers dispatchers and subclass dispatchers registered via ``register``.
212	"""
213
214	match_subclasses_of: set[type] = set()	1✔
215	dispatch_overwrite: dict[type, Callable] = {}	1✔
216
217	def __init__(self, args, *kwargs):	1✔
218	super().__init__(args, *kwargs)	1✔
219
220	# create the dispatch table (inherit the dill dispatchers)
221	dispatch = _SuperclassMatchingTypeDict(dill.Pickler.dispatch.copy())	1✔
222	dispatch.update(Pickler.dispatch_overwrite.copy()) # makes sure ours take precedence	1✔
223	dispatch.match_subclasses_of.update(Pickler.match_subclasses_of.copy())	1✔
224	self.dispatch = dispatch	1✔
225
226
227	class PickleEncoder(Encoder):	1✔
228	"""
229	An Encoder that use a dill pickling under the hood, and by default uses the custom ``Pickler`` that can be
230	extended with custom serializers.
231	"""
232
233	pickler_class: type[dill.Pickler]	1✔
234
235	def __init__(self, pickler_class: type[dill.Pickler] = None):	1✔
236	self.pickler_class = pickler_class or Pickler	1✔
237
238	def encode(self, obj: Any, file: BinaryIO, py_type: type = None) -> Any:	1✔
239	return self.pickler_class(file).dump(obj)	1✔
240
241
242	class PickleDecoder(Decoder):	1✔
243	"""
244	A Decoder that use a dill pickling under the hood, and by default uses the custom ``Unpickler`` that can be
245	extended with custom serializers.
246	"""
247
248	unpickler_class: type[dill.Unpickler]	1✔
249
250	def __init__(self, unpickler_class: type[dill.Unpickler] = None):	1✔
251	self.unpickler_class = unpickler_class or dill.Unpickler	1✔
252
253	def decode(self, file: BinaryIO, py_type=None) -> Any:	1✔
254	return self.unpickler_class(file).load()	1✔
255
256
257	def get_default_encoder() -> Encoder:	1✔
258	from .codecs import get_default_encoder	1✔
259
260	return get_default_encoder()	1✔
261
262
263	def get_default_decoder() -> Decoder:	1✔
264	from .codecs import get_default_decoder	1✔
265
266	return get_default_decoder()	1✔
267
268
269	class ObjectStateReducer[T]:	1✔
270	"""
271	A generalization of the following pattern::
272
273	def _create_something(cls: Type, state: dict):
274	obj = cls.__new__(self.cls)
275
276	# do stuff on the state (perhaps re-create some attributes)
277	state["this_one_doesnt_serialize"] = restore(state["this_one_serialized"])
278
279	obj.__dict__.update(state)
280	return obj
281
282	@register(Something)
283	def pickle_something(pickler, obj):
284	state = obj.__dict__.copy()
285	state.pop("this_one_doesnt_serialize")
286	return pickler.save_reduce(_create_something, (state,), obj=obj)
287
288
289	With the ObjectStateReducer, this can now be expressed as:
290
291	@register()
292	class SomethingPickler(ObjectStatePickler):
293	cls = Something
294
295	def prepare(state: dict):
296	state.pop("this_one_doesnt_serialize")
297
298	def restore(state: dict):
299	state["this_one_doesnt_serialize"] = restore(state["this_one_serialized"])
300	"""
301
302	cls: T	1✔
303
304	@classmethod	1✔
305	def create(cls):	1✔
306	return cls()	1✔
307
308	def register(self, subclasses=False):	1✔
309	"""
310	Registers this ObjectStateReducer's reducer function. See ``pickle.register``.
311	"""
UNCOV 312	add_dispatch_entry(self.cls, self._pickle, subclasses=subclasses)	×
313
314	def _pickle(self, pickler, obj: T):	1✔
315	state = self.get_state(obj)	1✔
316	self.prepare(obj, state)	1✔
317	return pickler.save_reduce(self._unpickle, (state,), obj=obj)	1✔
318
319	def _unpickle(self, state: dict) -> dict:	1✔
320	obj = self.cls.__new__(self.cls)	1✔
321	self.restore(obj, state)	1✔
322	self.set_state(obj, state)	1✔
323	return obj	1✔
324
325	def get_state(self, obj: T) -> Any:	1✔
326	"""
327	Return the objects state. Can be overwritten by subclasses to return custom state.
328
329	:param obj: the object
330	:return: the unprepared state
331	"""
332	return obj.__dict__.copy()	1✔
333
334	def set_state(self, obj: T, state: Any):	1✔
335	"""
336	Set the state of the object. Can be overwritten by subclasses to set custom state.
337
338	:param obj: the object
339	:param state: the restored object state.
340	"""
341	obj.__dict__.update(state)	1✔
342
343	def prepare(self, obj: T, state: Any):	1✔
344	"""
345	Can be overwritten by subclasses to prepare the object state for pickling.
346
347	:param obj: the object
348	:param state: the object state to serialize
349	"""
UNCOV 350	pass	×
351
352	def restore(self, obj: T, state: Any):	1✔
353	"""
354	Can be overwritten by subclasses to modify the object state to restore any previously removed attributes.
355
356	:param obj: the object
357	:param state: the object's state to restore
358	"""
UNCOV 359	pass	×

localstack / localstack / 21697093787

Source File Press 'n' to go to next uncovered line, 'b' for previous

Source File
Press 'n' to go to next uncovered line, 'b' for previous