deser_pickle/
types.rs

1//! Globals, references and the information passed on out of band.
2use alloc::borrow::Cow;
3use alloc::string::String;
4use core::fmt;
5
6use deser_core::State;
7use deser_core::de::{Deserialize, OwnedSink, Sink, SinkHandle, Slot, default_atom};
8use deser_core::ext::{ExtValue, Extension};
9use deser_core::ser::{Describe, Emit, Serialize};
10use deser_core::{Atom, ContainerShape, Error, Text};
11
12/// A global: a class or function referred to by its module and name.
13///
14/// Pickles refer to classes and functions by name (with the `GLOBAL` and
15/// `STACK_GLOBAL` opcodes), they are never imported.  A global is the
16/// class of an [object](crate#objects) (see [`take_class`]) and globals
17/// that are values themselves (a pickled class) are passed on as extension
18/// atoms of this type.  The fallback of the extension is the dotted path
19/// (`module.name`) as string:
20///
21/// ```
22/// use deser_pickle::Global;
23///
24/// // `pickle.dumps(collections.OrderedDict, 4)`
25/// let input = b"\x80\x04\x95\x1f\x00\x00\x00\x00\x00\x00\x00\x8c\x0bcollections\x94\x8c\x0bOrderedDict\x94\x93\x94.";
26/// let global: Global = deser_pickle::from_slice(input).unwrap();
27/// assert_eq!(global.module(), "collections");
28/// assert_eq!(global.name(), "OrderedDict");
29/// let path: String = deser_pickle::from_slice(input).unwrap();
30/// assert_eq!(path, "collections.OrderedDict");
31/// ```
32#[derive(Clone, PartialEq, Eq, PartialOrd, Ord, Hash)]
33pub struct Global {
34    /// `module.name`
35    path: String,
36    /// The length of the module in `path`.
37    split: usize,
38}
39
40impl Global {
41    /// Creates a global from its module and name.
42    pub fn new(module: &str, name: &str) -> Global {
43        let mut path = String::with_capacity(module.len() + name.len() + 1);
44        path.push_str(module);
45        path.push('.');
46        path.push_str(name);
47        Global {
48            path,
49            split: module.len(),
50        }
51    }
52
53    /// Returns the module (like `collections`).
54    pub fn module(&self) -> &str {
55        &self.path[..self.split]
56    }
57
58    /// Returns the name in the module (like `OrderedDict`).
59    ///
60    /// Names of nested classes contain dots (`Outer.Inner`).
61    pub fn name(&self) -> &str {
62        &self.path[self.split + 1..]
63    }
64
65    /// Returns the dotted path (`module.name`).
66    pub fn path(&self) -> &str {
67        &self.path
68    }
69}
70
71impl fmt::Debug for Global {
72    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
73        write!(f, "Global({})", self.path)
74    }
75}
76
77impl fmt::Display for Global {
78    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
79        f.write_str(&self.path)
80    }
81}
82
83impl Extension for Global {
84    fn name(&self) -> &str {
85        "python global"
86    }
87
88    fn fallback(&self) -> Atom<'_> {
89        Atom::Str(Text::borrowed(&self.path))
90    }
91}
92
93impl Serialize for Global {
94    fn serialize<'a>(value: &'a Self, _state: &mut State) -> Result<Emit<'a>, Error> {
95        Ok(Emit::Atom(Atom::Ext(ExtValue::borrowed(value))))
96    }
97}
98
99impl<'de> Deserialize<'de> for Global {
100    fn deserialize_atom(slot: &mut Slot<Self>, atom: Atom, state: &mut State) -> Result<(), Error> {
101        match atom {
102            Atom::Ext(ref ext) => match ext.downcast_ref::<Global>() {
103                Some(value) => {
104                    slot.set(value.clone());
105                    Ok(())
106                }
107                None => default_atom(slot, atom, state),
108            },
109            other => default_atom(slot, other, state),
110        }
111    }
112
113    fn expecting() -> Cow<'static, str> {
114        Cow::Borrowed("python global")
115    }
116}
117
118/// A reference to a value that contains it (a cycle).
119///
120/// Pickles are graphs: a value can be reached more than once and can
121/// contain itself.  Values that are reached more than once are emitted at
122/// every place (see [References](crate#references)), with an id as event
123/// data (see [`take_shared_id`]).  Where a value is reached again from
124/// within itself it cannot be emitted again, a reference with the id of the
125/// value is emitted instead.
126///
127/// The fallback of the extension is `null`: types that do not understand
128/// references see them as missing values (an `Option` is `None`).
129///
130/// ```
131/// use deser_pickle::Reference;
132///
133/// // `x = []; x.append(x)`
134/// let input = b"\x80\x04\x95\x06\x00\x00\x00\x00\x00\x00\x00]\x94h\x00a.";
135/// let value: Vec<Reference> = deser_pickle::from_slice(input).unwrap();
136/// assert_eq!(value.len(), 1);
137/// let value: Vec<Option<u32>> = deser_pickle::from_slice(input).unwrap();
138/// assert_eq!(value, [None]);
139/// ```
140#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash)]
141pub struct Reference {
142    id: u64,
143}
144
145impl Reference {
146    /// Creates a reference to the value with the given id.
147    pub const fn new(id: u64) -> Reference {
148        Reference { id }
149    }
150
151    /// Returns the id of the value it refers to (see [`take_shared_id`]).
152    pub const fn id(self) -> u64 {
153        self.id
154    }
155}
156
157impl Extension for Reference {
158    fn name(&self) -> &str {
159        "pickle reference"
160    }
161
162    fn fallback(&self) -> Atom<'_> {
163        Atom::Null
164    }
165}
166
167impl Serialize for Reference {
168    fn serialize<'a>(value: &'a Self, _state: &mut State) -> Result<Emit<'a>, Error> {
169        Ok(Emit::Atom(Atom::Ext(ExtValue::borrowed(value))))
170    }
171}
172
173impl<'de> Deserialize<'de> for Reference {
174    fn deserialize_atom(slot: &mut Slot<Self>, atom: Atom, state: &mut State) -> Result<(), Error> {
175        match atom {
176            Atom::Ext(ref ext) => match ext.downcast_ref::<Reference>() {
177                Some(&value) => {
178                    slot.set(value);
179                    Ok(())
180                }
181                None => default_atom(slot, atom, state),
182            },
183            other => default_atom(slot, other, state),
184        }
185    }
186
187    fn expecting() -> Cow<'static, str> {
188        Cow::Borrowed("pickle reference")
189    }
190}
191
192/// The Python type of a value where the data model does not tell.
193///
194/// Tuples, sets and frozensets are sequences and bytearrays are bytes.
195/// The deserializer attaches the kind to their first event (see
196/// [`take_kind`]), the serializer writes these types if it's set (see
197/// [`set_kind`]).  Lists, dicts and bytes have no kind.
198#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash)]
199pub enum Kind {
200    /// A `tuple` (a sequence).
201    Tuple,
202    /// A `set` (a sequence).
203    Set,
204    /// A `frozenset` (a sequence).
205    FrozenSet,
206    /// A `bytearray` (bytes).
207    ByteArray,
208}
209
210/// How an object is created from the value it's emitted as.
211///
212/// The deserializer attaches the form to the first event of an
213/// [object](crate#objects) together with its class (see [`take_form`]).
214/// The serializer creates objects in this form (see [`set_form`]) which
215/// makes them read back as they were.  Without a form, objects that are
216/// maps are created from their state, lists from their items, tuples from
217/// their arguments and other values are the argument.
218#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash)]
219pub enum Form {
220    /// The value is the state: `cls.__new__(cls)` and `__setstate__` (or
221    /// the attributes for maps).  This is how Python pickles instances of
222    /// classes (including dataclasses).
223    State,
224    /// The value is the state of an object with slots: `cls.__new__(cls)`
225    /// and the entries of the map are set as attributes.
226    Slots,
227    /// The value is the items: `cls.__new__(cls)` and `__setitem__` for
228    /// maps, `extend` (or `append`) for sequences.  This is how Python
229    /// pickles subclasses of `dict` and `list`.
230    Items,
231    /// The value is the arguments: `cls(*value)` for sequences and
232    /// `cls.__new__(cls, **value)` for maps (which requires protocol 4).
233    Arguments,
234    /// The value is the only argument: `cls(value)`.
235    Argument,
236}
237
238/// The class of a value, attached as event data to its first event.
239#[derive(Debug, Default, Clone)]
240pub(crate) struct ClassData(pub(crate) Option<Global>);
241
242/// The form of an object, attached as event data to its first event.
243#[derive(Debug, Default, Clone)]
244pub(crate) struct FormData(pub(crate) Option<Form>);
245
246/// The kind of a value, attached as event data to its first event.
247#[derive(Debug, Default, Clone)]
248pub(crate) struct KindData(pub(crate) Option<Kind>);
249
250/// The id of a value that is reached more than once.
251#[derive(Debug, Default, Clone)]
252pub(crate) struct SharedIdData(pub(crate) Option<u64>);
253
254/// Takes the class of the current value from the state.
255///
256/// This is intended to be called by sinks from within [`Sink::atom`],
257/// [`Sink::map`] or [`Sink::seq`].  Returns `None` if the value is not an
258/// [object](crate#objects) or the data format is not pickle.
259pub fn take_class(state: &mut State) -> Option<Global> {
260    state.take_event::<ClassData>().and_then(|class| class.0)
261}
262
263/// Sets the class of the value that is serialized.
264///
265/// It must be called from [`Serialize::serialize`] and applies to the
266/// value serialized from that call (see [`State::event`]).  The value is
267/// written as an object of the class (see
268/// [Serialization](crate#serialization)).  Serializers of other formats
269/// ignore it.
270pub fn set_class(state: &mut State, class: Global) {
271    state.event_mut::<ClassData>().0 = Some(class);
272}
273
274/// Takes the form of the current object from the state.
275///
276/// Returns `None` for values that are not objects.  See [`Form`].
277pub fn take_form(state: &mut State) -> Option<Form> {
278    state.take_event::<FormData>().and_then(|form| form.0)
279}
280
281/// Sets the form of the object that is serialized.
282///
283/// It must be called from [`Serialize::serialize`] together with
284/// [`set_class`].  See [`Form`].
285pub fn set_form(state: &mut State, form: Form) {
286    state.event_mut::<FormData>().0 = Some(form);
287}
288
289/// Takes the kind of the current value from the state.
290///
291/// Returns `None` for values that have no kind (see [`Kind`]).
292pub fn take_kind(state: &mut State) -> Option<Kind> {
293    state.take_event::<KindData>().and_then(|kind| kind.0)
294}
295
296/// Sets the kind of the value that is serialized.
297///
298/// It must be called from [`Serialize::serialize`].  Sequences are written
299/// as tuples, sets or frozensets and bytes as bytearrays.
300pub fn set_kind(state: &mut State, kind: Kind) {
301    state.event_mut::<KindData>().0 = Some(kind);
302}
303
304/// Takes the id of the current value if it's reached more than once.
305///
306/// Values that are reached more than once (shared values) are emitted at
307/// every place with the same id.  Returns `None` for values that are only
308/// reached once.  See [References](crate#references).
309pub fn take_shared_id(state: &mut State) -> Option<u64> {
310    state.take_event::<SharedIdData>().and_then(|id| id.0)
311}
312
313/// Sets the id of the value that is serialized.
314///
315/// The serializer writes a value with an id once: other values with the
316/// same id are written as references to the first and a [`Reference`] with
317/// the id refers to it.  It must be called from [`Serialize::serialize`].
318pub fn set_shared_id(state: &mut State, id: u64) {
319    state.event_mut::<SharedIdData>().0 = Some(id);
320}
321
322/// A value with the class of a Python object.
323///
324/// When deserialized the class of the value (if there is one) is captured,
325/// when serialized the value is written as an object of the class:
326///
327/// ```
328/// use std::collections::BTreeMap;
329/// use deser_pickle::{Global, Object};
330///
331/// // a `Point(x=1, y=2)` of a module `geometry`
332/// let input = b"\x80\x04\x95*\x00\x00\x00\x00\x00\x00\x00\x8c\x08geometry\x94\x8c\x05Point\x94\x93\x94)\x81\x94}\x94(\x8c\x01x\x94K\x01\x8c\x01y\x94K\x02ub.";
333/// let point: Object<BTreeMap<String, i32>> = deser_pickle::from_slice(input).unwrap();
334/// assert_eq!(point.class, Some(Global::new("geometry", "Point")));
335/// assert_eq!(point.value["y"], 2);
336/// ```
337///
338/// Other formats do not have classes.  When deserializing from such
339/// formats, the class is `None`, when serializing it's ignored.
340#[derive(Clone, PartialEq, Eq, PartialOrd, Ord, Hash, Default)]
341pub struct Object<T> {
342    /// The class of the value.
343    pub class: Option<Global>,
344    /// How the object is created from the value (see [`Form`]).
345    pub form: Option<Form>,
346    /// The value.
347    pub value: T,
348}
349
350impl<T> Object<T> {
351    /// Creates a value with a class.
352    pub fn new(class: Global, value: T) -> Object<T> {
353        Object {
354            class: Some(class),
355            form: None,
356            value,
357        }
358    }
359
360    /// Creates a value with a class and the form the object is created in.
361    pub fn with_form(class: Global, form: Form, value: T) -> Object<T> {
362        Object {
363            class: Some(class),
364            form: Some(form),
365            value,
366        }
367    }
368
369    /// Returns the inner value.
370    pub fn into_inner(self) -> T {
371        self.value
372    }
373}
374
375impl<T: fmt::Debug> fmt::Debug for Object<T> {
376    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
377        match self.class {
378            Some(ref class) => {
379                write!(f, "{} ", class)?;
380                fmt::Debug::fmt(&self.value, f)
381            }
382            None => fmt::Debug::fmt(&self.value, f),
383        }
384    }
385}
386
387impl<T: Serialize> Serialize for Object<T> {
388    fn serialize<'a>(this: &'a Self, state: &mut State) -> Result<Emit<'a>, Error> {
389        // the class is set after the value attached its data, it replaces it
390        let emit = T::serialize(&this.value, state)?;
391        if let Some(ref class) = this.class {
392            set_class(state, class.clone());
393            if let Some(form) = this.form {
394                set_form(state, form);
395            }
396        }
397        Ok(emit)
398    }
399
400    fn finish(this: &Self, state: &mut State) -> Result<(), Error> {
401        T::finish(&this.value, state)
402    }
403
404    fn is_optional(this: &Self) -> bool {
405        T::is_optional(&this.value)
406    }
407
408    fn container_shape(this: &Self) -> ContainerShape {
409        T::container_shape(&this.value)
410    }
411
412    fn describe(this: &Self, d: &mut dyn Describe) {
413        T::describe(&this.value, d)
414    }
415}
416
417impl<'de, T: Deserialize<'de>> Deserialize<'de> for Object<T> {
418    fn deserialize_into<'out>(
419        out: &'out mut Option<Self>,
420        state: &mut State,
421    ) -> SinkHandle<'out, 'de> {
422        SinkHandle::arena(
423            ObjectSink {
424                out,
425                slot: None,
426                compound: None,
427                class: None,
428                form: None,
429            },
430            state,
431        )
432    }
433
434    fn expecting() -> Cow<'static, str> {
435        T::expecting()
436    }
437
438    fn describe_type(d: &mut dyn Describe) {
439        T::describe_type(d)
440    }
441}
442
443struct ObjectSink<'a, 'de, T> {
444    out: &'a mut Option<Object<T>>,
445    // atoms are deserialized directly into this slot, maps and sequences
446    // need a sink that lives across calls
447    slot: Option<T>,
448    compound: Option<OwnedSink<'de, T>>,
449    class: Option<Global>,
450    form: Option<Form>,
451}
452
453impl<'a, 'de, T: Deserialize<'de>> ObjectSink<'a, 'de, T> {
454    /// Takes the class and form of the value.
455    fn take(&mut self, state: &mut State) {
456        self.class = take_class(state);
457        self.form = take_form(state);
458    }
459
460    fn compound(&mut self, state: &mut State) -> &mut dyn Sink<'de> {
461        self.compound
462            .get_or_insert_with(|| OwnedSink::deserialize(state))
463            .get_mut()
464    }
465}
466
467impl<'a, 'de, T: Deserialize<'de>> Sink<'de> for ObjectSink<'a, 'de, T> {
468    fn atom(&mut self, atom: Atom, state: &mut State) -> Result<(), Error> {
469        self.take(state);
470        let mut sink = T::deserialize_into(&mut self.slot, state);
471        sink.atom(atom, state)?;
472        sink.finish(state)
473    }
474
475    fn borrowed_atom(&mut self, atom: Atom<'de>, state: &mut State) -> Result<(), Error> {
476        self.take(state);
477        let mut sink = T::deserialize_into(&mut self.slot, state);
478        sink.borrowed_atom(atom, state)?;
479        sink.finish(state)
480    }
481
482    fn map(&mut self, state: &mut State) -> Result<(), Error> {
483        self.take(state);
484        self.compound(state).map(state)
485    }
486
487    fn seq(&mut self, state: &mut State) -> Result<(), Error> {
488        self.take(state);
489        self.compound(state).seq(state)
490    }
491
492    fn next_key(&mut self, state: &mut State) -> Result<SinkHandle<'_, 'de>, Error> {
493        self.compound(state).next_key(state)
494    }
495
496    fn next_value(&mut self, state: &mut State) -> Result<SinkHandle<'_, 'de>, Error> {
497        self.compound(state).next_value(state)
498    }
499
500    fn value_for_key(
501        &mut self,
502        key: &str,
503        state: &mut State,
504    ) -> Result<Option<SinkHandle<'_, 'de>>, Error> {
505        self.compound(state).value_for_key(key, state)
506    }
507
508    fn recover(&mut self, err: Error, state: &mut State) -> Result<(), Error> {
509        match self.compound {
510            Some(ref mut compound) => compound.get_mut().recover(err, state),
511            None => Err(err),
512        }
513    }
514
515    fn finish(&mut self, state: &mut State) -> Result<(), Error> {
516        let value = match self.compound {
517            Some(ref mut compound) => {
518                compound.get_mut().finish(state)?;
519                compound.take()
520            }
521            None => self.slot.take(),
522        };
523        let class = self.class.take();
524        let form = self.form.take();
525        *self.out = value.map(|value| Object { class, form, value });
526        Ok(())
527    }
528
529    fn expecting(&self) -> Cow<'_, str> {
530        if let Some(ref compound) = self.compound {
531            return compound.get().expecting();
532        }
533        T::expecting()
534    }
535}
deser_pickle/types.rs

deser_pickle/
types.rs
