Add import callbacks for import hook

* typo-fix

* add before and after import callbacks

2 files changed, 21 insertions, 2 deletions

diff --git a/library/python/runtime_py3/importer.pxi b/library/python/runtime_py3/importer.pxi
index 6b3c5889a5f..f3b353ef034 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 5110abc412c..9f794f03591 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