Add type annotations

colorlog/colorlog.py

@@ -4,6 +4,7 @@
 
 import logging
 import sys
+import typing
 
 from colorlog.escape_codes import escape_codes, parse_colors
 
@@ -15,6 +16,10 @@
     "TTYColoredFormatter",
 )
 
+# Type aliases used in function signatures.
+LogColors = typing.Mapping[str, str]
+SecondaryLogColors = typing.Mapping[str, LogColors]
+
 # The default colors to use for the debug levels
 default_log_colors = {
     "DEBUG": "white",
@@ -40,7 +45,7 @@ class ColoredRecord(object):
     StrFormatStyle, and StringTemplateStyle classes).
     """
 
-    def __init__(self, record):
+    def __init__(self, record: logging.LogRecord) -> None:
         """Add attributes from the escape_codes dict and the record."""
         self.__dict__.update(escape_codes)
         self.__dict__.update(record.__dict__)
@@ -49,7 +54,7 @@ def __init__(self, record):
         # access functions that are not in ``__dict__``
         self.__record = record
 
-    def __getattr__(self, name):
+    def __getattr__(self, name: str) -> typing.Any:
         return getattr(self.__record, name)
 
 
@@ -62,13 +67,13 @@ class ColoredFormatter(logging.Formatter):
 
     def __init__(
         self,
-        fmt=None,
-        datefmt=None,
-        style="%",
-        log_colors=None,
-        reset=True,
-        secondary_log_colors=None,
-    ):
+        fmt: typing.Optional[str] = None,
+        datefmt: typing.Optional[str] = None,
+        style: str = "%",
+        log_colors: typing.Optional[LogColors] = None,
+        reset: bool = True,
+        secondary_log_colors: typing.Optional[SecondaryLogColors] = None,
+    ) -> None:
         """
         Set the format and colors the ColoredFormatter will use.
 
@@ -108,11 +113,11 @@ def __init__(
         self.secondary_log_colors = secondary_log_colors
         self.reset = reset
 
-    def color(self, log_colors, level_name):
+    def color(self, log_colors: LogColors, level_name: str) -> str:
         """Return escape codes from a ``log_colors`` dict."""
         return parse_colors(log_colors.get(level_name, ""))
 
-    def format(self, record):
+    def format(self, record: logging.LogRecord) -> str:
         """Format a message from a record object."""
         record = ColoredRecord(record)
         record.log_color = self.color(self.log_colors, record.levelname)
@@ -139,12 +144,12 @@ class LevelFormatter(ColoredFormatter):
 
     def __init__(
         self,
-        fmt=None,
-        datefmt=None,
-        style="%",
-        log_colors=None,
-        reset=True,
-        secondary_log_colors=None,
+        fmt: typing.Optional[str] = None,
+        datefmt: typing.Optional[str] = None,
+        style: str = "%",
+        log_colors: typing.Optional[LogColors] = None,
+        reset: bool = True,
+        secondary_log_colors: typing.Optional[SecondaryLogColors] = None,
     ):
         """
         Set the per-loglevel format that will be used.
@@ -179,7 +184,7 @@ def __init__(
         self.style = style
         self.fmt = fmt
 
-    def format(self, record):
+    def format(self, record: logging.LogRecord) -> str:
         """Customize the message format based on the log level."""
         if isinstance(self.fmt, dict):
             self._fmt = self.fmt[record.levelname]
@@ -202,14 +207,31 @@ class TTYColoredFormatter(ColoredFormatter):
     This is useful when you want to be able to pipe colorlog output to a file.
     """
 
-    def __init__(self, *args, **kwargs):
+    def __init__(
+        self,
+        stream: typing.IO,
+        fmt: typing.Optional[typing.Mapping[str, str]] = None,
+        datefmt: typing.Optional[str] = None,
+        style: str = "%",
+        log_colors: typing.Optional[LogColors] = None,
+        reset: bool = True,
+        secondary_log_colors: typing.Optional[SecondaryLogColors] = None,
+    ) -> None:
         """Overwrite the `reset` argument to False if stream is not a TTY."""
-        self.stream = kwargs.pop("stream")
+        self.stream = stream
 
         # Both `reset` and `isatty` must be true to insert reset codes.
-        kwargs["reset"] = kwargs.get("reset", True) and self.stream.isatty()
+        reset = reset and self.stream.isatty()
 
-        ColoredFormatter.__init__(self, *args, **kwargs)
+        ColoredFormatter.__init__(
+            self,
+            fmt=fmt,
+            datefmt=datefmt,
+            style=style,
+            log_colors=log_colors,
+            reset=reset,
+            secondary_log_colors=secondary_log_colors,
+        )
 
     def color(self, log_colors, level_name):
         """Only returns colors if STDOUT is a TTY."""

colorlog/escape_codes.py

@@ -19,15 +19,28 @@
 
 
 # Returns escape codes from format codes
-def esc(*x):
+def esc(*x: str) -> str:
     return "\033[" + ";".join(x) + "m"
 
 
 # The initial list of escape codes
-escape_codes = {"reset": esc("0"), "bold": esc("01"), "thin": esc("02")}
+escape_codes = {
+    "reset": esc("0"),
+    "bold": esc("01"),
+    "thin": esc("02"),
+}
 
 # The color names
-COLORS = ["black", "red", "green", "yellow", "blue", "purple", "cyan", "white"]
+COLORS = [
+    "black",
+    "red",
+    "green",
+    "yellow",
+    "blue",
+    "purple",
+    "cyan",
+    "white",
+]
 
 PREFIXES = [
     # Foreground without prefix
@@ -48,6 +61,6 @@ def esc(*x):
         escape_codes[prefix_name + name] = esc(prefix + str(code))
 
 
-def parse_colors(sequence):
+def parse_colors(sequence: str) -> str:
     """Return escape codes from a color sequence."""
     return "".join(escape_codes[n] for n in sequence.split(",") if n)

colorlog/logging.py

@@ -4,21 +4,22 @@
 
 import functools
 import logging
+import typing
 
-from colorlog.colorlog import ColoredFormatter
+from colorlog.colorlog import ColoredFormatter, LogColors, SecondaryLogColors
 
 BASIC_FORMAT = "%(log_color)s%(levelname)s%(reset)s:%(name)s:%(message)s"
 
 
 def basicConfig(
-    style="%",
-    log_colors=None,
-    reset=True,
-    secondary_log_colors=None,
-    format=BASIC_FORMAT,
-    datefmt=None,
+    style: str = "%",
+    log_colors: typing.Optional[LogColors] = None,
+    reset: bool = True,
+    secondary_log_colors: typing.Optional[SecondaryLogColors] = None,
+    format: str = BASIC_FORMAT,
+    datefmt: typing.Optional[str] = None,
     **kwargs
-):
+) -> None:
     """Call ``logging.basicConfig`` and override the formatter it creates."""
     logging.basicConfig(**kwargs)
     logging._acquireLock()

setup.cfg

@@ -1,10 +1,13 @@
 [bdist_wheel]
 universal = 1
 
-[tool:pytest]
-addopts = -p no:logging
-
 [flake8]
-exclude = colorlog/tests
+exclude = .tox,colorlog/tests,venv
 ignore = W503
 max-line-length = 88
+
+[tool:pytest]
+addopts = -p no:logging
+
+[mypy]
+ignore_missing_imports = True

setup.py

@@ -12,7 +12,10 @@
     license="MIT License",
     packages=["colorlog"],
     setup_requires=["setuptools>=38.6.0"],
-    extras_require={':sys_platform=="win32"': ["colorama"]},
+    extras_require={
+        ':sys_platform=="win32"': ["colorama"],
+        "development": ["black", "flake8", "mypy", "pytest", "types-colorama"],
+    },
     classifiers=[
         "Development Status :: 5 - Production/Stable",
         "Environment :: Console",