Concurrent directory traversal

walk.py

import os
import argparse
import hashlib
import threading
from collections import defaultdict
from queue import Queue


rem = Queue(maxsize=0)
sem = threading.Semaphore(value=100)


_file_info = defaultdict(set)
_file_info_lock = threading.Lock()
_seen = set()
_seen_lock = threading.Lock()


def worker(path):
    while True:
        path = os.path.join(path, rem.get())
        with _seen_lock:
            if path in _seen:
                continue
            else:
                _seen.add(path)
        sem.acquire()
        try:
            if os.path.isdir(path):
                for _path in os.listdir(path):
                    if _path.startswith('.'):
                        continue
                    rem.put_nowait(os.path.join(path, _path))
            else:
                hasher = hashlib.md5()
                with open(path, 'rb') as f:
                    hasher.update(f.read())
                with _file_info_lock:
                    _file_info[hasher.hexdigest()].add(
                        (path, os.stat(path).st_size)
                    )
        finally:
            sem.release()
            rem.task_done()


if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument('-d')
    parser.add_argument('-n', type=int)

    args = parser.parse_args()

    for _ in range(args.n):
        threading.Thread(target=worker, args=(args.d,)).start()
    rem.put_nowait(os.path.join(os.getcwd(), args.d))
    rem.join()

    # print(_file_info)
    os._exit(0)