1import sys
2import types
3import toolz
4from importlib import import_module
5
6
7class TlzLoader(object):
8    """ Finds and loads ``tlz`` modules when added to sys.meta_path"""
9    def __init__(self):
10        self.always_from_toolz = {
11            toolz.pipe,
12        }
13
14    def _load_toolz(self, fullname):
15        rv = {}
16        package, dot, submodules = fullname.partition('.')
17        try:
18            module_name = ''.join(['cytoolz', dot, submodules])
19            rv['cytoolz'] = import_module(module_name)
20        except ImportError:
21            pass
22        try:
23            module_name = ''.join(['toolz', dot, submodules])
24            rv['toolz'] = import_module(module_name)
25        except ImportError:
26            pass
27        if not rv:
28            raise ImportError(fullname)
29        return rv
30
31    def find_module(self, fullname, path=None):  # pragma: py3 no cover
32        package, dot, submodules = fullname.partition('.')
33        if package == 'tlz':
34            return self
35
36    def load_module(self, fullname):  # pragma: py3 no cover
37        if fullname in sys.modules:  # pragma: no cover
38            return sys.modules[fullname]
39        spec = TlzSpec(fullname, self)
40        module = self.create_module(spec)
41        sys.modules[fullname] = module
42        self.exec_module(module)
43        return module
44
45    def find_spec(self, fullname, path, target=None):  # pragma: no cover
46        package, dot, submodules = fullname.partition('.')
47        if package == 'tlz':
48            return TlzSpec(fullname, self)
49
50    def create_module(self, spec):
51        return types.ModuleType(spec.name)
52
53    def exec_module(self, module):
54        toolz_mods = self._load_toolz(module.__name__)
55        fast_mod = toolz_mods.get('cytoolz') or toolz_mods['toolz']
56        slow_mod = toolz_mods.get('toolz') or toolz_mods['cytoolz']
57        module.__dict__.update(toolz.merge(fast_mod.__dict__, module.__dict__))
58        package = fast_mod.__package__
59        if package is not None:
60            package, dot, submodules = package.partition('.')
61            module.__package__ = ''.join(['tlz', dot, submodules])
62        if not module.__doc__:
63            module.__doc__ = fast_mod.__doc__
64
65        # show file from toolz during introspection
66        try:
67            module.__file__ = slow_mod.__file__
68        except AttributeError:
69            pass
70
71        for k, v in fast_mod.__dict__.items():
72            tv = slow_mod.__dict__.get(k)
73            try:
74                hash(tv)
75            except TypeError:
76                tv = None
77            if tv in self.always_from_toolz:
78                module.__dict__[k] = tv
79            elif (
80                isinstance(v, types.ModuleType)
81                and v.__package__ == fast_mod.__name__
82            ):
83                package, dot, submodules = v.__name__.partition('.')
84                module_name = ''.join(['tlz', dot, submodules])
85                submodule = import_module(module_name)
86                module.__dict__[k] = submodule
87
88
89class TlzSpec(object):
90    def __init__(self, name, loader):
91        self.name = name
92        self.loader = loader
93        self.origin = None
94        self.submodule_search_locations = []
95        self.loader_state = None
96        self.cached = None
97        self.parent = None
98        self.has_location = False
99
100
101tlz_loader = TlzLoader()
102sys.meta_path.append(tlz_loader)
103tlz_loader.exec_module(sys.modules['tlz'])
104