RustPython · youknowone · Apr 23, 2024 · Apr 21, 2024 · Apr 21, 2024 · Apr 21, 2024
diff --git a/Lib/test/test_builtin.py b/Lib/test/test_builtin.py
@@ -327,7 +327,7 @@ def test_chr(self):
     def test_cmp(self):
         self.assertTrue(not hasattr(builtins, "cmp"))
 
-    # TODO: RUSTPYTHON
+    # TODO: RUSTPYTHON optval=2 does not remove docstrings
     @unittest.expectedFailure
     def test_compile(self):
         compile('print(1)\n', '', 'exec')
@@ -340,11 +340,10 @@ def test_compile(self):
         self.assertRaises(TypeError, compile)
         self.assertRaises(ValueError, compile, 'print(42)\n', '<string>', 'badmode')
         self.assertRaises(ValueError, compile, 'print(42)\n', '<string>', 'single', 0xff)
-        self.assertRaises(ValueError, compile, chr(0), 'f', 'exec')
         self.assertRaises(TypeError, compile, 'pass', '?', 'exec',
                           mode='eval', source='0', filename='tmp')
         compile('print("\xe5")\n', '', 'exec')
-        self.assertRaises(ValueError, compile, chr(0), 'f', 'exec')
+        self.assertRaises(SyntaxError, compile, chr(0), 'f', 'exec')
         self.assertRaises(ValueError, compile, str('a = 1'), 'f', 'bad')
 
         # test the optimize argument

diff --git a/vm/src/stdlib/ast.rs b/vm/src/stdlib/ast.rs
@@ -341,8 +341,13 @@ pub(crate) fn compile(
     object: PyObjectRef,
     filename: &str,
     mode: crate::compiler::Mode,
+    optimize: Option<u8>,
 ) -> PyResult {
-    let opts = vm.compile_opts();
+    let mut opts = vm.compile_opts();
+    if let Some(optimize) = optimize {
+        opts.optimize = optimize;
+    }
+
     let ast = Node::ast_from_object(vm, object)?;
     let code = codegen::compile::compile_top(&ast, filename.to_owned(), mode, opts)
         .map_err(|err| (CompileError::from(err), None).to_pyexception(vm))?; // FIXME source
@@ -354,6 +359,43 @@ pub(crate) use _ast::NodeAst;
 // Used by builtins::compile()
 pub const PY_COMPILE_FLAG_AST_ONLY: i32 = 0x0400;
 
+// The following flags match the values from Include/cpython/compile.h
+// Caveat emptor: These flags are undocumented on purpose and depending
+// on their effect outside the standard library is **unsupported**.
+const PY_CF_DONT_IMPLY_DEDENT: i32 = 0x200;
+const PY_CF_ALLOW_INCOMPLETE_INPUT: i32 = 0x4000;
+
+// __future__ flags - sync with Lib/__future__.py
+// TODO: These flags aren't being used in rust code
+//       CO_FUTURE_ANNOTATIONS does make a difference in the codegen,
+//       so it should be used in compile().
+//       see compiler/codegen/src/compile.rs
+const CO_NESTED: i32 = 0x0010;
+const CO_GENERATOR_ALLOWED: i32 = 0;
+const CO_FUTURE_DIVISION: i32 = 0x20000;
+const CO_FUTURE_ABSOLUTE_IMPORT: i32 = 0x40000;
+const CO_FUTURE_WITH_STATEMENT: i32 = 0x80000;
+const CO_FUTURE_PRINT_FUNCTION: i32 = 0x100000;
+const CO_FUTURE_UNICODE_LITERALS: i32 = 0x200000;
+const CO_FUTURE_BARRY_AS_BDFL: i32 = 0x400000;
+const CO_FUTURE_GENERATOR_STOP: i32 = 0x800000;
+const CO_FUTURE_ANNOTATIONS: i32 = 0x1000000;
+
+// Used by builtins::compile() - the summary of all flags
+pub const PY_COMPILE_FLAGS_MASK: i32 = PY_COMPILE_FLAG_AST_ONLY
+    | PY_CF_DONT_IMPLY_DEDENT
+    | PY_CF_ALLOW_INCOMPLETE_INPUT
+    | CO_NESTED
+    | CO_GENERATOR_ALLOWED
+    | CO_FUTURE_DIVISION
+    | CO_FUTURE_ABSOLUTE_IMPORT
+    | CO_FUTURE_WITH_STATEMENT
+    | CO_FUTURE_PRINT_FUNCTION
+    | CO_FUTURE_UNICODE_LITERALS
+    | CO_FUTURE_BARRY_AS_BDFL
+    | CO_FUTURE_GENERATOR_STOP
+    | CO_FUTURE_ANNOTATIONS;
+
 pub fn make_module(vm: &VirtualMachine) -> PyRef<PyModule> {
     let module = _ast::make_module(vm);
     gen::extend_module_nodes(vm, &module);

diff --git a/vm/src/stdlib/builtins.rs b/vm/src/stdlib/builtins.rs
@@ -121,6 +121,15 @@ mod builtins {
 
             let mode_str = args.mode.as_str();
 
+            let optimize: i32 = args.optimize.map_or(Ok(-1), |v| v.try_to_primitive(vm))?;
+            let optimize: u8 = if optimize == -1 {
+                vm.state.settings.optimize
+            } else {
+                optimize.try_into().map_err(|_| {
+                    vm.new_value_error("compile() optimize value invalid".to_owned())
+                })?
+            };
+
             if args
                 .source
                 .fast_isinstance(&ast::NodeAst::make_class(&vm.ctx))
@@ -134,7 +143,13 @@ mod builtins {
                     let mode = mode_str
                         .parse::<crate::compiler::Mode>()
                         .map_err(|err| vm.new_value_error(err.to_string()))?;
-                    return ast::compile(vm, args.source, args.filename.as_str(), mode);
+                    return ast::compile(
+                        vm,
+                        args.source,
+                        args.filename.as_str(),
+                        mode,
+  
EB9C
                      Some(optimize),
+                    );
                 }
             }
 
@@ -146,20 +161,23 @@ mod builtins {
             }
             #[cfg(feature = "rustpython-parser")]
             {
-                use crate::{builtins::PyBytesRef, convert::ToPyException};
+                use crate::convert::ToPyException;
                 use num_traits::Zero;
                 use rustpython_parser as parser;
 
-                let source = Either::<PyStrRef, PyBytesRef>::try_from_object(vm, args.source)?;
+                let source = ArgStrOrBytesLike::try_from_object(vm, args.source)?;
+                let source = source.borrow_bytes();
+
                 // TODO: compiler::compile should probably get bytes
-                let source = match &source {
-                    Either::A(string) => string.as_str(),
-                    Either::B(bytes) => std::str::from_utf8(bytes)
-                        .map_err(|e| vm.new_unicode_decode_error(e.to_string()))?,
-                };
+                let source = std::str::from_utf8(&source)
+                    .map_err(|e| vm.new_unicode_decode_error(e.to_string()))?;
 
                 let flags = args.flags.map_or(Ok(0), |v| v.try_to_primitive(vm))?;
 
+                if !(flags & !ast::PY_COMPILE_FLAGS_MASK).is_zero() {
+                    return Err(vm.new_value_error("compile() unrecognized flags".to_owned()));
+                }
+
                 if (flags & ast::PY_COMPILE_FLAG_AST_ONLY).is_zero() {
                     #[cfg(not(feature = "rustpython-compiler"))]
                     {
@@ -170,8 +188,17 @@ mod builtins {
                         let mode = mode_str
                             .parse::<crate::compiler::Mode>()
                             .map_err(|err| vm.new_value_error(err.to_string()))?;
+
+                        let mut opts = vm.compile_opts();
+                        opts.optimize = optimize;
+
                         let code = vm
-                            .compile(source, mode, args.filename.as_str().to_owned())
+                            .compile_with_opts(
+                                source,
+                                mode,
+                                args.filename.as_str().to_owned(),
+                                opts,
+                            )
                             .map_err(|err| (err, Some(source)).to_pyexception(vm))?;
                         Ok(code.into())
                     }