import gc
import os
import random
import sys
import threading
import time
import weakref

# Create a module that will augment its own dictionary each time it is
# initialized. Coupled with releasing the GIL often enough, this
# allows testing that the import lock prevents initializing the same module
# from several threads at once.
module_name = "test_mt%d" % os.getpid()
module_path = module_name + ".py"

module_code = """if 1:
    import sys
    import threading
    import time

    def register_thread():
        dct = sys.modules[__name__].__dict__
        threads = dct.setdefault('threads', [])
        threads.append(threading.current_thread())
    
    register_thread()
    """

def create_module():
    with open(module_path, "w") as f:
        f.write(module_code)

def delete_module():
    os.unlink(module_path)
    for c in "co":
        try:
            os.unlink(module_path + c)
        except OSError:
            pass

def do_import(modules):
    modules.append(__import__(module_name))


class GILReleaser:
    """Release the GIL very frequently, even in the middle of pure C code.
    This relies on periodic interruption of memory allocations by the
    cyclic garbage collector."""

    class Cyclic:
        """An object with a cyclic reference to itself"""
        def __init__(self):
            self.x = self

    def __init__(self):
        self.refs = set()

    def __enter__(self):
        self.old_thres = gc.get_threshold()

        # We bootstrap the GIL releasing mechanism by creating an object with
        # a cyclic reference, and a weakref callback that will be executed
        # when the object is collected. This callback in turn creates another
        # similar object, and releases the GIL. To maximize impact, we set a
        # very small cyclic collection threshold as well.

        def cb(ref=None):
            if ref is not None:
                self.refs.remove(ref)
            c = self.Cyclic()
            self.refs.add(weakref.ref(c, cb))
            # Set a very small cyclic collection threshold
            gc.set_threshold(random.randrange(2, 10), 1, 1)
            # Release the GIL
            time.sleep(0.0001)

        cb()

    def __exit__(self, *e):
        gc.set_threshold(*self.old_thres)
        gc.disable()
        try:
            for wr in self.refs:
                obj = wr()
                if obj is not None:
                    obj.x = None
            self.refs.clear()
        finally:
            gc.enable()


def test_imports():
    create_module()
    sys.path.insert(0, ".")
    try:
        threads = []
        modules = []
        N = 20
        for i in range(N):
            threads.append(threading.Thread(target=do_import, args=(modules,)))
        with GILReleaser():
            for t in threads:
                t.start()
            for t in threads:
                t.join()
        # The module object was only created once
        assert len(set(modules)) == 1, len(set(modules))
        # The module object was only initialized once
        mod_threads = modules[0].threads
        assert len(mod_threads) == 1, len(mod_threads)
        # Sanity check: it was initialized from one of the subthreads
        assert mod_threads[0] in threads
    finally:
        del sys.path[0]
        del sys.modules[module_name]
        delete_module()

if __name__ == "__main__":
    test_imports()