Add lazy_static argument

coolreader18 · coolreader18 · commit b5ba62fe943a · 2019-06-14T21:17:12.000-05:00
diff --git a/derive/src/compile_bytecode.rs b/derive/src/compile_bytecode.rs
@@ -51,11 +51,17 @@ struct PyCompileInput {
     metas: Vec<Meta>,
 }
 
+struct PyCompileResult {
+    code_obj: CodeObject,
+    lazy_static: bool,
+}
+
 impl PyCompileInput {
-    fn compile(&self) -> Result<CodeObject, Diagnostic> {
+    fn compile(&self) -> Result<PyCompileResult, Diagnostic> {
         let mut source_path = None;
         let mut mode = None;
         let mut source: Option<CompilationSource> = None;
+        let mut lazy_static = false;
 
         fn assert_source_empty(source: &Option<CompilationSource>) -> Result<(), Diagnostic> {
             if let Some(source) = source {
@@ -108,11 +114,16 @@ impl PyCompileInput {
                         });
                     }
                 }
+                Meta::Word(ident) => {
+                    if ident == "lazy_static" {
+                        lazy_static = true;
+                    }
+                }
                 _ => {}
             }
         }
 
-        source
+        let code_obj = source
             .ok_or_else(|| {
                 Diagnostic::span_error(
                     self.span.clone(),
@@ -122,7 +133,12 @@ impl PyCompileInput {
             .compile(
                 &mode.unwrap_or(compile::Mode::Exec),
                 source_path.unwrap_or_else(|| "frozen".to_string()),
-            )
+            )?;
+
+        Ok(PyCompileResult {
+            code_obj,
+            lazy_static,
+        })
     }
 }
 
@@ -140,16 +156,32 @@ 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 PyCompileResult {
+        code_obj,
+        lazy_static,
+    } = input.compile()?;
+
     let bytes = bincode::serialize(&code_obj).expect("Failed to serialize");
     let bytes = LitByteStr::new(&bytes, Span::call_site());
 
     let output = quote! {
         ({
-            use bincode;
-            bincode::deserialize(#bytes).expect("Deserializing CodeObject failed")
+            use ::bincode;
+            bincode::deserialize::<::rustpython_vm::bytecode::CodeObject>(#bytes).expect("Deserializing CodeObject failed")
         })
     };
 
-    Ok(output)
+    if lazy_static {
+        Ok(quote! {
+            ({
+                use ::lazy_static::lazy_static;
+                lazy_static! {
+                    static ref STATIC: ::rustpython_vm::bytecode::CodeObject = #output;
+                }
+                &*STATIC
+            })
+        })
+    } else {
+        Ok(output)
+    }
 }
diff --git a/vm/src/frozen.rs b/vm/src/frozen.rs
@@ -1,22 +1,19 @@
 use crate::bytecode::CodeObject;
 use std::collections::HashMap;
 
-lazy_static! {
-    static ref HELLO: CodeObject = py_compile_bytecode!(
-        source = "initialized = True
-print(\"Hello world!\")
-",
-    );
-    static ref IMPORTLIB_BOOTSTRAP: CodeObject =
-        py_compile_bytecode!(file = "../Lib/importlib/_bootstrap.py");
-    static ref IMPORTLIB_BOOTSTRAP_EXTERNAL: CodeObject =
-        py_compile_bytecode!(file = "../Lib/importlib/_bootstrap_external.py");
-}
-
 pub fn get_module_inits() -> HashMap<&'static str, &'static CodeObject> {
     hashmap! {
-        "__hello__" => &*HELLO,
-        "_frozen_importlib_external" => &*IMPORTLIB_BOOTSTRAP_EXTERNAL,
-        "_frozen_importlib" => &*IMPORTLIB_BOOTSTRAP,
+        "__hello__" => py_compile_bytecode!(
+            lazy_static,
+            source = "initialized = True; print(\"Hello world!\")\n",
+        ),
+        "_frozen_importlib" => py_compile_bytecode!(
+            lazy_static,
+            file = "../Lib/importlib/_bootstrap.py",
+        ),
+        "_frozen_importlib_external" => py_compile_bytecode!(
+            lazy_static,
+            file = "../Lib/importlib/_bootstrap_external.py",
+        ),
     }
 }
