Add import callbacks for import hook

* typo-fix

* add before and after import callbacks

3 files changed, 36 insertions, 2 deletions

diff --git a/library/python/runtime/importer.pxi b/library/python/runtime/importer.pxi
index e27eaa0549..ba6c422bde 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]
diff --git a/library/python/runtime_py3/importer.pxi b/library/python/runtime_py3/importer.pxi
index 6b3c5889a5..f3b353ef03 100644
--- a/library/python/runtime_py3/importer.pxi
+++ b/library/python/runtime_py3/importer.pxi
@@ -185,6 +185,10 @@ class ResourceImporter(object):
         self.source_map = {}                  # Map from file names to module names.
         self._source_name = {}                # Map from original to altered module names.
         self._package_prefix = ''
+
+        self._before_import_callback = None
+        self._after_import_callback = None
+
         if Y_PYTHON_SOURCE_ROOT and Y_PYTHON_EXTENDED_SOURCE_SEARCH:
             self.arcadia_source_finder = ArcadiaSourceFinder(_s(Y_PYTHON_SOURCE_ROOT))
         else:
@@ -196,6 +200,11 @@ class ResourceImporter(object):
                 if k not in self.memory:
                     self.memory.add(k)
 
+    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
+
     def for_package(self, name):
         import copy
         importer = copy.copy(self)
@@ -230,7 +239,17 @@ class ResourceImporter(object):
         if self.is_package(module.__name__):
             module.__path__= [executable + path_sep + module.__name__.replace('.', path_sep)]
         # exec(code, module.__dict__)
-        _call_with_frames_removed(exec, code, module.__dict__)
+
+        if self._before_import_callback:
+            self._before_import_callback(module)
+
+        # “Zero-cost” exceptions are implemented.
+        # The cost of try statements is almost eliminated when no exception is raised
+        try:
+            _call_with_frames_removed(exec, code, module.__dict__)
+        finally:
+            if self._after_import_callback:
+                self._after_import_callback(module)
 
     # PEP-302 extension 1 of 3: data loader.
     def get_data(self, path):
diff --git a/library/python/runtime_py3/test/test_arcadia_source_finder.py b/library/python/runtime_py3/test/test_arcadia_source_finder.py
index 5110abc412..9f794f0359 100644
--- a/library/python/runtime_py3/test/test_arcadia_source_finder.py
+++ b/library/python/runtime_py3/test/test_arcadia_source_finder.py
@@ -383,7 +383,7 @@ class TestDictionaryChangedSizeDuringIteration(ArcadiaSourceFinderTestCase):
             b"py/namespace/unique_prefix1/project/lib2": b"project.lib2.",
         }
 
-    def test_no_crash_on_recusive_iter_modules(self):
+    def test_no_crash_on_recursive_iter_modules(self):
         for package in self.arcadia_source_finder.iter_modules("project.", ""):
             for _ in self.arcadia_source_finder.iter_modules(package[0], ""):
                 pass