pyc_editor/
lib.rs

1pub mod error;
2pub mod prelude;
3pub mod traits;
4pub mod utils;
5#[cfg(feature = "v310")]
6pub mod v310;
7#[cfg(feature = "v311")]
8pub mod v311;
9#[cfg(feature = "v312")]
10pub mod v312;
11#[cfg(feature = "v313")]
12pub mod v313;
13
14use error::Error;
15use python_marshal::{self, magic::PyVersion, minimize_references};
16use std::io::Read;
17
18#[derive(Debug, Clone)]
19pub enum PycFile {
20    #[cfg(feature = "v310")]
21    V310(v310::code_objects::Pyc),
22    #[cfg(feature = "v311")]
23    V311(v311::code_objects::Pyc),
24    #[cfg(feature = "v312")]
25    V312(v312::code_objects::Pyc),
26    #[cfg(feature = "v313")]
27    V313(v313::code_objects::Pyc),
28}
29
30impl From<PycFile> for python_marshal::PycFile {
31    fn from(val: PycFile) -> Self {
32        match val.clone() {
33            #[cfg(feature = "v310")]
34            PycFile::V310(pyc) => {
35                python_marshal::PycFile {
36                    python_version: pyc.python_version,
37                    timestamp: Some(pyc.timestamp),
38                    hash: pyc.hash,
39                    object: python_marshal::Object::Code(pyc.code_object.into()),
40                    references: Vec::new(), // All references are resolved in this editor.
41                }
42            }
43            #[cfg(feature = "v311")]
44            PycFile::V311(pyc) => {
45                python_marshal::PycFile {
46                    python_version: pyc.python_version,
47                    timestamp: Some(pyc.timestamp),
48                    hash: pyc.hash,
49                    object: python_marshal::Object::Code(pyc.code_object.into()),
50                    references: Vec::new(), // All references are resolved in this editor.
51                }
52            }
53            #[cfg(feature = "v312")]
54            PycFile::V312(pyc) => {
55                python_marshal::PycFile {
56                    python_version: pyc.python_version,
57                    timestamp: Some(pyc.timestamp),
58                    hash: pyc.hash,
59                    object: python_marshal::Object::Code(pyc.code_object.into()),
60                    references: Vec::new(), // All references are resolved in this editor.
61                }
62            }
63            #[cfg(feature = "v313")]
64            PycFile::V313(pyc) => {
65                python_marshal::PycFile {
66                    python_version: pyc.python_version,
67                    timestamp: Some(pyc.timestamp),
68                    hash: pyc.hash,
69                    object: python_marshal::Object::Code(pyc.code_object.into()),
70                    references: Vec::new(), // All references are resolved in this editor.
71                }
72            }
73        }
74    }
75}
76
77#[derive(Debug, Clone, PartialEq)]
78pub enum CodeObject {
79    #[cfg(feature = "v310")]
80    V310(v310::code_objects::Code),
81    #[cfg(feature = "v311")]
82    V311(v311::code_objects::Code),
83    #[cfg(feature = "v312")]
84    V312(v312::code_objects::Code),
85    #[cfg(feature = "v313")]
86    V313(v313::code_objects::Code),
87}
88
89pub fn load_pyc(data: impl Read) -> Result<PycFile, Error> {
90    let pyc_file = python_marshal::load_pyc(data)?;
91
92    match pyc_file.python_version {
93        #[cfg(feature = "v310")]
94        PyVersion {
95            major: 3,
96            minor: 10,
97            ..
98        } => {
99            let pyc = v310::code_objects::Pyc::try_from(pyc_file)?;
100            Ok(PycFile::V310(pyc))
101        }
102        #[cfg(feature = "v311")]
103        PyVersion {
104            major: 3,
105            minor: 11,
106            ..
107        } => {
108            let pyc = v311::code_objects::Pyc::try_from(pyc_file)?;
109            Ok(PycFile::V311(pyc))
110        }
111        #[cfg(feature = "v312")]
112        PyVersion {
113            major: 3,
114            minor: 12,
115            ..
116        } => {
117            let pyc = v312::code_objects::Pyc::try_from(pyc_file)?;
118            Ok(PycFile::V312(pyc))
119        }
120        #[cfg(feature = "v313")]
121        PyVersion {
122            major: 3,
123            minor: 13,
124            ..
125        } => {
126            let pyc = v313::code_objects::Pyc::try_from(pyc_file)?;
127            Ok(PycFile::V313(pyc))
128        }
129        _ => Err(Error::UnsupportedVersion(pyc_file.python_version)),
130    }
131}
132
133pub fn dump_pyc(pyc_file: PycFile) -> Result<Vec<u8>, Error> {
134    let mut pyc: python_marshal::PycFile = pyc_file.into();
135
136    let (obj, refs) = minimize_references(&pyc.object, pyc.references);
137
138    pyc.object = obj;
139    pyc.references = refs;
140
141    python_marshal::dump_pyc(pyc).map_err(|e| e.into())
142}
143
144pub fn load_code(mut data: impl Read, python_version: PyVersion) -> Result<CodeObject, Error> {
145    match python_version {
146        #[cfg(feature = "v310")]
147        PyVersion {
148            major: 3,
149            minor: 10,
150            ..
151        } => {
152            let mut buf = Vec::new();
153            data.read_to_end(&mut buf)?;
154            let code = python_marshal::load_bytes(&buf, python_version)?;
155            Ok(CodeObject::V310(code.try_into()?))
156        }
157        #[cfg(feature = "v311")]
158        PyVersion {
159            major: 3,
160            minor: 11,
161            ..
162        } => {
163            let mut buf = Vec::new();
164            data.read_to_end(&mut buf)?;
165            let code = python_marshal::load_bytes(&buf, python_version)?;
166            Ok(CodeObject::V311(code.try_into()?))
167        }
168        #[cfg(feature = "v312")]
169        PyVersion {
170            major: 3,
171            minor: 12,
172            ..
173        } => {
174            let mut buf = Vec::new();
175            data.read_to_end(&mut buf)?;
176            let code = python_marshal::load_bytes(&buf, python_version)?;
177            Ok(CodeObject::V312(code.try_into()?))
178        }
179        #[cfg(feature = "v313")]
180        PyVersion {
181            major: 3,
182            minor: 13,
183            ..
184        } => {
185            let mut buf = Vec::new();
186            data.read_to_end(&mut buf)?;
187            let code = python_marshal::load_bytes(&buf, python_version)?;
188            Ok(CodeObject::V313(code.try_into()?))
189        }
190        _ => Err(Error::UnsupportedVersion(python_version)),
191    }
192}
193
194pub fn dump_code(
195    code_object: CodeObject,
196    python_version: PyVersion,
197    marshal_version: u8,
198) -> Result<Vec<u8>, Error> {
199    let object = match code_object {
200        #[cfg(feature = "v310")]
201        CodeObject::V310(code) => python_marshal::Object::Code(code.into()),
202        #[cfg(feature = "v311")]
203        CodeObject::V311(code) => python_marshal::Object::Code(code.into()),
204        #[cfg(feature = "v312")]
205        CodeObject::V312(code) => python_marshal::Object::Code(code.into()),
206        #[cfg(feature = "v313")]
207        CodeObject::V313(code) => python_marshal::Object::Code(code.into()),
208    };
209
210    let (obj, refs) = minimize_references(&object, vec![]);
211
212    Ok(python_marshal::dump_bytes(
213        obj,
214        Some(refs),
215        python_version,
216        marshal_version,
217    )?)
218}