Add logger decorator to catch exceptions

loguru/__init__.py

@@ -175,15 +175,15 @@ def format_path(self):
         now_ = now()
         patch_datetime_file(now_)
 
-        info = {
+        record = {
             "time": now_,
             "start_time": self.start_time,
             "rotation_time": self.rotation_time,
             "n": self.created,
             "n+1": self.created + 1,
         }
 
-        return self.path.format_map(info)
+        return self.path.format_map(record)
 
     @staticmethod
     def make_regex_file_name(file_name):
@@ -601,6 +601,52 @@ def stop(self, handler_id=None):
 
         return 0
 
+    def catch(self, *args, **kwargs):
+
+        def catch_decorator(wrapped_function,
+                            message="An error has been caught in function '{function}', "
+                                    "process '{process.name}' ({process.id}), "
+                                    "thread '{thread.name}' ({thread.id}):", *,
+                                    level=None, re_raise=False):
+
+            if level is not None:
+                # TODO: Call log function accordingly
+                raise NotImplementedError
+
+            @functools.wraps(wrapped_function)
+            def catch_wrapper(*args, **kwargs):
+                try:
+                    wrapped_function(*args, **kwargs)
+                except BaseException as e:
+                    thread = current_thread()
+                    thread_recattr = ThreadRecattr(thread.ident)
+                    thread_recattr.id, thread_recattr.name = thread.ident, thread.name
+
+                    process = current_process()
+                    process_recattr = ProcessRecattr(process.ident)
+                    process_recattr.id, process_recattr.name = process.ident, process.name
+
+                    function_name = wrapped_function.__name__
+
+                    record = {
+                        'process': process_recattr,
+                        'thread': thread_recattr,
+                        'function': function_name,
+                    }
+
+                    # TODO: Use the stacktrace from 'e' rather than calling sys.exc_info() in
+                    self.exception(message.format_map(record))
+
+                    if re_raise:
+                        raise
+
+            return catch_wrapper
+
+        if not kwargs and len(args) == 1 and callable(args[0]):
+            return catch_decorator(args[0])
+        else:
+            return lambda f: catch_decorator(f, *args, **kwargs)
+
     @staticmethod
     def make_log_function(level, log_exception=False):
 

tests/test_catch_wrapper.py

@@ -0,0 +1,60 @@
+import pytest
+
+@pytest.mark.parametrize('args, kwargs', [
+    ([], {}),
+    ([2, 0], {}),
+    ([4], {'b': 0}),
+    ([], {'a': 8}),
+])
+def test_wrapped(logger, writer, args, kwargs):
+    logger.log_to(writer)
+
+    @logger.catch
+    def c(a=1, b=0):
+        a / b
+
+    c(*args, **kwargs)
+
+    assert writer.read().endswith('ZeroDivisionError: division by zero\n')
+
+@pytest.mark.parametrize('better_exceptions', [True, False])
+def test_wrapped_better_exceptions(logger, writer, better_exceptions):
+    logger.log_to(writer, better_exceptions=better_exceptions)
+
+    @logger.catch()
+    def c():
+        a = 2
+        b = 0
+        a / b
+
+    c()
+
+    length = len(writer.read().splitlines())
+
+    if better_exceptions:
+        assert length == 15
+    else:
+        assert length == 7
+
+def test_custom_message(logger, writer):
+    logger.log_to(writer, format='{message}')
+
+    @logger.catch(message='An error occured:')
+    def a():
+        1 / 0
+
+    a()
+
+    assert writer.read().startswith('An error occured:\n')
+
+def test_re_raise(logger, writer):
+    logger.log_to(writer)
+
+    @logger.catch(re_raise=True)
+    def a():
+        1 / 0
+
+    with pytest.raises(ZeroDivisionError):
+        a()
+
+    assert writer.read().endswith('ZeroDivisionError: division by zero\n')