Distribution.py

import matplotlib.pyplot as plt
import os
import sys

COLORS = ['#1f77b4', '#ff7f0e', '#2ca02c', '#d62728', '#9467bd',
          '#8c564b', '#e377c2', '#7f7f7f', '#bcbd22', '#17becf']


def get_files_count(directory):
    count = 0

    directories = sorted(os.listdir(directory))
    for path in directories:
        if os.path.isfile(os.path.join(directory, path)):
            count += 1
    return count


def find_subs(directory):
    dir_names = []
    values = []

    directories = sorted(os.listdir(directory))
    for path in directories:
        dir_path = os.path.join(directory, path)
        if os.path.isdir(dir_path):
            try:
                values.append(get_files_count(dir_path))
                dir_names.append(path)
            except PermissionError:
                print(f"Permission denied: '{dir_path}'")
    return dir_names, values


def main():
    directory = sys.argv[1] if len(sys.argv) > 1 else None
    if directory is None:
        return print("Directory argument needed: "
                     "./Distribution.py <directory>")

    try:
        dir_names, values = find_subs(directory)

        fig, ax = plt.subplots(1, 2)
        fig.tight_layout()
        fig.suptitle(f"Files distribution inside {directory}")

        if sum(values) == 0:
            return print("No images found.")

        ax[0].pie(values, labels=dir_names, autopct='%1.1f%%')
        ax[1].bar(dir_names, values, color=COLORS)
        plt.show()
    except NotADirectoryError:
        return print(f"Not a directory: '{directory}'")
    except PermissionError:
        return print(f"Permission denied: '{directory}'")
    except FileNotFoundError:
        return print(f"Could not find directory: '{directory}'")

    return


if __name__ == '__main__':
    main()