From 891f0d28e5f0cc2b7e306cb3c304ff388849ba2a Mon Sep 17 00:00:00 2001
From: Jason Paryani <github@jparyani.com>
Date: Sun, 1 Sep 2013 23:55:29 -0700
Subject: [PATCH] Change import hook to require modules to end in '_capnp'

---
 capnp/capnp.pyx   |  5 +++++
 test/test_load.py | 23 +++++++++++++++++++++--
 2 files changed, 26 insertions(+), 2 deletions(-)

diff --git a/capnp/capnp.pyx b/capnp/capnp.pyx
index f1c5a71..96d32ce 100644
--- a/capnp/capnp.pyx
+++ b/capnp/capnp.pyx
@@ -1165,6 +1165,11 @@ class _Importer:
             module_name = mod_parts[-1]
         else:
             module_name = fullname
+
+        if not module_name.endswith('_capnp'):
+            return None
+
+        module_name = module_name.rstrip('_capnp')
         capnp_module_name = module_name + self.extension
 
         if package_path:
diff --git a/test/test_load.py b/test/test_load.py
index e596314..23bbd07 100644
--- a/test/test_load.py
+++ b/test/test_load.py
@@ -1,6 +1,7 @@
 import pytest
 import capnp
 import os
+import sys
 
 this_dir = os.path.dirname(__file__)
 
@@ -57,5 +58,23 @@ def test_failed_import():
 def test_add_import_hook():
     capnp.add_import_hook([this_dir])
 
-    import addressbook
-    addressbook.AddressBook.new_message()
+    import addressbook_capnp
+    addressbook_capnp.AddressBook.new_message()
+
+def test_multiple_add_import_hook():
+    capnp.add_import_hook()
+    capnp.add_import_hook()
+    capnp.add_import_hook([this_dir])
+
+    import addressbook_capnp
+    addressbook_capnp.AddressBook.new_message()
+
+def test_remove_import_hook():
+    capnp.add_import_hook([this_dir])
+    capnp.remove_import_hook()
+
+    if 'addressbook_capnp' in sys.modules:
+        del sys.modules['addressbook_capnp'] # hack to deal with it being imported already
+
+    with pytest.raises(ImportError):
+        import addressbook_capnp