hook.py 886 B

1234567891011121314151617181920212223242526272829303132333435
  1. # -*- coding: utf-8 -*-
  2. from __future__ import absolute_import
  3. import sys
  4. from .parser import load_module
  5. class ThriftImporter(object):
  6. def __init__(self, extension="_thrift"):
  7. self.extension = extension
  8. def __eq__(self, other):
  9. return self.__class__.__module__ == other.__class__.__module__ and \
  10. self.__class__.__name__ == other.__class__.__name__ and \
  11. self.extension == other.extension
  12. def find_module(self, fullname, path=None):
  13. if fullname.endswith(self.extension):
  14. return self
  15. def load_module(self, fullname):
  16. return load_module(fullname)
  17. _imp = ThriftImporter()
  18. def install_import_hook():
  19. global _imp
  20. sys.meta_path[:] = [x for x in sys.meta_path if _imp != x] + [_imp]
  21. def remove_import_hook():
  22. global _imp
  23. sys.meta_path[:] = [x for x in sys.meta_path if _imp != x]