1##############################################################################
2#
3# Copyright (c) 2001-2012 Zope Foundation and Contributors.
4# All Rights Reserved.
5#
6# This software is subject to the provisions of the Zope Public License,
7# Version 2.1 (ZPL).  A copy of the ZPL should accompany this distribution.
8# THIS SOFTWARE IS PROVIDED "AS IS" AND ANY AND ALL EXPRESS OR IMPLIED
9# WARRANTIES ARE DISCLAIMED, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED
10# WARRANTIES OF TITLE, MERCHANTABILITY, AGAINST INFRINGEMENT, AND FITNESS
11# FOR A PARTICULAR PURPOSE
12#
13##############################################################################
14import os
15import sys
16
17PYPY = hasattr(sys, 'pypy_version_info')
18
19
20if sys.version_info[0] < 3: # pragma: no cover Python2
21
22    PY2 = True
23    PY3 = False
24
25    int_types = int, long
26    xrange = xrange
27    def compare(x, y):
28        if x is None:
29            if y is None:
30                return 0
31            else:
32                return -1
33        elif y is None:
34            return 1
35        else:
36            return cmp(x, y)
37
38    _bytes = str
39    def _ascii(x):
40        return bytes(x)
41
42else: # pragma: no cover Python3
43
44    PY2 = False
45    PY3 = True
46
47    int_types = int,
48    xrange = range
49
50    def compare(x, y):
51        if x is None:
52            if y is None:
53                return 0
54            else:
55                return -1
56        elif y is None:
57            return 1
58        else:
59            return (x > y) - (y > x)
60
61    _bytes = bytes
62    def _ascii(x):
63        return bytes(x, 'ascii')
64
65try:
66    from collections import abc
67except ImportError:
68    import collections as abc
69
70collections_abc = abc
71del abc
72
73def _c_optimizations_required():
74    """
75    Return a true value if the C optimizations are required.
76
77    This uses the ``PURE_PYTHON`` variable as documented in `import_c_extension`.
78    """
79    pure_env = os.environ.get('PURE_PYTHON')
80    require_c = pure_env == "0"
81    return require_c
82
83
84def _c_optimizations_available(module_name):
85    """
86    Return the C optimization module, if available, otherwise
87    a false value.
88
89    If the optimizations are required but not available, this
90    raises the ImportError.
91
92    This does not say whether they should be used or not.
93    """
94    import importlib
95    catch = () if _c_optimizations_required() else (ImportError,)
96    try:
97        return importlib.import_module('BTrees._' + module_name)
98    except catch: # pragma: no cover
99        return False
100
101
102def _c_optimizations_ignored():
103    """
104    The opposite of `_c_optimizations_required`.
105    """
106    pure_env = os.environ.get('PURE_PYTHON')
107    return pure_env != "0" if pure_env is not None else PYPY
108
109
110def _should_attempt_c_optimizations():
111    """
112    Return a true value if we should attempt to use the C optimizations.
113
114    This takes into account whether we're on PyPy and the value of the
115    ``PURE_PYTHON`` environment variable, as defined in `import_c_extension`.
116    """
117    if PYPY:
118        return False
119
120    if _c_optimizations_required():
121        return True
122    return not _c_optimizations_ignored()
123
124
125def import_c_extension(mod_globals):
126    """
127    Call this function with the globals of a module that implements
128    Python versions of a BTree family to find the C optimizations.
129
130    If the ``PURE_PYTHON`` environment variable is set to any value
131    other than ``"0"``, or we're on PyPy, ignore the C implementation.
132    If the C implementation cannot be imported, return the Python
133    version. If ``PURE_PYTHON`` is set to ``"0"``, *require* the C
134    implementation (let the ImportError propagate); the exception again
135    is PyPy, where we never use the C extension (although it builds here, the
136    ``persistent`` library doesn't provide native extensions for PyPy).
137
138    """
139    c_module = None
140    module_name = mod_globals['__name__']
141    assert module_name.startswith('BTrees.')
142    module_name = module_name.split('.')[1]
143    if _should_attempt_c_optimizations():
144        c_module = _c_optimizations_available(module_name)
145
146    if c_module:
147        new_values = dict(c_module.__dict__)
148        new_values.pop("__name__", None)
149        new_values.pop('__file__', None)
150        new_values.pop('__doc__', None)
151        mod_globals.update(new_values)
152    else:
153        # No C extension, make the Py versions available without that
154        # extension. The list comprehension both filters and prevents
155        # concurrent modification errors.
156        for py in [k for k in mod_globals if k.endswith('Py')]:
157            mod_globals[py[:-2]] = mod_globals[py]
158
159    # Assign the global aliases
160    prefix = module_name[:2]
161    for name in ('Bucket', 'Set', 'BTree', 'TreeSet'):
162        mod_globals[name] = mod_globals[prefix + name]
163
164    # Cleanup
165    mod_globals.pop('import_c_extension', None)
166