From 98c9b959f7efefa5d62bd2a69e22134eabeea292 Mon Sep 17 00:00:00 2001
From: Jason Paryani <github@jparyani.com>
Date: Sun, 1 Sep 2013 22:24:46 -0700
Subject: [PATCH] Fix up bug with import hook on python3/pypy

---
 capnp/capnp.pyx   | 11 +++++++++--
 test/test_load.py |  9 +--------
 2 files changed, 10 insertions(+), 10 deletions(-)

diff --git a/capnp/capnp.pyx b/capnp/capnp.pyx
index be44446..f1c5a71 100644
--- a/capnp/capnp.pyx
+++ b/capnp/capnp.pyx
@@ -850,6 +850,10 @@ cdef class SchemaParser:
         fileSchema = parser._parse_disk_file(display_name, file_name, imports)
         _load(fileSchema, module)
 
+        abs_path = _os.path.abspath(file_name)
+        module.__path__ = _os.path.dirname(abs_path)
+        module.__file__ = abs_path
+
         return module
 
 cdef class _MessageBuilder:
@@ -1143,7 +1147,10 @@ class _Loader:
             self.fullname, fullname))
 
         imports = self.additional_paths + _sys.path
-        return load(self.path, fullname, imports=imports)
+        module = load(self.path, fullname, imports=imports)
+        _sys.modules[fullname] = module
+
+        return module
 
 class _Importer:
     def __init__(self, additional_paths):
@@ -1198,7 +1205,7 @@ def add_import_hook(additional_paths=[]):
     """
     global _importer
     if _importer is not None:
-        return
+        remove_import_hook()
 
     _importer = _Importer(additional_paths)
     _sys.meta_path.append(_importer)
diff --git a/test/test_load.py b/test/test_load.py
index bba7d26..e596314 100644
--- a/test/test_load.py
+++ b/test/test_load.py
@@ -56,13 +56,6 @@ def test_failed_import():
 
 def test_add_import_hook():
     capnp.add_import_hook([this_dir])
-    
+
     import addressbook
     addressbook.AddressBook.new_message()
-
-def test_remove_import_hook():
-    capnp.add_import_hook([this_dir])
-    capnp.remove_import_hook()
-
-    with pytest.raises(ImportError):
-        import addressbook