# Re-adapted from https://github.com/benhoyt/scandir/blob/master/benchmark.py

import os
import timeit
import shutil


DEPTH = 4
NUM_DIRS = 5
NUM_FILES = 50
SRC_NAME = "copytree"
DST_NAME = "copytree-copy"
_total = 0


def create_tree(path, depth=DEPTH):
    global _total
    _total += 1
    os.mkdir(path)
    for i in range(NUM_FILES):
        filename = os.path.join(path, 'file{0:03}.txt'.format(i))
        with open(filename, 'wb') as f:
            f.write(b'foo')
        _total += 1
    if depth <= 1:
        return
    for i in range(NUM_DIRS):
        dirname = os.path.join(path, 'dir{0:03}'.format(i))
        create_tree(dirname, depth - 1)


def benchmark(path):
    def doit():
        shutil.copytree(path, DST_NAME, symlinks=True)

    # Run this once first to cache things, so we're not benchmarking I/O
    shutil.rmtree(DST_NAME, ignore_errors=True)
    print("Priming the system's cache...")
    doit()
    # Use the best of 3 time for each of them to eliminate high outliers
    copytree_time = 1000000
    N = 3
    for i in range(N):
        print('{0} files and dirs, repeat {1}/{2}... '.format(
            _total, i + 1, N), end='')
        shutil.rmtree(DST_NAME, ignore_errors=True)
        copytree_time = min(copytree_time, timeit.timeit(doit, number=1))
        print('min = {0:.3f}s'.format(copytree_time))
    print('best result = {0:.3f}s'.format(copytree_time))


def main():
    shutil.rmtree(DST_NAME, ignore_errors=True)
    create_tree(SRC_NAME)
    benchmark(SRC_NAME)


if __name__ == '__main__':
    try:
        main()
    finally:
        shutil.rmtree(SRC_NAME, ignore_errors=True)
        shutil.rmtree(DST_NAME, ignore_errors=True)