changeset 26:abd21cfb67f4

EnigmaMachine.from_key_sheet can now take a string for rotor settings.
author Brian Neal <bgneal@gmail.com>
date Tue, 29 May 2012 12:43:17 -0500
parents dbc0c73735d6
children dc7f939a2ebf
files enigma/machine.py enigma/tests/test_enigma.py
diffstat 2 files changed, 19 insertions(+), 8 deletions(-) [+]
line wrap: on
line diff
--- a/enigma/machine.py	Mon May 28 20:43:31 2012 -0500
+++ b/enigma/machine.py	Tue May 29 12:43:17 2012 -0500
@@ -61,9 +61,13 @@
         or a single string:
             e.g. ["I", "III", "IV"] or "I III IV"
 
-        ring_settings: an iterable of integers or None representing the ring
-        settings to be applied to the rotors in the rotors list. None means all
-        ring settings are 0.
+        ring_settings: either a list/tuple of integers, a string, or None to
+        represent the ring settings to be applied to the rotors in the rotors
+        list. The acceptable values are:
+            - A list/tuple of integers with values between 0-25
+            - A string; either space separated letters or numbers, e.g. 'B U L'
+              or '1 20 11'
+            - None means all ring settings are 0.
 
         reflector: a string that names the reflector to use
 
@@ -86,6 +90,16 @@
 
         if ring_settings is None:
             ring_settings = [0] * num_rotors
+        elif isinstance(ring_settings, str):
+            strings = ring_settings.split()
+            ring_settings = []
+            for s in strings:
+                if s.isalpha():
+                    ring_settings.append(ord(s.upper()) - ord('A'))
+                elif s.isdigit():
+                    ring_settings.append(int(s))
+                else:
+                    raise EnigmaError('invalid ring setting: %s' % s)
 
         if num_rotors != len(ring_settings):
             raise EnigmaError("# of rotors doesn't match # of ring settings")
--- a/enigma/tests/test_enigma.py	Mon May 28 20:43:31 2012 -0500
+++ b/enigma/tests/test_enigma.py	Tue May 29 12:43:17 2012 -0500
@@ -59,12 +59,10 @@
     """
     def setUp(self):
 
-        ring_settings = [ord(c) - ord('A') for c in 'BUL']
-
         self.machine = EnigmaMachine.from_key_sheet(
                 rotors='II IV V',
                 reflector='B',
-                ring_settings=ring_settings,
+                ring_settings='B U L',
                 plugboard_settings='AV BS CG DL FU HZ IN KM OW RX')
 
     def decrypt(self, start, enc_key, ciphertext, truth_data):
@@ -147,12 +145,11 @@
     """
     def test_decrypt(self):
 
-        ring_settings = [ord(c) - ord('A') for c in 'AAAV']
         stecker ='1/20 2/12 4/6 7/10 8/13 14/23 15/16 17/25 18/26 22/24'
 
         machine = EnigmaMachine.from_key_sheet(
                     rotors='Beta II IV I',
-                    ring_settings=ring_settings,
+                    ring_settings='A A A V',
                     reflector='B-Thin',
                     plugboard_settings=stecker)