Add shutdown method

Author: aminalaee

src/picologging/__init__.py

@@ -1,4 +1,6 @@
+import atexit
 import io
+import itertools
 import os
 import sys
 import warnings
@@ -503,3 +505,42 @@ def makeLogRecord(dict):
     for k, v in dict.items():
         setattr(rv, k, v)
     return rv
+
+
+def shutdown(handlerList=None):
+    """
+    Perform any cleanup actions in the logging system (e.g. flushing
+    buffers).
+    Should be called at application exit.
+    """
+    handlers = [logger.handlers for logger in root.manager.loggerDict.values()] + [
+        root.handlers
+    ]
+    handlerList = list(itertools.chain(*handlers))
+
+    for h in reversed(handlerList):
+        # errors might occur, for example, if files are locked
+        # we just ignore them if raiseExceptions is not set
+        try:
+            if h:
+                try:
+                    h.acquire()
+                    # MemoryHandlers might not want to be flushed on close,
+                    # but circular imports prevent us scoping this to just
+                    # those handlers.  hence the default to True.
+                    if getattr(h, "flushOnClose", True):
+                        h.flush()
+                    h.close()
+                except (OSError, ValueError):
+                    # Ignore errors which might be caused
+                    # because handlers have been closed but
+                    # references to them are still around at
+                    # application exit.
+                    pass
+                finally:
+                    h.release()
+        except:  # ignore everything, as we're shutting down
+            raise
+
+
+atexit.register(shutdown)

src/picologging/__init__.pyi

@@ -3,7 +3,17 @@ from io import TextIOWrapper
 from multiprocessing import Manager
 from string import Template
 from types import TracebackType
-from typing import Any, Generic, Optional, Pattern, TextIO, TypeVar, Union, overload
+from typing import (
+    Any,
+    Generic,
+    Optional,
+    Pattern,
+    TextIO,
+    TypeVar,
+    Union,
+    overload,
+    Sequence,
+)
 
 from _typeshed import StrPath, SupportsWrite
 from typing_extensions import Literal, TypeAlias
@@ -362,3 +372,6 @@ BASIC_FORMAT: str
 
 def getLevelName(level: _Level) -> Any: ...
 def makeLogRecord(dict: Mapping[str, object]) -> LogRecord: ...
+def shutdown(
+    handlerList: Sequence[Any] = ...,
+) -> None: ...  # handlerList is undocumented

tests/unit/test_picologging.py

@@ -1,9 +1,9 @@
-import logging
 import sys
 
 import pytest
 
 import picologging
+from picologging.handlers import BufferingHandler
 
 levels = [
     (picologging.DEBUG, "DEBUG"),
@@ -146,3 +146,15 @@ def test_make_log_record():
     log_record = picologging.makeLogRecord({"levelno": picologging.WARNING})
 
     assert log_record.levelno == picologging.WARNING
+
+
+def test_shutdown():
+    handler = BufferingHandler(capacity=1)
+    logger = picologging.getLogger("test")
+    logger.setLevel(picologging.DEBUG)
+    logger.addHandler(handler)
+    logger.debug("test")
+
+    picologging.shutdown()
+
+    assert handler.buffer == []