From 5aa22375de5516915bb1cb9168e04430933e86a6 Mon Sep 17 00:00:00 2001
From: Min RK <benjaminrk@gmail.com>
Date: Mon, 14 Mar 2016 16:00:41 +0100
Subject: [PATCH] Make it an IPython extension

---
 Demo.ipynb   | 48 +++++++++++++++++++++++++++++++++++++++++++++++-
 README.md    |  8 ++++++++
 wurlitzer.py | 19 +++++++++++++++++++
 3 files changed, 74 insertions(+), 1 deletion(-)

diff --git a/Demo.ipynb b/Demo.ipynb
index 7e84b68..1f49bb4 100644
--- a/Demo.ipynb
+++ b/Demo.ipynb
@@ -51,7 +51,7 @@
    "cell_type": "code",
    "execution_count": 2,
    "metadata": {
-    "collapsed": true
+    "collapsed": false
    },
    "outputs": [],
    "source": [
@@ -198,6 +198,52 @@
     "\n",
     "print(stdout.getvalue())"
    ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## IPython extension\n",
+    "\n",
+    "You can also enable wurlitzer as an IPython extension,\n",
+    "so that it always forwards C-level output during execution:"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 8,
+   "metadata": {
+    "collapsed": true
+   },
+   "outputs": [],
+   "source": [
+    "%load_ext wurlitzer"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 9,
+   "metadata": {
+    "collapsed": false
+   },
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Hello from C, 0!\n",
+      "Hello from C, 1!\n",
+      "Hello from C, 2!\n",
+      "Hello from C, 3!\n",
+      "Hello from C, 4!\n"
+     ]
+    }
+   ],
+   "source": [
+    "for i in range(5):\n",
+    "    time.sleep(1)\n",
+    "    printf(\"Hello from C, %i!\" % i)"
+   ]
   }
  ],
  "metadata": {
diff --git a/README.md b/README.md
index 30b303c..bb97710 100644
--- a/README.md
+++ b/README.md
@@ -42,6 +42,14 @@ with sys_pipes():
     call_some_c_function()
 ```
 
+Or even simpler, enable it as an IPython extension:
+
+```
+%load_ext wurlitzer
+```
+
+To forward all C-level output to IPython during execution.
+
 ## Acknowledgments
 
 This package is based on stuff we learned with @takluyver and @karies while working on capturing output from the [Cling Kernel](https://github.com/root-mirror/cling/tree/master/tools/Jupyter/kernel) for Jupyter.
diff --git a/wurlitzer.py b/wurlitzer.py
index d726f8a..05212dc 100644
--- a/wurlitzer.py
+++ b/wurlitzer.py
@@ -261,3 +261,22 @@ def stop_sys_pipes():
         _mighty_wurlitzer.__exit__(None, None, None)
         _mighty_wurlitzer = None
 
+
+def load_ipython_extension(ip):
+    """Register me as an IPython extension
+    
+    Captures all C output during execution and forwards to sys.
+    
+    Use: %load_ext wurlitzer
+    """
+    ip.events.register('pre_execute', sys_pipes_forever)
+    ip.events.register('post_execute', stop_sys_pipes)
+
+
+def unload_ipython_extension(ip):
+    """Unload me as an IPython extension
+    
+    Use: %unload_ext wurlitzer
+    """
+    ip.events.unregister('pre_execute', sys_pipes_forever)
+    ip.events.unregister('post_execute', stop_sys_pipes)