File count

file_count.py

# Usage:
# $python file_count.py [option]
# options:
# -r      count recursively
# -s[dir] workdir
# eg.
# #work on curent workdir:
# $python file_count.py
# #count recursively
# $python file_count.py -r
# #work on other dir
# $python file_count.py -r -s/home/user/images

import os
import sys

import re

work_dir = './'
recursion = False
args = sys.argv
for arg in args[1:]:
    if re.search('^-s', arg):
        work_dir = arg[2:]
    if re.search('^-r', arg):
        recursion = True


def is_img_file(file):
    file = str(file)
    file = file.lower()
    return (not file.startswith('.')) and (re.search('\.bmp$', file) or re.search('\.png$', file) or
                                           re.search('\.jpg$', file) or re.search('\.jpeg$', file) or re.search('\.tiff$', file))


def count_on_dir(dir, recursion=False, is_file_type=lambda f: True):
    if (not os.path.exists(dir)) or not os.path.isdir(dir):
        return 0
    count = 0
    for file in os.listdir(dir):
        if is_file_type(file):
            count += 1
        elif recursion and os.path.isdir(os.path.join(dir, file)):
            count += count_on_dir(os.path.join(dir, file), recursion, is_file_type)
    return count

if __name__ == '__main__':
    img_file_num = count_on_dir(work_dir, recursion)
    print('Image files number:{}'.format(img_file_num))