Configured flask8 to use max_complexity=10

Also reorganised the only function that had a higher complexity.

1 files changed, 29 insertions, 23 deletions

diff --git a/rsa/pem.py b/rsa/pem.py
index 02c7691..a50a5e8 100644
--- a/rsa/pem.py
+++ b/rsa/pem.py
@@ -35,29 +35,11 @@ def _markers(pem_marker: FlexiText) -> typing.Tuple[bytes, bytes]:
             b'-----END ' + pem_marker + b'-----')
 
 
-def load_pem(contents: FlexiText, pem_marker: FlexiText) -> bytes:
-    """Loads a PEM file.
+def _pem_lines(contents: bytes, pem_start: bytes, pem_end: bytes) -> typing.Iterator[bytes]:
+    """Generator over PEM lines between pem_start and pem_end."""
 
-    :param contents: the contents of the file to interpret
-    :param pem_marker: the marker of the PEM content, such as 'RSA PRIVATE KEY'
-        when your file has '-----BEGIN RSA PRIVATE KEY-----' and
-        '-----END RSA PRIVATE KEY-----' markers.
-
-    :return: the base64-decoded content between the start and end markers.
-
-    @raise ValueError: when the content is invalid, for example when the start
-        marker cannot be found.
-
-    """
-
-    # We want bytes, not text. If it's text, it can be converted to ASCII bytes.
-    if not isinstance(contents, bytes):
-        contents = contents.encode('ascii')
-
-    (pem_start, pem_end) = _markers(pem_marker)
-
-    pem_lines = []
     in_pem_part = False
+    seen_pem_start = False
 
     for line in contents.splitlines():
         line = line.strip()
@@ -72,6 +54,7 @@ def load_pem(contents: FlexiText, pem_marker: FlexiText) -> bytes:
                 raise ValueError('Seen start marker "%s" twice' % pem_start)
 
             in_pem_part = True
+            seen_pem_start = True
             continue
 
         # Skip stuff before first marker
@@ -87,15 +70,38 @@ def load_pem(contents: FlexiText, pem_marker: FlexiText) -> bytes:
         if b':' in line:
             continue
 
-        pem_lines.append(line)
+        yield line
 
     # Do some sanity checks
-    if not pem_lines:
+    if not seen_pem_start:
         raise ValueError('No PEM start marker "%s" found' % pem_start)
 
     if in_pem_part:
         raise ValueError('No PEM end marker "%s" found' % pem_end)
 
+
+def load_pem(contents: FlexiText, pem_marker: FlexiText) -> bytes:
+    """Loads a PEM file.
+
+    :param contents: the contents of the file to interpret
+    :param pem_marker: the marker of the PEM content, such as 'RSA PRIVATE KEY'
+        when your file has '-----BEGIN RSA PRIVATE KEY-----' and
+        '-----END RSA PRIVATE KEY-----' markers.
+
+    :return: the base64-decoded content between the start and end markers.
+
+    @raise ValueError: when the content is invalid, for example when the start
+        marker cannot be found.
+
+    """
+
+    # We want bytes, not text. If it's text, it can be converted to ASCII bytes.
+    if not isinstance(contents, bytes):
+        contents = contents.encode('ascii')
+
+    (pem_start, pem_end) = _markers(pem_marker)
+    pem_lines = [line for line in _pem_lines(contents, pem_start, pem_end)]
+
     # Base64-decode the contents
     pem = b''.join(pem_lines)
     return base64.standard_b64decode(pem)