Merge pull request #1166 from palaviv/write-bytecode

Write bytecode

src/main.rs

@@ -84,6 +84,11 @@ fn parse_arguments<'a>(app: App<'a, '_>) -> ArgMatches<'a> {
                 .short("S")
                 .help("don't imply 'import site' on initialization"),
         )
+        .arg(
+            Arg::with_name("dont-write-bytecode")
+                .short("B")
+                .help("don't write .pyc files on import"),
+        )
         .arg(
             Arg::with_name("ignore-environment")
                 .short("E")
@@ -171,6 +176,12 @@ fn create_settings(matches: &ArgMatches) -> PySettings {
         settings.quiet = true;
     }
 
+    if matches.is_present("dont-write-bytecode")
+        || (!ignore_environment && env::var_os("PYTHONDONTWRITEBYTECODE").is_some())
+    {
+        settings.dont_write_bytecode = true;
+    }
+
     settings
 }
 

vm/src/import.rs

@@ -1,11 +1,13 @@
 /*
  * Import mechanics
  */
+use rand::Rng;
 
 use crate::bytecode::CodeObject;
 use crate::obj::{objcode, objsequence, objstr, objtype};
 use crate::pyobject::{ItemProtocol, PyObjectRef, PyResult, PyValue};
 use crate::scope::Scope;
+use crate::version::get_git_revision;
 use crate::vm::VirtualMachine;
 #[cfg(feature = "rustpython-compiler")]
 use rustpython_compiler::compile;
@@ -23,6 +25,15 @@ pub fn init_importlib(vm: &VirtualMachine, external: bool) -> PyResult {
         let install_external =
             vm.get_attribute(importlib.clone(), "_install_external_importers")?;
         vm.invoke(install_external, vec![])?;
+        // Set pyc magic number to commit hash. Should be changed when bytecode will be more stable.
+        let importlib_external =
+            vm.import("_frozen_importlib_external", &vm.ctx.new_tuple(vec![]), 0)?;
+        let mut magic = get_git_revision().into_bytes();
+        magic.truncate(4);
+        if magic.len() != 4 {
+            magic = rand::thread_rng().gen::<[u8; 4]>().to_vec();
+        }
+        vm.set_attr(&importlib_external, "MAGIC_NUMBER", vm.ctx.new_bytes(magic))?;
     }
     Ok(vm.get_none())
 }

vm/src/sysmodule.rs

@@ -88,6 +88,7 @@ impl SysFlags {
         flags.ignore_environment = settings.ignore_environment;
         flags.verbose = settings.verbose;
         flags.quiet = settings.quiet;
+        flags.dont_write_bytecode = settings.dont_write_bytecode;
         flags
     }
 }
@@ -322,7 +323,7 @@ settrace() -- set the global debug tracing function
       "path_hooks" => ctx.new_list(vec![]),
       "path_importer_cache" => ctx.new_dict(),
       "pycache_prefix" => vm.get_none(),
-      "dont_write_bytecode" => vm.new_bool(true),
+      "dont_write_bytecode" => vm.new_bool(vm.settings.dont_write_bytecode),
       "setprofile" => ctx.new_rustfunc(sys_setprofile),
       "settrace" => ctx.new_rustfunc(sys_settrace),
       "version" => vm.new_str(version::get_version()),

vm/src/vm.rs

@@ -90,6 +90,9 @@ pub struct PySettings {
     /// -q
     pub quiet: bool,
 
+    /// -B
+    pub dont_write_bytecode: bool,
+
     /// Environment PYTHONPATH and RUSTPYTHONPATH:
     pub path_list: Vec<String>,
 }
@@ -122,6 +125,7 @@ impl Default for PySettings {
             ignore_environment: false,
             verbose: 0,
             quiet: false,
+            dont_write_bytecode: false,
             path_list: vec![],
         }
     }