Support freezing of stdlib

palaviv · palaviv · commit 5c240d71a885 · 2019-07-20T18:06:46.000+03:00
diff --git a/Cargo.lock b/Cargo.lock
diff --git a/derive/Cargo.toml b/derive/Cargo.toml
@@ -21,3 +21,4 @@ rustpython-compiler = { path = "../compiler", version = "0.1.0" }
 rustpython-bytecode = { path = "../bytecode", version = "0.1.0" }
 bincode =  "1.1"
 proc-macro-hack = { version = "0.5", optional = true }
+maplit = "1.0"
diff --git a/derive/src/compile_bytecode.rs b/derive/src/compile_bytecode.rs
@@ -19,6 +19,7 @@ use proc_macro2::{Span, TokenStream as TokenStream2};
 use quote::quote;
 use rustpython_bytecode::bytecode::CodeObject;
 use rustpython_compiler::compile;
+use std::collections::HashMap;
 use std::env;
 use std::fs;
 use std::path::PathBuf;
@@ -28,6 +29,7 @@ use syn::{self, parse2, Lit, LitByteStr, LitStr, Meta, Token};
 enum CompilationSourceKind {
     File(PathBuf),
     SourceCode(String),
+    Dir(PathBuf),
 }
 
 struct CompilationSource {
@@ -36,14 +38,22 @@ struct CompilationSource {
 }
 
 impl CompilationSource {
-    fn compile(self, mode: &compile::Mode, module_name: String) -> Result<CodeObject, Diagnostic> {
-        let compile = |source| {
-            compile::compile(source, mode, module_name, 0).map_err(|err| {
-                Diagnostic::spans_error(self.span, format!("Compile error: {}", err))
-            })
-        };
-
-        match &self.kind {
+    fn compile_string(
+        &self,
+        source: &str,
+        mode: &compile::Mode,
+        module_name: String,
+    ) -> Result<CodeObject, Diagnostic> {
+        compile::compile(source, mode, module_name, 0)
+            .map_err(|err| Diagnostic::spans_error(self.span, format!("Compile error: {}", err)))
+    }
+
+    fn compile(
+        &self,
+        mode: &compile::Mode,
+        module_name: String,
+    ) -> Result<HashMap<String, CodeObject>, Diagnostic> {
+        Ok(match &self.kind {
             CompilationSourceKind::File(rel_path) => {
                 let mut path = PathBuf::from(
                     env::var_os("CARGO_MANIFEST_DIR").expect("CARGO_MANIFEST_DIR is not present"),
@@ -55,10 +65,61 @@ impl CompilationSource {
                         format!("Error reading file {:?}: {}", path, err),
                     )
                 })?;
-                compile(&source)
+                hashmap! {module_name.clone() => self.compile_string(&source, mode, module_name.clone())?}
+            }
+            CompilationSourceKind::SourceCode(code) => {
+                hashmap! {module_name.clone() => self.compile_string(code, mode, module_name.clone())?}
+            }
+            CompilationSourceKind::Dir(rel_path) => {
+                let mut path = PathBuf::from(
+                    env::var_os("CARGO_MANIFEST_DIR").expect("CARGO_MANIFEST_DIR is not present"),
+                );
+                path.push(rel_path);
+                self.compile_dir(&path, "".to_string(), mode)?
+            }
+        })
+    }
+
+    fn compile_dir(
+        &self,
+        path: &PathBuf,
+        parent: String,
+        mode: &compile::Mode,
+    ) -> Result<HashMap<String, CodeObject>, Diagnostic> {
+        let mut code_map = HashMap::new();
+        let paths = fs::read_dir(&path).map_err(|err| {
+            Diagnostic::spans_error(self.span, format!("Error listing dir {:?}: {}", path, err))
+        })?;
+        for path in paths {
+            let path = path.map_err(|err| {
+                Diagnostic::spans_error(self.span, format!("Failed to list file: {}", err))
+            })?;
+            let path = path.path();
+            let file_name = path.file_name().unwrap().to_str().unwrap();
+            if path.is_dir() {
+                code_map.extend(self.compile_dir(
+                    &path,
+                    format!("{}{}.", parent, file_name),
+                    mode,
+                )?);
+            } else {
+                if file_name.ends_with(".py") {
+                    let source = fs::read_to_string(&path).map_err(|err| {
+                        Diagnostic::spans_error(
+                            self.span,
+                            format!("Error reading file {:?}: {}", path, err),
+                        )
+                    })?;
+                    let file_name_splitte: Vec<&str> = file_name.splitn(2, ".").collect();
+                    let module_name = format!("{}{}", parent, file_name_splitte[0]);
+                    code_map.insert(
+                        module_name.clone(),
+                        self.compile_string(&source, mode, module_name)?,
+                    );
+                }
             }
-            CompilationSourceKind::SourceCode(code) => compile(code),
         }
+        Ok(code_map)
     }
 }
 
@@ -69,7 +130,7 @@ struct PyCompileInput {
 }
 
 impl PyCompileInput {
-    fn compile(&self) -> Result<CodeObject, Diagnostic> {
+    fn compile(&self) -> Result<HashMap<String, CodeObject>, Diagnostic> {
         let mut module_name = None;
         let mut mode = None;
         let mut source: Option<CompilationSource> = None;
@@ -122,6 +183,16 @@ impl PyCompileInput {
                         kind: CompilationSourceKind::File(path),
                         span: extract_spans(&name_value).unwrap(),
                     });
+                } else if name_value.ident == "dir" {
+                    assert_source_empty(&source)?;
+                    let path = match &name_value.lit {
+                        Lit::Str(s) => PathBuf::from(s.value()),
+                        _ => bail_span!(name_value.lit, "source must be a string"),
+                    };
+                    source = Some(CompilationSource {
+                        kind: CompilationSourceKind::Dir(path),
+                        span: extract_spans(&name_value).unwrap(),
+                    });
                 }
             }
         }
@@ -154,18 +225,22 @@ impl Parse for PyCompileInput {
 pub fn impl_py_compile_bytecode(input: TokenStream2) -> Result<TokenStream2, Diagnostic> {
     let input: PyCompileInput = parse2(input)?;
 
-    let code_obj = input.compile()?;
-
-    let module_name = LitStr::new(&code_obj.source_path, Span::call_site());
+    let code_map = input.compile()?;
 
-    let bytes = bincode::serialize(&code_obj).expect("Failed to serialize");
-    let bytes = LitByteStr::new(&bytes, Span::call_site());
+    let modules = code_map.iter().map(|(module_name, code_obj)| {
+        let module_name = LitStr::new(&module_name, Span::call_site());
+        let bytes = bincode::serialize(&code_obj).expect("Failed to serialize");
+        let bytes = LitByteStr::new(&bytes, Span::call_site());
+        quote! {  #module_name.into() => bincode::deserialize::<::rustpython_vm::bytecode::CodeObject>(#bytes)
+                .expect("Deserializing CodeObject failed") }
+    });
 
     let output = quote! {
         ({
             use ::rustpython_vm::__exports::bincode;
-            hashmap! { #module_name.into() => bincode::deserialize::<::rustpython_vm::bytecode::CodeObject>(#bytes)
-                .expect("Deserializing CodeObject failed")}
+            hashmap! {
+                #(#modules),*
+            }
         })
     };
 
diff --git a/derive/src/lib.rs b/derive/src/lib.rs
@@ -4,6 +4,9 @@
 
 extern crate proc_macro;
 
+#[macro_use]
+extern crate maplit;
+
 #[macro_use]
 mod error;
 mod compile_bytecode;
diff --git a/vm/src/frozen.rs b/vm/src/frozen.rs
@@ -16,5 +16,10 @@ pub fn get_module_inits() -> HashMap<String, CodeObject> {
         module_name = "_frozen_importlib_external",
     ));
 
+    #[cfg(feature = "freeze-stdlib")]
+    {
+        modules.extend(py_compile_bytecode!(dir = "../Lib/",));
+    }
+
     modules
 }
