Implement import hook to transform modules on load

Apply an arbitrary transformation to module source code at load time
before compiling it to bytecode and executing it, of the form:

  transform(source : str, name : str, filename : str) -> str | AST

The result is passed to builtin compile() function wich accepts both
str and AST objects.

my2to3/hook.py

+import imp
+import os
+import sys
+
+
+class Finder(object):
+  INIT_DOT_PY = '__init__%spy' % os.extsep
+
+  def __init__(self, prefix, transform):
+    self._prefix = prefix.split('.')
+    self._transform = transform
+
+  def find_module(self, name, path=None):
+    components = name.split('.')
+    it = iter(components)
+    if not any(c in it for c in self._prefix):
+      return None
+    subname = components[-1]
+    file, path, description = imp.find_module(subname, path)
+    if description[-1] == imp.PKG_DIRECTORY:
+      filename = os.path.join(path, self.INIT_DOT_PY)
+      file = open(filename)
+    elif file is None:
+      return None
+    else:
+      filename = path
+    return Loader(file, path, filename, self._transform)
+
+
+class Loader(object):
+  def __init__(self, file, path, filename, transform):
+    self._file = file
+    self._path = path
+    self._filename = filename
+    self._transform = transform
+
+  def compile(self, name):
+    filename = self._filename
+    source = self._file.read()
+    self._file.close()
+    return compile(self._transform(source, name, filename), filename, 'exec')
+
+  def prepare_module(self, module):
+    filename, path = self._filename, self._path
+    module.__file__ = filename
+    module.__loader__ = self
+    if len(filename) > len(path):
+      module.__path__ = [path]
+      module.__package__ = module.__name__
+    else:
+      module.__package__ = module.__name__.rpartition('.')[0]
+    return module
+
+  def load_module(self, name):
+    imp.acquire_lock()
+    try:
+      try:
+        module = sys.modules[name]
+      except KeyError:
+        module = imp.new_module(name)
+        sys.modules[name] = module
+      self.prepare_module(module)
+      code = self.compile(name)
+      try:
+        exec code in module.__dict__
+        return module
+      except:
+        del sys.modules[name]
+        raise
+    finally:
+      imp.release_lock()
+      self._file.close()

tests/context.py

+import os
+import sys
+
+path = os.path.abspath(os.path.join(os.path.dirname(__file__), os.pardir))
+sys.path.insert(0, path)

tests/sample/sample.py

+import my2to3.trace
+
+def f(): pass
+
+b = 2 / 3

tests/test_hook.py

+import unittest
+import sys
+
+import context
+from my2to3.hook import Finder
+
+
+class TestHook(unittest.TestCase):
+  def setUp(self):
+    self._calls = []
+    self._finder = None
+
+  def transform(self, source, name, filename):
+    self._calls.append((source, name, filename))
+    return source
+
+  def test_hook_sample(self):
+    self._finder = Finder('sample', self.transform)
+    sys.meta_path.append(self._finder)
+
+    import sample.sample
+
+    self.assertEqual(len(self._calls), 2)
+    for i, s in enumerate((sample, sample.sample)):
+      with open(s.__file__) as f:
+        source = f.read()
+      self.assertEqual(self._calls[i], (source, s.__name__, s.__file__))
+
+    calls = self._calls
+    self._calls = []
+
+    reload(sample)
+    reload(sample.sample)
+
+    self.assertEqual(calls, self._calls)
+
+  def tearDown(self):
+    sys.meta_path.remove(self._finder)
+    del sys.modules['sample']
+    del sys.modules['sample.sample']