Merge pull request fluent#101 from arcivanov/issue_100

repeatedly · web-flow · commit d99bca7dd82c · 2017-12-06T09:12:24.000+09:00
Allow formatter to use attribute exclusion set
diff --git a/fluent/handler.py b/fluent/handler.py
@@ -29,9 +29,14 @@ class FluentRecordFormatter(logging.Formatter, object):
         key is not found. Put None if not found.
     :param format_json: if True, will attempt to parse message as json. If not,
         will use message as-is. Defaults to True
+    :param exclude_attrs: switches this formatter into a mode where all attributes
+        except the ones specified by `exclude_attrs` are logged with the record as is.
+        If `None`, operates as before, otherwise `fmt` is ignored.
+        Can be a `list`, `tuple` or a `set`.
     """
 
-    def __init__(self, fmt=None, datefmt=None, style='%', fill_missing_fmt_key=False, format_json=True):
+    def __init__(self, fmt=None, datefmt=None, style='%', fill_missing_fmt_key=False, format_json=True,
+                 exclude_attrs=None):
         super(FluentRecordFormatter, self).__init__(None, datefmt)
 
         if sys.version_info[0:2] >= (3, 2) and style != '%':
@@ -55,10 +60,15 @@ def __init__(self, fmt=None, datefmt=None, style='%', fill_missing_fmt_key=False
                 'sys_module': '%(module)s',
             }
 
-        if not fmt:
-            self._fmt_dict = basic_fmt_dict
+        if exclude_attrs is not None:
+            self._exc_attrs = set(exclude_attrs)
+            self._fmt_dict = None
         else:
-            self._fmt_dict = fmt
+            self._exc_attrs = None
+            if not fmt:
+                self._fmt_dict = basic_fmt_dict
+            else:
+                self._fmt_dict = fmt
 
         if format_json:
             self._format_msg = self._format_msg_json
@@ -81,25 +91,33 @@ def format(self, record):
 
         # Apply format
         data = {}
-        for key, value in self._fmt_dict.items():
-            try:
-                if self.__style:
-                    value = self.__style(value).format(record)
-                else:
-                    value = value % record.__dict__
-            except KeyError as exc:
-                value = None
-                if not self.fill_missing_fmt_key:
-                    raise exc
-
-            data[key] = value
+        if self._exc_attrs is not None:
+            for key, value in record.__dict__.items():
+                if key not in self._exc_attrs:
+                    data[key] = value
+        else:
+            for key, value in self._fmt_dict.items():
+                try:
+                    if self.__style:
+                        value = self.__style(value).format(record)
+                    else:
+                        value = value % record.__dict__
+                except KeyError as exc:
+                    value = None
+                    if not self.fill_missing_fmt_key:
+                        raise exc
+
+                data[key] = value
 
         self._structuring(data, record)
         return data
 
     def usesTime(self):
-        return any([value.find('%(asctime)') >= 0
-                    for value in self._fmt_dict.values()])
+        if self._exc_attrs is not None:
+            return super(FluentRecordFormatter, self).usesTime()
+        else:
+            return any([value.find('%(asctime)') >= 0
+                        for value in self._fmt_dict.values()])
 
     def _structuring(self, data, record):
         """ Melds `msg` into `data`.
diff --git a/tests/test_handler.py b/tests/test_handler.py
@@ -1,11 +1,10 @@
-# -*- coding: utf-8 -*-
+#  -*- coding: utf-8 -*-
 
 import logging
 import sys
 import unittest
 
 import fluent.handler
-
 from tests import mockserver
 
 
@@ -63,6 +62,42 @@ def test_custom_fmt(self):
         self.assertTrue('lineno' in data[0][2])
         self.assertTrue('emitted_at' in data[0][2])
 
+    def test_exclude_attrs(self):
+        handler = fluent.handler.FluentHandler('app.follow', port=self._port)
+
+        logging.basicConfig(level=logging.INFO)
+        log = logging.getLogger('fluent.test')
+        handler.setFormatter(
+            fluent.handler.FluentRecordFormatter(exclude_attrs=[])
+        )
+        log.addHandler(handler)
+        log.info({'sample': 'value'})
+        handler.close()
+
+        data = self.get_data()
+        self.assertTrue('name' in data[0][2])
+        self.assertEqual('fluent.test', data[0][2]['name'])
+        self.assertTrue('lineno' in data[0][2])
+
+    def test_exclude_attrs_with_extra(self):
+        handler = fluent.handler.FluentHandler('app.follow', port=self._port)
+
+        logging.basicConfig(level=logging.INFO)
+        log = logging.getLogger('fluent.test')
+        handler.setFormatter(
+            fluent.handler.FluentRecordFormatter(exclude_attrs=[])
+        )
+        log.addHandler(handler)
+        log.info("Test with value '%s'", "test value", extra={"x": 1234})
+        handler.close()
+
+        data = self.get_data()
+        self.assertTrue('name' in data[0][2])
+        self.assertEqual('fluent.test', data[0][2]['name'])
+        self.assertTrue('lineno' in data[0][2])
+        self.assertEqual("Test with value 'test value'", data[0][2]['message'])
+        self.assertEqual(1234, data[0][2]['x'])
+
     @unittest.skipUnless(sys.version_info[0:2] >= (3, 2), 'supported with Python 3.2 or above')
     def test_custom_fmt_with_format_style(self):
         handler = fluent.handler.FluentHandler('app.follow', port=self._port)
@@ -135,7 +170,7 @@ def test_custom_field_fill_missing_fmt_key_is_true(self):
             fluent.handler.FluentRecordFormatter(fmt={
                 'name': '%(name)s',
                 'custom_field': '%(custom_field)s'
-                },
+            },
                 fill_missing_fmt_key=True
             )
         )
