From 965d02ad41a5269e0b858a7e983c7d7e24e1ee33 Mon Sep 17 00:00:00 2001
From: mateuszz0000 <mtszzwdzk@gmail.com>
Date: Wed, 20 May 2020 08:23:17 +0200
Subject: [PATCH] Update atbash cipher (doc, doctest, performance) (#2017)

* Update atbash

* Add benchmark() to quantify the performance improvement

Co-authored-by: Christian Clauss <cclauss@me.com>
---
 ciphers/atbash.py | 59 +++++++++++++++++++++++++++++++++++++++++++----
 1 file changed, 55 insertions(+), 4 deletions(-)

diff --git a/ciphers/atbash.py b/ciphers/atbash.py
index 4cf003859..c17d1e34f 100644
--- a/ciphers/atbash.py
+++ b/ciphers/atbash.py
@@ -1,6 +1,17 @@
-def atbash():
+""" https://en.wikipedia.org/wiki/Atbash """
+import string
+
+
+def atbash_slow(sequence: str) -> str:
+    """
+    >>> atbash_slow("ABCDEFG")
+    'ZYXWVUT'
+
+    >>> atbash_slow("aW;;123BX")
+    'zD;;123YC'
+    """
     output = ""
-    for i in input("Enter the sentence to be encrypted ").strip():
+    for i in sequence:
         extract = ord(i)
         if 65 <= extract <= 90:
             output += chr(155 - extract)
@@ -8,8 +19,48 @@ def atbash():
             output += chr(219 - extract)
         else:
             output += i
-    print(output)
+    return output
+
+
+def atbash(sequence: str) -> str:
+    """
+    >>> atbash("ABCDEFG")
+    'ZYXWVUT'
+
+    >>> atbash("aW;;123BX")
+    'zD;;123YC'
+    """
+    letters = string.ascii_letters
+    letters_reversed = string.ascii_lowercase[::-1] + string.ascii_uppercase[::-1]
+    return "".join(
+        letters_reversed[letters.index(c)] if c in letters else c for c in sequence
+    )
+
+
+def benchmark() -> None:
+    """Let's benchmark them side-by-side..."""
+    from timeit import timeit
+
+    print("Running performance benchmarks...")
+    print(
+        "> atbash_slow()",
+        timeit(
+            "atbash_slow(printable)",
+            setup="from string import printable ; from __main__ import atbash_slow",
+        ),
+        "seconds",
+    )
+    print(
+        ">      atbash()",
+        timeit(
+            "atbash(printable)",
+            setup="from string import printable ; from __main__ import atbash",
+        ),
+        "seconds",
+    )
 
 
 if __name__ == "__main__":
-    atbash()
+    for sequence in ("ABCDEFGH", "123GGjj", "testStringtest", "with space"):
+        print(f"{sequence} encrypted in atbash: {atbash(sequence)}")
+    benchmark()