Add str.encode for utf-8

youknowone · youknowone · commit 7f2560c9e1b8 · 2019-05-29T21:25:34.000+09:00
diff --git a/tests/snippets/strings.py b/tests/snippets/strings.py
@@ -206,3 +206,13 @@ def try_mutate_str():
    word[0] = 'x'
 
 assert_raises(TypeError, try_mutate_str)
+
+ss = ['Hello', '안녕', '👋']
+bs = [b'Hello', b'\xec\x95\x88\xeb\x85\x95', b'\xf0\x9f\x91\x8b']
+
+for s, b in zip(ss, bs):
+    assert s.encode() == b
+
+for s, b, e in zip(ss, bs, ['u8', 'U8', 'utf-8', 'UTF-8', 'utf_8']):
+    assert s.encode(e) == b
+    # assert s.encode(encoding=e) == b
diff --git a/vm/src/function.rs b/vm/src/function.rs
@@ -374,6 +374,17 @@ impl<T> OptionalArg<T> {
             Missing => f(),
         }
     }
+
+    pub fn map_or_else<U, D, F>(self, default: D, f: F) -> U
+    where
+        D: FnOnce() -> U,
+        F: FnOnce(T) -> U,
+    {
+        match self {
+            Present(value) => f(value),
+            Missing => default(),
+        }
+    }
 }
 
 impl<T> FromArgs for OptionalArg<T>
diff --git a/vm/src/obj/objbyteinner.rs b/vm/src/obj/objbyteinner.rs
@@ -93,24 +93,28 @@ pub fn normalize_encoding(encoding: &str) -> String {
     res
 }
 
+pub fn encode_to_vec(value: &str, encoding: &str, vm: &VirtualMachine) -> PyResult<Vec<u8>> {
+    let encoding = normalize_encoding(encoding);
+    if encoding == "utf_8" || encoding == "u8" {
+        Ok(value.as_bytes().to_vec())
+    } else {
+        // TODO: different encoding
+        return Err(
+            vm.new_value_error(format!("unknown encoding: {}", encoding)), //should be lookup error
+        );
+    }
+}
+
 impl ByteInnerNewOptions {
     pub fn get_value(self, vm: &VirtualMachine) -> PyResult<PyByteInner> {
         // First handle bytes(string, encoding[, errors])
         if let OptionalArg::Present(enc) = self.encoding {
             if let OptionalArg::Present(eval) = self.val_option {
                 if let Ok(input) = eval.downcast::<PyString>() {
                     let encoding = enc.as_str();
-                    if encoding.to_lowercase() == "utf8" || encoding.to_lowercase() == "utf-8"
-                    // TODO: different encoding
-                    {
-                        return Ok(PyByteInner {
-                            elements: input.value.as_bytes().to_vec(),
-                        });
-                    } else {
-                        return Err(
-                            vm.new_value_error(format!("unknown encoding: {}", encoding)), //should be lookup error
-                        );
-                    }
+                    return Ok(PyByteInner {
+                        elements: encode_to_vec(&input.value, &encoding, vm)?,
+                    });
                 } else {
                     return Err(vm.new_type_error("encoding without a string argument".to_string()));
                 }
diff --git a/vm/src/obj/objstr.rs b/vm/src/obj/objstr.rs
@@ -19,6 +19,7 @@ use crate::pyobject::{
 };
 use crate::vm::VirtualMachine;
 
+use super::objbyteinner;
 use super::objdict::PyDict;
 use super::objint::{self, PyInt};
 use super::objnone::PyNone;
@@ -957,6 +958,31 @@ impl PyString {
             }
         }
     }
+
+    #[pymethod]
+    fn encode(
+        &self,
+        encoding: OptionalArg<PyObjectRef>,
+        _errors: OptionalArg<PyObjectRef>,
+        vm: &VirtualMachine,
+    ) -> PyResult {
+        let encoding = encoding.map_or_else(
+            || Ok("utf-8".to_string()),
+            |v| {
+                if objtype::isinstance(&v, &vm.ctx.str_type()) {
+                    Ok(get_value(&v))
+                } else {
+                    Err(vm.new_type_error(format!(
+                        "encode() argument 1 must be str, not {}",
+                        v.class().name
+                    )))
+                }
+            },
+        )?;
+
+        let encoded = objbyteinner::encode_to_vec(&self.value, &encoding, vm)?;
+        Ok(vm.ctx.new_bytes(encoded))
+    }
 }
 
 impl PyValue for PyString {

Original file line number	Diff line number	Diff line change
`@@ -374,6 +374,17 @@ impl<T> OptionalArg<T> {`
`374`	`374`	`Missing => f(),`
`375`	`375`	`}`
`376`	`376`	`}`
	`377`	`+`
	`378`	`+ pub fn map_or_else<U, D, F>(self, default: D, f: F) -> U`
	`379`	`+ where`
	`380`	`+ D: FnOnce() -> U,`
	`381`	`+ F: FnOnce(T) -> U,`
	`382`	`+ {`
	`383`	`+ match self {`
	`384`	`+ Present(value) => f(value),`
	`385`	`+ Missing => default(),`
	`386`	`+ }`
	`387`	`+ }`
`377`	`388`	`}`
`378`	`389`
`379`	`390`	`impl<T> FromArgs for OptionalArg<T>`