Add import callbacks for import hook

* typo-fix

* add before and after import callbacks

1 files changed, 15 insertions, 0 deletions

diff --git a/library/python/runtime/importer.pxi b/library/python/runtime/importer.pxi
index e27eaa0549e..ba6c422bdea 100644
--- a/library/python/runtime/importer.pxi
+++ b/library/python/runtime/importer.pxi
@@ -108,6 +108,9 @@ class ResourceImporter(object):
         self._source_name = {}                # Map from original to altered module names.
         self._package_prefix = ''
 
+        self._before_import_callback = None
+        self._after_import_callback = None
+
         for p in list(self.memory) + list(sys.builtin_module_names):
             for pp in iter_prefixes(p):
                 k = pp + '.__init__'
@@ -119,6 +122,11 @@ class ResourceImporter(object):
         importer._package_prefix = name + '.'
         return importer
 
+    def set_callbacks(self, before_import=None, after_import=None):
+        """Callable[[module], None]"""
+        self._before_import_callback = before_import
+        self._after_import_callback = after_import
+
     # PEP-302 finder.
     def find_module(self, fullname, path=None):
         try:
@@ -233,12 +241,19 @@ class ResourceImporter(object):
         old_mod = sys.modules.get(mod_name, None)
         sys.modules[mod_name] = mod
 
+        if self._before_import_callback:
+            self._before_import_callback(mod)
+
         try:
             exec code in mod.__dict__
             old_mod = sys.modules[mod_name]
         finally:
             sys.modules[mod_name] = old_mod
 
+            # "Zero-cost". Just in case import error occurs
+            if self._after_import_callback:
+                self._after_import_callback(mod)
+
         # Some hacky modules (e.g. pygments.lexers) replace themselves in
         # `sys.modules` with proxies.
         return sys.modules[mod_name]