Merge pull request #41 from StasEvseev/feature/custom_hooks_for_time_…

…execution

NEW: Added specifying set of hooks for time_execution decorator

Dockerfile

@@ -1,4 +1,4 @@
-FROM kpndigital/tox:latest
+FROM kpndigital/tox:py27_py35
 
 WORKDIR /app
 

Makefile

@@ -32,7 +32,7 @@ clean: pyclean docsclean
 	@rm -rf .tox
 
 venv:
-	@virtualenv -p python2.7 venv
+	@python3.6 -m venv venv
 	@$(PIP) install -U "pip>=7.0" -q
 	@$(PIP) install -r $(DEPS)
 

README.rst

@@ -144,7 +144,7 @@ It's also possible to run backend in different thread with logic behind it, to s
 
 For example:
 
- .. code-block:: python
+.. code-block:: python
 
     from time_execution import settings, time_execution
     from time_execution.backends.threaded import ThreadedBackend
@@ -239,6 +239,49 @@ See the following example how to setup hooks.
     # Configure the time_execution decorator, but now with hooks
     settings.configure(backends=[backend], hooks=[my_hook])
 
+
+There is also possibility to create decorator with custom set of hooks. It is needed for example to track `celery` tasks.
+
+.. code-block:: python
+
+    from multiprocessing import current_process
+    # Hook for celery tasks
+    def celery_hook(response, exception, metric, func_args, func_kwargs):
+        """
+        Add celery worker-specific details into response.
+        """
+        p = current_process()
+        hook = {
+            'name': metric.get('name'),
+            'value': metric.get('value'),
+            'success': exception is None,
+            'process_name': p.name,
+            'process_pid': p.pid,
+        }
+        return hook
+
+    # Create time_execution decorator with extra hooks
+    time_execution_celery = time_execution(extra_hooks=[celery_hook])
+
+    @celery.task
+    @time_execution_celery
+    def celery_task(self, **kwargs):
+        return True
+
+    # Or do it in place where it is needed
+    @celery.task
+    @time_execution(extra_hooks=[celery_hook])
+    def celery_task(self, **kwargs):
+        return True
+
+    # Or override default hooks by custom ones. Just setup `disable_default_hooks` flag
+    @celery.task
+    @time_execution(extra_hooks=[celery_hook], disable_default_hooks=True)
+    def celery_task(self, **kwargs):
+        return True
+
+
+
 Manually sending metrics
 ------------------------
 

setup.cfg

@@ -1,6 +1,6 @@
 [flake8]
 exclude = build,.git
-ignore = E902
+ignore = E902,W503
 max-line-length = 119
 
 [isort]

tests/test_hooks.py

@@ -1,5 +1,6 @@
 import unittest
 
+from fqn_decorators import get_fqn
 from tests.conftest import go
 from time_execution import settings, time_execution
 from time_execution.backends.base import BaseMetricsBackend
@@ -13,7 +14,138 @@ def write(self, name, **data):
         return self.callback(name, **data)
 
 
+class CollectorBackend(BaseMetricsBackend):
+    def __init__(self):
+        self.metrics = []
+
+    def write(self, name, **data):
+        self.metrics.append({name: data})
+
+    def clean(self):
+        self.metrics = []
+
+
+def local_hook(**kwargs):
+    return dict(local_hook_key='local hook value')
+
+
+def global_hook(**kwargs):
+    return dict(global_hook_key='global hook value')
+
+
 class TestTimeExecution(unittest.TestCase):
+    def test_custom_hook(self):
+        with settings(backends=[CollectorBackend()], hooks=[global_hook]):
+            collector = settings.backends[0]
+
+            @time_execution(extra_hooks=[local_hook])
+            def func_local_hook(*args, **kwargs):
+                return True
+
+            func_local_hook()
+            assert len(collector.metrics) == 1
+            metadata = collector.metrics[0][func_local_hook.get_fqn()]
+            assert metadata["local_hook_key"] == "local hook value"
+            assert metadata["global_hook_key"] == "global hook value"
+            collector.clean()
+
+            @time_execution(extra_hooks=[local_hook], disable_default_hooks=True)
+            def func_local_hook_disable_default_hooks(*args, **kwargs):
+                return True
+
+            func_local_hook_disable_default_hooks()
+            assert len(collector.metrics) == 1
+            metadata = collector.metrics[0][func_local_hook_disable_default_hooks.get_fqn()]
+            assert metadata["local_hook_key"] == "local hook value"
+            assert "global_hook_key" not in metadata
+            collector.clean()
+
+            @time_execution
+            def func_global_hook(*args, **kwargs):
+                return True
+
+            func_global_hook()
+            assert len(collector.metrics) == 1
+            metadata = collector.metrics[0][func_global_hook.get_fqn()]
+            assert metadata["global_hook_key"] == "global hook value"
+            assert "local_hook_key" not in metadata
+            collector.clean()
+
+            class ClassNoHooks:
+                @time_execution(extra_hooks=[local_hook])
+                def method_local_hook(self):
+                    return True
+
+                @time_execution
+                def method_global_hook(self):
+                    return True
+
+            ClassNoHooks().method_local_hook()
+            assert len(collector.metrics) == 1
+            metadata = collector.metrics[0][ClassNoHooks().method_local_hook.get_fqn()]
+            assert metadata["global_hook_key"] == "global hook value"
+            assert metadata["local_hook_key"] == "local hook value"
+            collector.clean()
+
+            ClassNoHooks().method_global_hook()
+            assert len(collector.metrics) == 1
+            metadata = collector.metrics[0][ClassNoHooks().method_global_hook.get_fqn()]
+            assert metadata["global_hook_key"] == "global hook value"
+            assert "local_hook_key" not in metadata
+            collector.clean()
+
+            @time_execution(extra_hooks=[local_hook])
+            class ClassLocalHook:
+                def method(self):
+                    return True
+
+                @time_execution
+                def method_global_hook(self):
+                    return True
+
+            ClassLocalHook().method()
+            assert len(collector.metrics) == 1
+            metadata = collector.metrics[0][get_fqn(ClassLocalHook)]
+            assert metadata["global_hook_key"] == "global hook value"
+            assert metadata["local_hook_key"] == "local hook value"
+            collector.clean()
+
+            ClassLocalHook().method_global_hook()
+            assert len(collector.metrics) == 2
+            metadata_class = collector.metrics[0][get_fqn(ClassLocalHook)]
+            assert metadata_class["local_hook_key"] == "local hook value"
+            assert metadata_class["global_hook_key"] == "global hook value"
+            metadata = collector.metrics[1][ClassLocalHook().method_global_hook.get_fqn()]
+            assert metadata["global_hook_key"] == "global hook value"
+            assert "local_hook_key" not in metadata
+            collector.clean()
+
+            @time_execution
+            class ClassGlobalHook:
+                def method(self):
+                    return True
+
+                @time_execution(extra_hooks=[local_hook])
+                def method_local_hook(self):
+                    return True
+
+            ClassGlobalHook().method()
+            assert len(collector.metrics) == 1
+            metadata = collector.metrics[0][get_fqn(ClassGlobalHook)]
+            assert metadata["global_hook_key"] == "global hook value"
+            assert "local_hook_key" not in metadata
+            collector.clean()
+
+            ClassGlobalHook().method_local_hook()
+            assert len(collector.metrics) == 2
+            metadata_class = collector.metrics[0][get_fqn(ClassGlobalHook)]
+            assert metadata_class["global_hook_key"] == "global hook value"
+            assert "local_hook_key" not in metadata_class
+            metadata = collector.metrics[1][ClassGlobalHook().method_local_hook.get_fqn()]
+            assert metadata["global_hook_key"] == "global hook value"
+            assert metadata["local_hook_key"] == "local hook value"
+            collector.clean()
+
     def test_hook(self):
         def test_args(**kwargs):
             self.assertIn('response', kwargs)

time_execution/decorator.py

@@ -25,9 +25,9 @@ def write_metric(name, **metric):
         backend.write(name, **metric)
 
 
-def _apply_hooks(**kwargs):
+def _apply_hooks(hooks, **kwargs):
     metadata = dict()
-    for hook in settings.hooks:
+    for hook in hooks:
         hook_result = hook(**kwargs)
         if hook_result:
             metadata.update(hook_result)
@@ -56,9 +56,16 @@ def after(self):
         if origin:
             metric['origin'] = origin
 
+        hooks = self.params.get('extra_hooks', [])
+        disable_default_hooks = self.params.get('disable_default_hooks', False)
+
+        if not disable_default_hooks:
+            hooks = settings.hooks + hooks
+
         # Apply the registered hooks, and collect the metadata they might
         # return to be stored with the metrics
         metadata = _apply_hooks(
+            hooks=hooks,
             response=self.result,
             exception=self.get_exception(),
             metric=metric,

tox.ini

@@ -4,7 +4,7 @@ addopts=
 
 [tox]
 envlist=
-    py27,py35,isort-check,isort-fix,lint,docs
+    py27,py35,py36,isort-check,isort-fix,lint,docs
 skipsdist=
     true
 
@@ -26,15 +26,15 @@ setenv=
 
 [testenv:lint]
 basepython=
-    python3.5
+    python3.6
 commands=
     flake8 time_execution tests --exclude=time_execution/__init__.py
 deps=
     flake8
 
 [testenv:docs]
 basepython=
-    python3.5
+    python3.6
 commands=
     python {toxinidir}/docs/apidoc.py -T -M -d 2 -o {toxinidir}/docs/api {toxinidir}/time_execution
     sphinx-build -W -b html {toxinidir}/docs {toxinidir}/docs/_build/html