Add type-hints to more files

imapclient/datetime_util.py

@@ -11,7 +11,7 @@
 _SHORT_MONTHS = " Jan Feb Mar Apr May Jun Jul Aug Sep Oct Nov Dec".split(" ")
 
 
-def parse_to_datetime(timestamp, normalise=True):
+def parse_to_datetime(timestamp: bytes, normalise: bool = True) -> datetime:
     """Convert an IMAP datetime string to a datetime.
 
     If normalise is True (the default), then the returned datetime
@@ -36,11 +36,11 @@ def parse_to_datetime(timestamp, normalise=True):
     return dt
 
 
-def datetime_to_native(dt):
+def datetime_to_native(dt: datetime) -> datetime:
     return dt.astimezone(FixedOffset.for_system()).replace(tzinfo=None)
 
 
-def datetime_to_INTERNALDATE(dt):
+def datetime_to_INTERNALDATE(dt: datetime) -> str:
     """Convert a datetime instance to a IMAP INTERNALDATE string.
 
     If timezone information is missing the current system
@@ -57,14 +57,14 @@ def datetime_to_INTERNALDATE(dt):
 _rfc822_dotted_time = re.compile(r"\w+, ?\d{1,2} \w+ \d\d(\d\d)? \d\d?\.\d\d?\.\d\d?.*")
 
 
-def _munge(s):
-    s = s.decode("latin-1")  # parsedate_tz only works with strings
+def _munge(timestamp: bytes) -> str:
+    s = timestamp.decode("latin-1")  # parsedate_tz only works with strings
     if _rfc822_dotted_time.match(s):
         return s.replace(".", ":")
     return s
 
 
-def format_criteria_date(dt):
+def format_criteria_date(dt: datetime) -> bytes:
     """Format a date or datetime instance for use in IMAP search criteria."""
     out = "%02d-%s-%d" % (dt.day, _SHORT_MONTHS[dt.month], dt.year)
     return out.encode("ascii")

imapclient/imap_utf7.py

@@ -8,9 +8,10 @@
 # Other variations and examples can be found in the RFC 3501, section 5.1.3.
 
 import binascii
+from typing import List, Union
 
 
-def encode(s):
+def encode(s: Union[str, bytes]) -> bytes:
     """Encode a folder name using IMAP modified UTF-7 encoding.
 
     Input is unicode; output is bytes (Python 3) or str (Python 2). If
@@ -21,9 +22,9 @@ def encode(s):
 
     res = bytearray()
 
-    b64_buffer = []
+    b64_buffer: List[str] = []
 
-    def consume_b64_buffer(buf):
+    def consume_b64_buffer(buf: List[str]) -> None:
         """
         Consume the buffer by encoding it into a modified base 64 representation
         and surround it with shift characters & and -
@@ -58,7 +59,7 @@ def consume_b64_buffer(buf):
 DASH_ORD = ord("-")
 
 
-def decode(s):
+def decode(s: Union[bytes, str]) -> str:
     """Decode a folder name from IMAP modified UTF-7 encoding to unicode.
 
     Input is bytes (Python 3) or str (Python 2); output is always
@@ -97,11 +98,11 @@ def decode(s):
     return "".join(res)
 
 
-def base64_utf7_encode(buffer):
+def base64_utf7_encode(buffer: List[str]) -> bytes:
     s = "".join(buffer).encode("utf-16be")
     return binascii.b2a_base64(s).rstrip(b"\n=").replace(b"/", b",")
 
 
-def base64_utf7_decode(s):
+def base64_utf7_decode(s: bytearray) -> str:
     s_utf7 = b"+" + s.replace(b",", b"/") + b"-"
     return s_utf7.decode("utf-7")

imapclient/tls.py

@@ -8,11 +8,18 @@
 """
 
 import imaplib
+import io
 import socket
 import ssl
+from typing import Optional, TYPE_CHECKING
 
+if TYPE_CHECKING:
+    from typing_extensions import Buffer
 
-def wrap_socket(sock, ssl_context, host):
+
+def wrap_socket(
+    sock: socket.socket, ssl_context: Optional[ssl.SSLContext], host: str
+) -> socket.socket:
     if not hasattr(ssl, "create_default_context"):
         # Python 2.7.0 - 2.7.8 do not have the concept of ssl contexts.
         # Thus we have to use the less flexible and legacy way of wrapping the
@@ -37,12 +44,21 @@ class IMAP4_TLS(imaplib.IMAP4):
     Adapted from imaplib.IMAP4_SSL.
     """
 
-    def __init__(self, host, port, ssl_context, timeout=None):
+    def __init__(
+        self,
+        host: str,
+        port: int,
+        ssl_context: Optional[ssl.SSLContext],
+        timeout: Optional[float] = None,
+    ):
         self.ssl_context = ssl_context
         self._timeout = timeout
         imaplib.IMAP4.__init__(self, host, port)
+        self.file: io.BufferedReader
 
-    def open(self, host, port=993, timeout=None):
+    def open(
+        self, host: str = "", port: int = 993, timeout: Optional[float] = None
+    ) -> None:
         self.host = host
         self.port = port
         sock = socket.create_connection(
@@ -51,14 +67,14 @@ def open(self, host, port=993, timeout=None):
         self.sock = wrap_socket(sock, self.ssl_context, host)
         self.file = self.sock.makefile("rb")
 
-    def read(self, size):
+    def read(self, size: int) -> bytes:
         return self.file.read(size)
 
-    def readline(self):
+    def readline(self) -> bytes:
         return self.file.readline()
 
-    def send(self, data):
+    def send(self, data: "Buffer") -> None:
         self.sock.sendall(data)
 
-    def shutdown(self):
+    def shutdown(self) -> None:
         imaplib.IMAP4.shutdown(self)

imapclient/util.py

@@ -3,13 +3,14 @@
 # Please see http://en.wikipedia.org/wiki/BSD_licenses
 
 import logging
+from typing import Iterator, Optional, Tuple, Union
 
 from . import exceptions
 
 logger = logging.getLogger(__name__)
 
 
-def to_unicode(s):
+def to_unicode(s: Union[bytes, str]) -> str:
     if isinstance(s, bytes):
         try:
             return s.decode("ascii")
@@ -23,20 +24,24 @@ def to_unicode(s):
     return s
 
 
-def to_bytes(s, charset="ascii"):
+def to_bytes(s: Union[bytes, str], charset: str = "ascii") -> bytes:
     if isinstance(s, str):
         return s.encode(charset)
     return s
 
 
-def assert_imap_protocol(condition, message=None):
+def assert_imap_protocol(condition: bool, message: Optional[str] = None) -> None:
     if not condition:
         msg = "Server replied with a response that violates the IMAP protocol"
         if message:
             msg += "{}: {}".format(msg, message)
         raise exceptions.ProtocolError(msg)
 
 
-def chunk(lst, size):
+_AtomPart = Tuple[Union[None, int, bytes], ...]
+_Atom = Tuple[Union[_AtomPart, "_Atom"], ...]
+
+
+def chunk(lst: _Atom, size: int) -> Iterator[_Atom]:
     for i in range(0, len(lst), size):
         yield lst[i : i + size]

pyproject.toml

@@ -48,15 +48,11 @@ warn_unused_ignores = true
 [[tool.mypy.overrides]]
 module = [
     "imapclient.config",
-    "imapclient.datetime_util",
-    "imapclient.imap_utf7",
     "imapclient.imapclient",
     "imapclient.interact",
     "imapclient.response_lexer",
     "imapclient.response_parser",
     "imapclient.response_types",
-    "imapclient.tls",
-    "imapclient.util",
     "interact",
     "livetest",
     "setup",