diff --git a/benchmark.py b/benchmark.py index 4e1a975..4b38dd9 100644 --- a/benchmark.py +++ b/benchmark.py @@ -123,25 +123,82 @@ def get_tree_size(path): try: for entry in scandir.scandir(path): if entry.is_dir(follow_symlinks=False): - size += get_tree_size(os.path.join(path, entry.name)) + size += get_tree_size(entry.path) else: size += entry.stat(follow_symlinks=False).st_size except OSError: pass return size +def get_tree_size_listdir(path): + """Return total size of all files in directory tree at path.""" + size = 0 + try: + for name in os.listdir(path): + fullpath = os.path.join(path, name) + if os.path.isdir(fullpath) and not os.path.islink(fullpath): + size += get_tree_size_listdir(fullpath) + else: + size += os.stat(fullpath, follow_symlinks=False).st_size + except OSError: + pass + return size + +def get_tree_size_listdir_stat(path): + """Return total size of all files in directory tree at path.""" + size = 0 + try: + for name in os.listdir(path): + fullpath = os.path.join(path, name) + st = os.stat(fullpath, follow_symlinks=False) + if stat.S_ISDIR(st.st_mode): + size += get_tree_size_listdir_stat(fullpath) + else: + size += st.st_size + except OSError: + pass + return size + +if hasattr(os, 'fwalk'): + _dir_flags = os.O_RDONLY + def get_tree_size_listdir_fd(fd): + """Return total size of all files in directory tree at path.""" + size = 0 + try: + for name in os.listdir(fd): + st = os.stat(name, follow_symlinks=False, dir_fd=fd) + if stat.S_ISDIR(st.st_mode): + top_fd = os.open(name, _dir_flags, dir_fd=fd) + try: + size += get_tree_size_listdir_fd(top_fd) + finally: + os.close(top_fd) + else: + size += st.st_size + except OSError: + pass + return size def benchmark(path, get_size=False): sizes = {} if get_size: - def do_os_walk(): - size = 0 - for root, dirs, files in os_walk_func(path): - for filename in files: - fullname = os.path.join(root, filename) - size += os.path.getsize(fullname) - sizes['os_walk'] = size + if get_size == 'fd' or (hasattr(os, 'fwalk') and get_size == 'best'): + print('Comparing against get_tree_size_listdir_fd') + def do_os_walk(): + fd = os.open(path, _dir_flags) + try: + sizes['os_walk'] = get_tree_size_listdir_fd(fd) + finally: + os.close(fd) + elif get_size == 'stat': + print('Comparing against get_tree_size_listdir_stat') + def do_os_walk(): + sizes['os_walk'] = get_tree_size_listdir_stat(path) + else: + print('Comparing against get_tree_size_listdir') + def do_os_walk(): + sizes['os_walk'] = get_tree_size_listdir(path) def do_scandir_walk(): sizes['scandir_walk'] = get_tree_size(path) @@ -195,6 +252,8 @@ using it instead of creating a tree.""" help='version of scandir() to use, default "%default"') parser.add_option('-o', '--os-walk', type='choice', choices=['match', 'builtin', 'python'], default='builtin', help='version of os.walk() to use, default "%default"') + parser.add_option('--get-size-type', type='choice', choices=['best', 'stat', 'fd', 'isdir'], default='best', + help='version of get_tree_size_listdir_*() to use, default "%default"') options, args = parser.parse_args() if args: @@ -252,4 +311,4 @@ using it instead of creating a tree.""" else: print('Comparing against ctypes emulation of os.walk()') - benchmark(tree_dir, get_size=options.size) + benchmark(tree_dir, get_size=options.size and options.get_size_type)