Refactored to make IPython/Djnango imports optional

.gitignore

@@ -102,3 +102,6 @@ venv.bak/
 
 # mypy
 .mypy_cache/
+
+# VS Code
+.vscode/

context_managers.py

@@ -1,23 +1,18 @@
 from abc import abstractmethod, ABC
 import copy
-from datetime import datetime
-
-from django import db
-from IPython.core import display
 
+from .hooks import SetHookMixin
 from .sample import Sampler
 from .utils import (
     format_options,
     generate_flame_graph_html,
-    humanize_timedelta,
 )
 
-FLAME_FILE_WIDTH = 1904
-FLAME_IPYTHON_WIDTH = 1904
+DEFAULT_FLAME_WIDTH = 1904
 SAMPLE_INTERVAL = 0.001
 
 
-class BaseFlame(ABC):
+class BaseFlame(SetHookMixin, ABC):
     """
     Base context manager to sample a block of code.
     """
@@ -34,24 +29,19 @@ def __enter__(self):
         track functional and db calls.
         """
         self.sampler.start()
-        self.started = datetime.now()
-        db.reset_queries()
+
+        # Run the set enter hook. Allows for different
+        # user defined functionality.
+        self.enter_section_hook()
 
     def __exit__(self, type, value, traceback):
         """
         On Context completion dump the summed functional calls for display.
         """
         self.sampler.stop()  # Always stop the sampler.
-        time_taken = humanize_timedelta(
-            datetime.now() - self.started, precision='ms'
-        )
-
-        num_sql_queries = len(db.connection.queries)
 
         default_args, default_kwargs = self.get_defaults()
-        default_kwargs['title'] = (
-            f'Made {num_sql_queries} SQL queries in {time_taken}.'
-        )
+        default_kwargs = self.exit_section_hook(default_kwargs)
 
         options = format_options(
             default_args, default_kwargs, self.option_args, self.option_kwargs
@@ -79,19 +69,29 @@ def __init__(self, path, **kwargs):
         super().__init__(**kwargs)
 
     def get_defaults(self):
-        return [], {'width': FLAME_FILE_WIDTH}
+        return [], {'width': DEFAULT_FLAME_WIDTH}
 
     def output(self, data):
         with open(self.path, 'w') as f:
             f.write(data)
 
+##
+# Support Ipython only if library is installed
+##
 
-class InlineFlame(BaseFlame):
-    """
-    Render FlameGraph HTML / SVG in an IPythonNotebook.
-    """
-    def get_defaults(self):
-        return [], {'width': FLAME_IPYTHON_WIDTH}
 
-    def output(self, data):
-        return display.display(display.SVG(data=data))
+try:
+    from IPython.core import display
+
+    class InlineFlame(BaseFlame):
+        """
+        Render FlameGraph HTML / SVG in an IPythonNotebook.
+        """
+        def get_defaults(self):
+            return [], {'width': DEFAULT_FLAME_WIDTH}
+
+        def output(self, data):
+            return display.display(display.SVG(data=data))
+
+except ImportError:
+    pass

hooks.py

@@ -0,0 +1,67 @@
+from datetime import datetime
+
+from humanize import naturaldelta
+
+##
+# Defined hook classes for framework independent integrations.
+# At the moment only supports Django, but implementers can substitute
+# their own mixins for different use cases.
+##
+
+
+class BaseHookMixin:
+    """
+    Base mixin supporting two break out methods for montiroing code execution
+    within the selected context. Can be overridden depending on hook to provide
+    enhanced functionality.
+    """
+    started = None
+    time_take = None
+
+    def enter_section_hook(self):
+        self.started = datetime.now()
+        return
+
+    def exit_section_hook(self, default_kwargs):
+        self.time_taken = naturaldelta(datetime.now() - self.started)
+
+        default_kwargs['title'] = (
+            f'Excecuted context in {self.time_taken}.'
+        )
+
+        return default_kwargs
+
+##
+# Set django as default hook if available.
+# Use a basic generic mixin if not.
+##
+
+
+try:
+    from django import db
+
+    class DjangoHookMixin(BaseHookMixin):
+        """
+        Django specific hook, enhances the end flame graph with an output of
+        total database queries.
+        """
+
+        def enter_section_hook(self):
+            super().enter_section_hook()
+            db.reset_queries()
+
+        def exit_section_hook(self, default_kwargs):
+            super().exit_section_hook(default_kwargs)
+            num_sql_queries = len(db.connection.queries)
+
+            default_kwargs['title'] = (
+                f'Executed context in {self.time_taken}. '
+                f'Made {num_sql_queries} SQL queries.'
+            )
+
+            return default_kwargs
+
+    SetHookMixin = DjangoHookMixin
+
+except ImportError:
+    SetHookMixin = BaseHookMixin

setup.py

@@ -8,17 +8,17 @@
 os.chdir(os.path.normpath(os.path.join(os.path.abspath(__file__), os.pardir)))
 
 setup(
-    name = 'publons-flame',
-    version = '0.1',
-    packages = ['publons_flame'],
-    include_package_data = True,
-    license = 'MIT License',
-    description = 'A small Django and IPython compatible application for benchmarking database and IO heavy work.',
-    long_description = README,
-    url = 'https://github.com/pulbons/publons-flame',
-    author = 'Matthew Betts, Aidan Houlihan',
-    author_email = '[email protected]',
-    classifiers =[
+    name='publons-flame',
+    version='0.1',
+    packages=['publons_flame'],
+    include_package_data=True,
+    license='MIT License',
+    description='A small Django and IPython compatible application for benchmarking database and IO heavy work.',
+    long_description=README,
+    url='https://github.com/pulbons/publons-flame',
+    author='Matthew Betts, Aidan Houlihan',
+    author_email='[email protected]',
+    classifiers=[
         'Environment :: Web Environment',
         'Framework :: Django',
         'Intended Audience :: Developers',