Make the package typed: ship py.typed, annotate the public API, type-check in CI

.github/workflows/main.yml

@@ -29,3 +29,18 @@ jobs:
         pip install tox tox-gh-actions
     - name: Test with tox
       run: tox
+
+  typecheck:
+    runs-on: ubuntu-24.04
+    steps:
+    - uses: actions/checkout@v7
+    - name: Set up Python
+      uses: actions/setup-python@v6
+      with:
+        python-version: '3.x'
+    - name: Install dependencies
+      run: |
+        python -m pip install --upgrade pip
+        pip install -r requirements.txt -r test-requirements.txt
+    - name: Type-check with mypy
+      run: mypy

.gitignore

@@ -5,6 +5,7 @@ __pycache__
 htmlcov
 .coverage
 .tox
+.idea
 *.egg-info
 build/*
 dist/*

logtail/compat.py

@@ -4,4 +4,4 @@
 try:
     import queue
 except ImportError:
-    import Queue as queue
+    import Queue as queue # type: ignore[import-not-found, no-redef]

logtail/flusher.py

@@ -3,13 +3,14 @@
 
 import threading
 import time
+from typing import Any, Callable, Optional
 
 from .compat import queue
 
 RETRY_SCHEDULE = (1, 10, 60)  # seconds
 
 
-class TransportFrame(dict):
+class TransportFrame(dict[str, Any]):
     """A record logged by a flush worker itself, i.e. urllib3 describing our own upload.
 
     It rides along with the next batch of regular records but never triggers an upload
@@ -18,7 +19,7 @@ class TransportFrame(dict):
 
 
 class FlushWorker(threading.Thread):
-    def __init__(self, upload, pipe, buffer_capacity, flush_interval, check_interval):
+    def __init__(self, upload: Callable[[list[Any]], Any], pipe: 'queue.Queue[Any]', buffer_capacity: int, flush_interval: float, check_interval: float) -> None:
         threading.Thread.__init__(self)
         self.parent_thread = threading.current_thread()
         self.upload = upload
@@ -30,11 +31,11 @@ def __init__(self, upload, pipe, buffer_capacity, flush_interval, check_interval
         self._flushing = False
         self._clean = True
 
-    def run(self):
+    def run(self) -> None:
         while self.should_run:
             self.step()
 
-    def _is_parent_alive(self):
+    def _is_parent_alive(self) -> bool:
         try:
             return self.parent_thread.is_alive()
         except RuntimeError:
@@ -44,10 +45,10 @@ def _is_parent_alive(self):
             # interpreter shutdown). Treat that as "no longer alive".
             return False
 
-    def step(self):
+    def step(self) -> None:
         last_flush = time.time()
         time_remaining = _initial_time_remaining(self.flush_interval)
-        frame = []
+        frame: list[Any] = []
         self._clean = True
 
         # If the parent thread has exited but there are still outstanding
@@ -80,7 +81,7 @@ def step(self):
         # request fails in a way that can be retried, it is retried with an
         # exponential backoff in between attempts.
         if frame and not all(isinstance(entry, TransportFrame) for entry in frame):
-            response = None
+            response: Any = None
             for delay in RETRY_SCHEDULE + (None, ):
                 response = self.upload(frame)
                 if not _should_retry(response.status_code):
@@ -100,7 +101,7 @@ def step(self):
         if shutdown and self.pipe.empty():
             self.should_run = False
 
-    def flush(self, timeout=None):
+    def flush(self, timeout: Optional[float] = None) -> bool:
         """Block until the worker has drained the queue.
 
         If ``timeout`` is given (in seconds), return after at most that much
@@ -125,19 +126,19 @@ def flush(self, timeout=None):
         finally:
             self._flushing = False
 
-def in_flush_worker():
+def in_flush_worker() -> bool:
     return isinstance(threading.current_thread(), FlushWorker)
 
 
-def _initial_time_remaining(flush_interval):
+def _initial_time_remaining(flush_interval: float) -> float:
     return flush_interval
 
 
-def _calculate_time_remaining(last_flush, flush_interval):
+def _calculate_time_remaining(last_flush: float, flush_interval: float) -> float:
     elapsed = time.time() - last_flush
     time_remaining = max(flush_interval - elapsed, 0)
     return time_remaining
 
 
-def _should_retry(status_code):
+def _should_retry(status_code: int) -> bool:
     return 500 <= status_code < 600

logtail/formatter.py

@@ -2,21 +2,22 @@
 from __future__ import print_function, unicode_literals
 import logging
 import json
+from typing import Any, Callable, Optional
 
-from .helpers import DEFAULT_CONTEXT
+from .helpers import DEFAULT_CONTEXT, LogtailContext
 from .frame import create_frame
 
 
 class LogtailFormatter(logging.Formatter):
     def __init__(self,
-                 context=DEFAULT_CONTEXT,
-                 json_default=None,
-                 json_encoder=None):
+                 context: LogtailContext = DEFAULT_CONTEXT,
+                 json_default: Optional[Callable[[Any], Any]] = None,
+                 json_encoder: Optional[type[json.JSONEncoder]] = None) -> None:
         self.context = context
         self.json_default = json_default
         self.json_encoder = json_encoder
 
-    def format(self, record):
+    def format(self, record: logging.LogRecord) -> str:
         # Because the formatter does not have an underlying format string for
         # which `extra` may be used to substitute arguments (see
         # https://docs.python.org/2/library/logging.html#logging.debug ), we

logtail/frame.py

@@ -1,24 +1,29 @@
 # coding: utf-8
 from __future__ import print_function, unicode_literals
 from datetime import datetime, timezone
-
+import logging
 from os import path
+from typing import Any, Optional
 import __main__
 
-def create_frame(record, message, context, include_extra_attributes=False):
+from .helpers import LogtailContext
+
+def create_frame(record: logging.LogRecord, message: str, context: LogtailContext, include_extra_attributes: bool = False) -> dict[str, Any]:
     r = record.__dict__
     # Django sends a request object in the record, which is not JSON serializable
     if "request" in r and not isinstance(r["request"], (dict, list, bool, int, float, str)) :
         del r["request"]
-    frame = {}
+    frame: dict[str, Any] = {}
     frame['dt'] = datetime.fromtimestamp(r['created'], timezone.utc).isoformat()
     frame['level'] = _levelname(r['levelname'])
     frame['severity'] = int(r['levelno'] / 10)
     frame['message'] = message
-    frame['context'] = ctx = {}
+    ctx: dict[str, Any] = {}
+    frame['context'] = ctx
 
     # Runtime context
-    ctx['runtime'] = runtime = {}
+    runtime: dict[str, Any] = {}
+    ctx['runtime'] = runtime
     runtime['function'] = r['funcName']
     runtime['file'] = _relative_to_main_module_if_possible(r['pathname'])
     runtime['line'] = r['lineno']
@@ -27,7 +32,8 @@ def create_frame(record, message, context, include_extra_attributes=False):
     runtime['logger_name'] = r['name']
 
     # Runtime context
-    ctx['system'] = system = {}
+    system: dict[str, Any] = {}
+    ctx['system'] = system
     system['pid'] = r['process']
     system['process_name'] = r['processName']
 
@@ -41,14 +47,14 @@ def create_frame(record, message, context, include_extra_attributes=False):
 
     return _remove_circular_dependencies(frame)
 
-def _parse_custom_events(record, include_extra_attributes):
+def _parse_custom_events(record: logging.LogRecord, include_extra_attributes: bool) -> dict[str, Any]:
     default_keys = {
         'args', 'asctime', 'created', 'exc_info', 'exc_text', 'pathname',
         'funcName', 'levelname', 'levelno', 'lineno', 'module', 'msecs',
         'message', 'msg', 'name', 'process', 'processName',
         'relativeCreated', 'thread', 'threadName'
     }
-    events = {}
+    events: dict[str, Any] = {}
     for key, val in record.__dict__.items():
         if key in default_keys:
             continue
@@ -57,7 +63,7 @@ def _parse_custom_events(record, include_extra_attributes):
         events[key] = val
     return events
 
-def _remove_circular_dependencies(obj, memo=None):
+def _remove_circular_dependencies(obj: Any, memo: Optional[set[int]] = None) -> Any:
     if memo is None:
         memo = set()
 
@@ -85,14 +91,14 @@ def _remove_circular_dependencies(obj, memo=None):
     else:
         return obj
 
-def _levelname(level):
+def _levelname(level: str) -> str:
     return level.lower()
 
-def _relative_to_main_module_if_possible(pathname):
+def _relative_to_main_module_if_possible(pathname: str) -> str:
     has_main_module = hasattr(__main__, '__file__')
     return _relative_to_main_module(pathname) if has_main_module else pathname
 
-def _relative_to_main_module(pathname):
+def _relative_to_main_module(pathname: str) -> str:
     try:
         return path.relpath(pathname, path.dirname(__main__.__file__))
     except ValueError:

logtail/handler.py

@@ -4,9 +4,10 @@
 import json
 import os
 import weakref
+from typing import Any, Optional, Union
 
 from .compat import queue
-from .helpers import DEFAULT_CONTEXT
+from .helpers import DEFAULT_CONTEXT, LogtailContext
 from .flusher import FlushWorker, TransportFrame, in_flush_worker
 from .uploader import Uploader
 from .frame import create_frame
@@ -22,31 +23,31 @@
 DEFAULT_FLUSH_TIMEOUT = 30
 
 
-_handlers = weakref.WeakSet()
+_handlers: 'weakref.WeakSet[LogtailHandler]' = weakref.WeakSet()
 
 
 class LogtailHandler(logging.Handler):
     def __init__(self,
-                 source_token,
-                 host=DEFAULT_HOST,
-                 buffer_capacity=DEFAULT_BUFFER_CAPACITY,
-                 flush_interval=DEFAULT_FLUSH_INTERVAL,
-                 check_interval=DEFAULT_CHECK_INTERVAL,
-                 raise_exceptions=DEFAULT_RAISE_EXCEPTIONS,
-                 drop_extra_events=DEFAULT_DROP_EXTRA_EVENTS,
-                 include_extra_attributes=DEFAULT_INCLUDE_EXTRA_ATTRIBUTES,
-                 context=DEFAULT_CONTEXT,
-                 timeout=DEFAULT_TIMEOUT,
-                 flush_timeout=DEFAULT_FLUSH_TIMEOUT,
-                 level=logging.NOTSET):
+                 source_token: str,
+                 host: str = DEFAULT_HOST,
+                 buffer_capacity: int = DEFAULT_BUFFER_CAPACITY,
+                 flush_interval: float = DEFAULT_FLUSH_INTERVAL,
+                 check_interval: float = DEFAULT_CHECK_INTERVAL,
+                 raise_exceptions: bool = DEFAULT_RAISE_EXCEPTIONS,
+                 drop_extra_events: bool = DEFAULT_DROP_EXTRA_EVENTS,
+                 include_extra_attributes: bool = DEFAULT_INCLUDE_EXTRA_ATTRIBUTES,
+                 context: LogtailContext = DEFAULT_CONTEXT,
+                 timeout: Union[float, tuple[float, float], None] = DEFAULT_TIMEOUT,
+                 flush_timeout: Optional[float] = DEFAULT_FLUSH_TIMEOUT,
+                 level: Union[int, str] = logging.NOTSET) -> None:
         super(LogtailHandler, self).__init__(level=level)
         self.source_token = source_token
         if host.startswith('https://') or host.startswith('http://'):
             self.host = host
         else:
             self.host = "https://" + host
         self.context = context
-        self.pipe = queue.Queue(maxsize=buffer_capacity)
+        self.pipe: 'queue.Queue[dict[str, Any]]' = queue.Queue(maxsize=buffer_capacity)
         self.uploader = Uploader(self.source_token, self.host, timeout)
         self.drop_extra_events = drop_extra_events
         self.include_extra_attributes = include_extra_attributes
@@ -57,10 +58,10 @@ def __init__(self,
         self.flush_timeout = flush_timeout
         self.dropcount = 0
         # Do not initialize the flush thread yet because it causes issues on Render.
-        self.flush_thread = None
+        self.flush_thread: Optional[FlushWorker] = None
         _handlers.add(self)
 
-    def ensure_flush_thread_alive(self):
+    def ensure_flush_thread_alive(self) -> None:
         if self.flush_thread and self.flush_thread.is_alive():
             return
 
@@ -73,7 +74,7 @@ def ensure_flush_thread_alive(self):
         )
         self.flush_thread.start()
 
-    def emit(self, record):
+    def emit(self, record: logging.LogRecord) -> None:
         try:
             self.ensure_flush_thread_alive()
 
@@ -94,10 +95,10 @@ def emit(self, record):
             if self.raise_exceptions:
                 raise e
 
-    def flush(self):
+    def flush(self) -> None:
         if not (self.flush_thread and self.flush_thread.is_alive()):
             return
-        if logging._lock._is_owned():
+        if logging._lock._is_owned():  # type: ignore[attr-defined]  # private, see below
             # logging.config.dictConfig() flushes the handlers it replaces while holding the global
             # logging lock, and the worker's upload needs that lock whenever urllib3's logger has a
             # level-cache miss, so waiting here could only deadlock. Nothing is lost: the worker
@@ -106,7 +107,7 @@ def flush(self):
         if not self.flush_thread.flush(timeout=self.flush_timeout):
             print('Gave up waiting for Better Stack uploads after {}s, logs are still buffered'.format(self.flush_timeout))
 
-    def _reset_after_fork(self):
+    def _reset_after_fork(self) -> None:
         # A forked child inherits the parent's queue with whatever was still buffered in it,
         # a flush thread object whose thread does not exist in the child, and an HTTP session
         # whose socket it shares with the parent, so it starts over with fresh ones.
@@ -115,7 +116,7 @@ def _reset_after_fork(self):
         self.uploader.reset()
 
 
-def _reset_handlers_after_fork():
+def _reset_handlers_after_fork() -> None:
     for handler in _handlers:
         handler._reset_after_fork()