From cee73c66c5c34a37a166d7eb8815f896e62b4def Mon Sep 17 00:00:00 2001
From: Dmitry Bashkatov <dbashkatov@gmail.com>
Date: Sun, 5 May 2024 20:06:23 +0300
Subject: [PATCH] add more features for template context

---
 kew/config.py        | 29 +++++------------------------
 kew/handlers/base.py | 13 +++++++++++--
 kew/matcher.py       | 24 ++++++++++++++++++++++++
 3 files changed, 40 insertions(+), 26 deletions(-)
 create mode 100644 kew/matcher.py

diff --git a/kew/config.py b/kew/config.py
index 0125445..b0af567 100644
--- a/kew/config.py
+++ b/kew/config.py
@@ -1,12 +1,9 @@
-import re
-from fnmatch import fnmatch
-
 import yaml
 
 from .handlers import get_handler
+from .matcher import prepare_patterns, event_does_match_any_pattern
 
 VERSION = 1
-extended_pat_re = re.compile(r'.*\(.*\)')
 
 
 def load_config(filename):
@@ -36,30 +33,14 @@ def __init__(self, sinks, mappings):
 class Mapping:
     def __init__(self, sink, include, exclude):
         self.sink = sink
-        self.include = _clean_patterns(include)
-        self.exclude = _clean_patterns(exclude)
+        self.include = prepare_patterns(include)
+        self.exclude = prepare_patterns(exclude)
 
     def does_match(self, event):
         if self.include is not None:
-            if not any((fnmatch(event._formatted, pat) for pat in self.include)):
+            if not event_does_match_any_pattern(event, self.include):
                 return False
         if self.exclude is not None:
-            if any((fnmatch(event._formatted, pat) for pat in self.exclude)):
+            if event_does_match_any_pattern(event, self.exclude):
                 return False
         return True
-
-
-def _clean_patterns(pats):
-    if not pats:
-        return pats
-    return [b for b in (_clean_pattern(a) for a in pats) if b]
-
-
-def _clean_pattern(pat):
-    pat = pat.strip()
-    if not pat:
-        return pat
-
-    if not extended_pat_re.fullmatch(pat):
-        pat += '(*)'
-    return pat
diff --git a/kew/handlers/base.py b/kew/handlers/base.py
index a2432cf..ab51843 100644
--- a/kew/handlers/base.py
+++ b/kew/handlers/base.py
@@ -1,6 +1,7 @@
 import json
 
 from kew.format import format_involved_object, format_involved_object_kind, format_event_age, format_event_source
+from kew.matcher import prepare_patterns, event_does_match_any_pattern
 
 TEMPLATE_GLOBALS = {
     'escape_json': lambda s: json.dumps(s)[1:-1],
@@ -8,13 +9,21 @@
     'format_involved_object_kind': format_involved_object_kind,
     'format_event_age': format_event_age,
     'format_event_source': format_event_source,
+    'prepare_patterns': prepare_patterns,
+    'event_does_match_any_pattern': event_does_match_any_pattern,
 }
 
 
 class BaseHandler:
     def __init__(self, config):
+        self.globals = dict(TEMPLATE_GLOBALS)
         self.template_context_script = config.pop('template_context', None)
         self.template = config.pop('template', '{event.message}')
+        template_context_init_script = config.pop('template_context_init', None)
+        if template_context_init_script:
+            ctx_locals = {}
+            exec(template_context_init_script, self.globals, ctx_locals)
+            self.globals.update(ctx_locals)
 
     def format(self, event):
         context = {
@@ -22,9 +31,9 @@ def format(self, event):
         }
 
         if self.template_context_script:
-            exec(self.template_context_script, TEMPLATE_GLOBALS, context)
+            exec(self.template_context_script, self.globals, context)
 
-        return eval(f'f{self.template!r}', TEMPLATE_GLOBALS, context)
+        return eval(f'f{self.template!r}', self.globals, context)
 
     def __call__(self, event):
         raise NotImplementedError
diff --git a/kew/matcher.py b/kew/matcher.py
new file mode 100644
index 0000000..a3fb765
--- /dev/null
+++ b/kew/matcher.py
@@ -0,0 +1,24 @@
+import re
+from fnmatch import fnmatch
+
+extended_pat_re = re.compile(r'.*\(.*\)')
+
+
+def prepare_patterns(pats):
+    if not pats:
+        return pats
+    return [b for b in (_clean_pattern(a) for a in pats) if b]
+
+
+def _clean_pattern(pat):
+    pat = pat.strip()
+    if not pat:
+        return pat
+
+    if not extended_pat_re.fullmatch(pat):
+        pat += '(*)'
+    return pat
+
+
+def event_does_match_any_pattern(event, patterns):
+    return any((fnmatch(event._formatted, pat) for pat in patterns))