From 18ac16828d80820949ea4c02b0312f6d31bf8b51 Mon Sep 17 00:00:00 2001
From: Jeong YunWon <jeong@youknowone.org>
Date: Thu, 25 Apr 2019 21:02:06 +0900
Subject: [PATCH] Add bytes.maketrans

---
 tests/snippets/bytes.py |  8 ++++++++
 vm/src/obj/objbytes.rs  | 20 +++++++++++++++++++-
 2 files changed, 27 insertions(+), 1 deletion(-)

diff --git a/tests/snippets/bytes.py b/tests/snippets/bytes.py
index c496bfeffb..b822a6d100 100644
--- a/tests/snippets/bytes.py
+++ b/tests/snippets/bytes.py
@@ -162,3 +162,11 @@
 with assertRaises(TypeError):
     b"b".center(2, b"ba")
 b"kok".center(5, bytearray(b"x"))
+
+assert bytes.maketrans(b'', b'') == bytes(range(0, 256))
+trans1, trans2 = b'dcba0123', b'8900xyzz'
+manual_map = list(range(0, 256))
+for c1, c2 in zip(trans1, trans2):
+    manual_map[c1] = c2
+
+assert bytes(manual_map) == bytes.maketrans(trans1, trans2)
diff --git a/vm/src/obj/objbytes.rs b/vm/src/obj/objbytes.rs
index 0556de0535..caf4c3a1d3 100644
--- a/vm/src/obj/objbytes.rs
+++ b/vm/src/obj/objbytes.rs
@@ -10,7 +10,7 @@ use crate::pyobject::{PyClassImpl, PyContext, PyObjectRef, PyRef, PyResult, PyVa
 use super::objbyteinner::{is_byte, PyByteInner};
 use super::objiter;
 use super::objslice::PySlice;
-use super::objtype::PyClassRef;
+use super::objtype::{self, PyClassRef};
 
 /// "bytes(iterable_of_ints) -> bytes\n\
 /// bytes(string, encoding[, errors]) -> bytes\n\
@@ -63,6 +63,7 @@ pub fn init(context: &PyContext) {
     let bytes_type = &context.bytes_type;
     extend_class!(context, bytes_type, {
     "fromhex" => context.new_rustfunc(PyBytesRef::fromhex),
+    "maketrans" => context.new_rustfunc(PyBytesRef::maketrans),
     });
     let bytesiterator_type = &context.bytesiterator_type;
     extend_class!(context, bytesiterator_type, {
@@ -234,6 +235,23 @@ impl PyBytesRef {
         )
     }
 
+    fn maketrans(frm: PyObjectRef, to: PyObjectRef, vm: &VirtualMachine) -> PyResult {
+        if !objtype::isinstance(&frm, &vm.ctx.bytes_type())
+            || !objtype::isinstance(&to, &vm.ctx.bytes_type())
+        {
+            return Err(vm.new_type_error(format!("TypeError: a bytes-like object is required")));
+        }
+        let frm = get_value(&frm);
+        let to = get_value(&to);
+
+        let mut table: Vec<u8> = (0..=255).collect();
+        for (i, c) in frm.iter().enumerate() {
+            table[*c as usize] = to[i];
+        }
+
+        Ok(vm.ctx.new_bytes(table))
+    }
+
     #[pymethod(name = "center")]
     fn center(
         self,