Implement more complete version of str_repr.

Author: cthulahoops

tests/snippets/strings.py

@@ -16,4 +16,9 @@
 assert str() == ""
 assert str("abc") == "abc"
 
-assert str(["a", "b", "can't"]) == "['a', 'b', 'can\\'t']"
+assert repr("a") == "'a'"
+assert repr("can't") == '"can\'t"'
+assert repr('"won\'t"') == "'\"won\\'t\"'"
+assert repr('\n\t') == "'\\n\\t'"
+
+assert str(["a", "b", "can't"]) == "['a', 'b', \"can't\"]"

vm/src/builtins.rs

@@ -288,6 +288,10 @@ fn builtin_range(vm: &mut VirtualMachine, args: PyFuncArgs) -> PyResult {
 }
 
 // builtin_repr
+fn builtin_repr(vm: &mut VirtualMachine, args: PyFuncArgs) -> PyResult {
+    arg_check!(vm, args, required = [(obj, None)]);
+    vm.to_repr(obj.clone())
+}
 // builtin_reversed
 // builtin_round
 // builtin_set
@@ -346,6 +350,7 @@ pub fn make_module(ctx: &PyContext) -> PyObjectRef {
     dict.insert(String::from("locals"), ctx.new_rustfunc(builtin_locals));
     dict.insert(String::from("print"), ctx.new_rustfunc(builtin_print));
     dict.insert(String::from("range"), ctx.new_rustfunc(builtin_range));
+    dict.insert(String::from("repr"), ctx.new_rustfunc(builtin_repr));
     dict.insert(String::from("setattr"), ctx.new_rustfunc(builtin_setattr));
     dict.insert(String::from("str"), ctx.str_type()); // new_rustfunc(builtin_str));
     dict.insert(String::from("tuple"), ctx.tuple_type());

vm/src/obj/objstr.rs

@@ -45,22 +45,38 @@ fn str_str(vm: &mut VirtualMachine, args: PyFuncArgs) -> PyResult {
     Ok(s.clone())
 }
 
+fn count_char(s: &str, c: char) -> usize {
+    s.chars().filter(|x| *x == c).count()
+}
+
 fn str_repr(vm: &mut VirtualMachine, args: PyFuncArgs) -> PyResult {
     arg_check!(vm, args, required = [(s, Some(vm.ctx.str_type()))]);
     let value = get_value(s);
-    let mut formatted = String::from("'");
+    let quote_char = if count_char(&value, '\'') > count_char(&value, '"') {
+        '"'
+    } else {
+        '\''
+    };
+    let mut formatted = String::new();
+    formatted.push(quote_char);
     for c in value.chars() {
-        match c {
-            '\'' | '\\' => {
-                formatted.push('\\');
-                formatted.push(c);
-            }
-            _ => {
-                formatted.push(c);
-            }
+        if c == quote_char || c == '\\' {
+            formatted.push('\\');
+            formatted.push(c);
+        } else if c == '\n' {
+            formatted.push('\\');
+            formatted.push('n');
+        } else if c == '\t' {
+            formatted.push('\\');
+            formatted.push('t');
+        } else if c == '\r' {
+            formatted.push('\\');
+            formatted.push('r');
+        } else {
+            formatted.push(c);
         }
     }
-    formatted.push('\'');
+    formatted.push(quote_char);
     Ok(vm.ctx.new_str(formatted))
 }