import os
import time
import tempfile

NDIR = 10**4
NFILE = 10**5 # x 2 (file + symlink)

TEST_STAT = True
RUNS = 3
LOOPS = 5

if TEST_STAT:
    def test_listdir():
        for name in os.listdir():
            os.stat(name)

    def test_scandir():
        for entry in os.scandir():
            entry.is_dir()
else:
    def test_listdir():
        for entry in os.listdir():
            pass

    def test_scandir():
        for entry in os.scandir():
            pass

def bench(times, func):
    t0 = time.perf_counter()
    for loop in range(LOOPS):
        func()
    t1 = time.perf_counter()
    dt = (t1 - t0) / LOOPS
    times.append(dt)

dt_listdir = []
dt_scandir = []
counter = 1
old_dir = os.getcwd()

if TEST_STAT:
    print("Test listdir+stat vs scandir+is_dir")
else:
    print("Test listdir vs scandir")

with tempfile.TemporaryDirectory() as tmp:
    os.chdir(tmp)
    print("Temporary directory: %s" % tmp)

    print("Create %s files+symlinks..." % NFILE)
    for file in range(NFILE):
        filename = "file-%08d" % counter
        linkname = "symlink-%08d" % counter
        counter += 1

        open(filename, "xb").close()
        os.symlink(filename, linkname)

    print("Create %s directories..." % NDIR)
    for dir in range(NDIR):
        dirname = "dir-%08d" % counter
        counter += 1
        os.mkdir(dirname)

    nentries = 0
    for entry in os.scandir():
        nentries += 1
    print("# entries: %s" % nentries)

    print("Benchmark...")
    for run in range(RUNS):
        bench(dt_listdir, test_listdir)
        print("listdir: %.1f ms" % (dt_listdir[-1]*1e3))
        bench(dt_scandir, test_scandir)
        print("scandir: %.1f ms" % (dt_scandir[-1]*1e3))
    print("")

    print("Remove the temporary directory...")
    print("")

os.chdir(old_dir)
best_listdir = min(dt_listdir)
best_scandir = min(dt_scandir)
worst_listdir = max(dt_listdir)
worst_scandir = max(dt_scandir)
print("Result:")
print("listdir: min=%.1f ms (%.1f us per file), max=%.1f ms (%.1f us per file)"
      % (best_listdir*1e3, best_listdir/nentries*1e6,
         worst_listdir*1e3, worst_listdir/nentries*1e6))
print("scandir: min=%.1f ms (%.1f us per file), max=%.1f ms (%.1f us per file)"
      % (best_scandir*1e3, best_scandir/nentries*1e6,
         worst_scandir*1e3, worst_scandir/nentries*1e6))

print("scandir is between %.1fx and %.1fx faster"
      % (best_listdir / worst_scandir, worst_listdir / best_scandir))