Added type annotations + some fixes to get them correct

One functional change: `CryptoOperation.read_infile()` now reads bytes
from `sys.stdin` instead of text. This is necessary to be consistent with
the rest of the code, which all deals with bytes.

1 files changed, 7 insertions, 3 deletions

diff --git a/rsa/pem.py b/rsa/pem.py
index 0650e64..02c7691 100644
--- a/rsa/pem.py
+++ b/rsa/pem.py
@@ -17,9 +17,13 @@
 """Functions that load and write PEM-encoded files."""
 
 import base64
+import typing
 
+# Should either be ASCII strings or bytes.
+FlexiText = typing.Union[str, bytes]
 
-def _markers(pem_marker):
+
+def _markers(pem_marker: FlexiText) -> typing.Tuple[bytes, bytes]:
     """
     Returns the start and end PEM markers, as bytes.
     """
@@ -31,7 +35,7 @@ def _markers(pem_marker):
             b'-----END ' + pem_marker + b'-----')
 
 
-def load_pem(contents, pem_marker):
+def load_pem(contents: FlexiText, pem_marker: FlexiText) -> bytes:
     """Loads a PEM file.
 
     :param contents: the contents of the file to interpret
@@ -97,7 +101,7 @@ def load_pem(contents, pem_marker):
     return base64.standard_b64decode(pem)
 
 
-def save_pem(contents, pem_marker):
+def save_pem(contents: bytes, pem_marker: FlexiText) -> bytes:
     """Saves a PEM file.
 
     :param contents: the contents to encode in PEM format