Как оптимизировать алгоритм

#!/usr/bin/env python
#
# Дампер почты
# Сохраняет все письма и сортирует по папкам
# Сортировка: по отправителям
#
 
import imaplib
import getpass
import os
import email
import time
import glob
import shutil
import threading
 
from samba.common import raw_input
 
 
def imap_predictor(email_account: str):
    ''' Предсказатель imap сервера'''
    imap_server = 'imap.{}'.format(email_account.split('@')[1])
    return imap_server
 
 
def mail_count(mail: imaplib):
    ''' Вывод количества писем'''
    try:
        rv, data = mail.search(None, "ALL")
        if rv != 'OK':
            print("No messages found!")
            return 0
        count = len(data[0].split())
        return count
    except Exception as e:
        print('Error: {}'.format(e))
 
 
def get_folder(mail: imaplib):
    ''' Выбирает папку на почтовом ящике'''
    list_folders = mail.list()
    folders = list_folders[1]
    new_folders = []
    for folder in folders:
        new_folders.append(folder.decode('utf-8').split('/')[1])
    folders = []
    print('Folders: ')
    for i in range(len(new_folders)):
        new_folders[i] = new_folders[i].strip('" ')
        folders.append(new_folders[i])
        print('\t{}) {}'.format(i + 1, new_folders[i]))
    try:
        i = int(input('Choose folder: '))
        return folders[i - 1]
    except IndexError:
        print('Folder number {} not found!'.format(i))
        exit(1)
 
 
def get_all_folders(mail: imaplib):
    '''Возвращает все папки на почте'''
    list_folders = mail.list()
    folders = list_folders[1]
    new_folders = []
    for folder in folders:
        new_folders.append(folder.decode('utf-8').split('/')[1].strip('" '))
    return new_folders
 
 
def make_folder(main_directory: str,
                email_account: str,
                email_folder: str
                ):
    ''' Создаёт папку для писем '''
    try:
        account_dir = '{}/{}'.format(main_directory, email_account)
        email_folder_dir = '{}/{}'.format(account_dir, email_folder)
        if (not os.path.exists(account_dir)):
            os.mkdir(email_account)
        if (not os.path.exists(email_folder_dir)):
            os.mkdir(email_folder_dir)
        output_directory = email_folder_dir
        return output_directory
    except Exception as e:
        print('Error: {}'.format(e))
 
 
def get_first_message(data: list, mail: imaplib):
    sender: str = ''
    for num in data[0].split():
        rv, data = mail.fetch(num, '(RFC822)')
        try:
            raw_email = data[0][1].decode('utf-8')
            email_msg = email.message_from_string(raw_email)
            sender = email_msg['From']
            if (sender != ''):
                break
        except UnicodeError:
            pass
    return sender
 
 
def get_senders(mail: imaplib):
    ''' Возвращает отправителей писем'''
    start_time = time.time()
    rv, data = mail.search(None, "ALL")
    if rv != 'OK':
        print("No messages found!")
        return
    senders = []
    for num in data[0].split():
        rv, data = mail.fetch(num, '(RFC822)')
        try:
            raw_email = data[0][1].decode('utf-8')
            email_msg = email.message_from_string(raw_email)
            sender: str = email_msg['From']
            senders.append(sender)
        except UnicodeError:
            pass
    unique_senders = list(set(senders))
    print('Done! {} s.'.format(time.time() - start_time))
    return unique_senders
 
 
def select_sender(senders: list):
    '''Выбор отправителя '''
    if (senders == []):
        sender = input('Input sender: ')
        return sender
    else:
        print('Senders: ')
        for i in range(len(senders)):
            print('{}: {}'.format(i+1, senders[i]))
        i = int(input('Select sender: '))
        return senders[i-1]
 
 
def write_senders(output_directory: str,
                  senders: list
                  ):
    '''Запись отправителей в файл'''
    if (not os.path.exists('{}/senders.txt'.format(output_directory))):
        sw = int(input('Write senders to file? (0/1): '))
        if (sw == 0):
            return
        if (sw == 1):
            file = open('{}/senders.txt'.format(output_directory), 'w')
            for sender in senders:
                file.write('{}\n'.format(sender))
            file.close()
            return
    else:
        print('File senders.txt is already exists!')
        sw = int(input('Rewrite file? (0/1): '))
        if (sw == 0):
            return
        if (sw == 1):
            file = open('{}/senders.txt'.format(output_directory), 'w')
            for sender in senders:
                file.write('{}\n'.format(sender))
            file.close()
            return
 
 
def read_senders(output_directory: str):
    ''' Чтение отправителей из файла'''
    senders = []
    if (not os.path.exists('{}/senders.txt'.format(output_directory))):
        print('File senders.txt is not exists!')
        return
    else:
        file = open('{}/senders.txt'.format(output_directory), 'r')
        while True:
            sender = file.readline()
            if not sender:
                break
            senders.append(sender)
        return senders
 
 
def get_senders_local(output_directory: str):
    '''Возвращает отправителей писем (считывает из письма).
        Пишет отправителей в файл'''
    try:
        list_of_files = glob.glob('{}/*.eml'.format(output_directory))
        senders = []
        for file in list_of_files:
            with open(file, 'rb') as fp:
                msg = email.message_from_binary_file(fp)
                text_from: str = str(msg.get('From'))
                senders.append(text_from)
        if (senders != []):
            senders = list(set(senders))
            with open('{}/senders.txt'.format(output_directory), 'w') as fp:
                for sender in senders:
                    fp.write('{}\n'.format(sender))
            return senders
    except Exception as e:
        print('Error: {}'.format(e))
    except FileNotFoundError:
        pass
    except IsADirectoryError:
        pass
 
 
def filter_sender_remote(sender: str, mail: imaplib, output_directory: str):
    ''' Сортировка по отправителям (удалённая)'''
    sender = sender.lower()
    sender = sender.strip('\n')
    rv, data = mail.search(None, 'ALL')
    if rv != 'OK':
        print('No messages found!')
        return
    for num in data[0].split():
        try:
            rv, data = mail.fetch(num, '(RFC822)')
            raw_email = data[0][1].decode('utf-8')
            email_message = email.message_from_string(raw_email)
            msg_from: str = email_message['From']
            msg_from = msg_from.lower()
            if rv != 'OK':
                print('ERROR getting message ', num)
                return
            if (msg_from.find(sender) != -1):
                if (not os.path.exists('%s/%s.eml' % (output_directory, num))):
                    print("Writing message ", num)
                    file = open('%s/%s.eml' % (output_directory, num), 'wb')
                    file.write(data[0][1])
                    file.close()
                else:
                    print('File {} is already exists!'.format('%s/%s.eml' % (output_directory, num)))
        except UnicodeError:
            pass
 
 
def filter_sender_local(sender: str,
                        output_directory: str,
                        ):
    ''' Сортировка по отправителям (локальная)'''
    try:
        sender = sender.strip('\n')
        list_of_files = glob.glob('{}/*.eml'.format(output_directory))
        filtered_path = '{}/{}'.format(output_directory, sender)
        files = []
        # files = {}
        if (not os.path.exists(filtered_path)):
            os.mkdir(filtered_path)
        for file in list_of_files:
            with open(file, 'rb') as fp:
                msg = email.message_from_binary_file(fp)
                text_from: str = str(msg.get(u'From'))
                files.append({'file':fp.name, 'from':text_from})
                # files[text_from] = {'file':fp.name, 'from':text_from}
        # for file in files.items():
            # if (file[0].find(sender) != -1):
                # shutil.move(file[1]['file'], filtered_path)
        # for i in range(len(files.keys())):
        #     if (files.keys()[i].find(sender) != -1): # Алгоритм на словарях
        #         shutil.move(files.values()[i]['file'], filtered_path)
        #     if (i == len(files.keys()) - 1 and
        #             files.keys()[len(files.keys()-1)].find(sender) != -1):
        #         shutil.move(files.values()[len(files.keys()-1)]['file'], filtered_path)
        #         files.pop(len(files.keys()-1))
        #     if (i == len(files.keys()) - 1 and
        #             files.keys()[len(files.keys()-1)].find(sender) == -1):
        #         files.pop(len(files.keys()-1))
        for i in range(len(files)):
            if (files[i]['from'].find(sender) != -1): # Алгоритм на списках
                shutil.move(files[i]['file'], filtered_path)
            if (i == len(files) - 1 and files[len(files)-1]['from'].find(sender) != -1):
                shutil.move(files[len(files)-1]['file'], filtered_path)
                files.pop(len(files)-1) # удаление элемента из списка
            if (i == len(files) - 1 and files[len(files)-1]['from'].find(sender) == -1):
                files.pop(len(files)-1) # удаление элемента из списка
 
            # if (text_from.find(sender) != -1):
            #     shutil.move(fp.name, filtered_path)
    except FileNotFoundError:
        pass
    except IsADirectoryError:
        pass
    except shutil.Error:
        print('Path {} is already exists!'.format(filtered_path))
    except OSError:
        print('Error! File name too long!')
    except AttributeError as e:
        print('Attribute error: {}'.format(e))
    except Exception as e:
        print('Error: {}'.format(e))
 
 
def select_and_process(mail: imaplib, main_directory, email_account, folder):
    '''Выбор папки и дамп почты. Возвращает путь созданной папки'''
    try:
        rv, data = mail.select(folder)
        msg_count = mail_count(mail)
        print('Count of messages: {}'.format(msg_count))
        if (msg_count != 0):
            output_directory = make_folder(main_directory, email_account, folder)
        if rv == 'OK':
            print('Processing mailbox: {}'.format(folder))
            process_mailbox(mail, output_directory)
        else:
            print('ERROR: Unable to open mailbox: {}'.format(rv))
        return output_directory
    except Exception as e:
        print('Error: {}'.format(e))
 
 
def process_mailbox(mail: imaplib, output_directory):
    '''Дампит всю почту в папку'''
    try:
        rv, data = mail.search(None, "ALL")
        # mail.ru не умеет в поиск по письмам
        if rv != 'OK':
            print("No messages found!")
            return
        for num in data[0].split():
            rv, data = mail.fetch(num, '(RFC822)')
            if rv != 'OK':
                print("ERROR getting message", num)
                return
            if (not os.path.exists('%s/%s.eml' % (output_directory, num))):
                print("Writing message ", num)
                file = open('%s/%s.eml' % (output_directory, num), 'wb')
                file.write(data[0][1])
                file.close()
            else:
                print('File {} is already exists!'.format('%s/%s.eml' % (output_directory, num)))
    except Exception as e:
        print('Error: {}'.format(e))
 
 
def multiprocess_filter(senders, output_directories, dir):
    '''Многопоточный фильтр'''
    try:
        start = time.time()
        for i in range(len(senders)):
            filter_sender_local(senders[i], output_directories[dir])
        print('Time: {} s.'.format(time.time() - start))
    except Exception as e:
        print('Error: {}'.format(e))
 
 
def menu():
    '''Меню'''
    print('Email dumper\n'
          '\t1 - dump all folders\n'
          '\t2 - dump folder\n'
          '\t3 - exit')
    sw = int(input('Choose: '))
    return sw
 
 
def main():
    '''Главная функция.
        Время исполнения слишком большое. Нужно оптимизировать сортировку
        2137 писем сортируется за 2224 с. на старом алгоритме'''
    try:
        main_directory = os.path.abspath(os.getcwd()) # директория с программой
        email_account = raw_input('email: ') # почта
        password = getpass.getpass() # пароль
        imap_server = imap_predictor(email_account)
        mail = imaplib.IMAP4_SSL(imap_server)
        mail.login(email_account, password)
        sw = menu() # меню с выбором дампа всех папок, или конкретной
 
        if (sw == 1):
            folders = get_all_folders(mail) # получить все папки
            output_directories = []
            for folder in folders:
                # Выкачивание писем не может в мультипоток (ошибка с запросами)
                output_directory = select_and_process(mail, main_directory, email_account, folder)
                output_directories.append(output_directory)
 
            # for dir in output_directories:
                # senders = get_senders_local(dir)
                # for sender in senders:
                    # filter_sender_local(sender, dir)
 
            for dir in range(len(output_directories)):
                senders1 = get_senders_local(output_directories[dir])
                senders2 = get_senders_local(output_directories[dir+1])
                senders3 = get_senders_local(output_directories[dir+2])
                senders4 = get_senders_local(output_directories[dir+3])
 
                # # # multiprocess_filter(senders, output_directories, dir) # сигнатура для подсказки
                t1 = threading.Thread(target=multiprocess_filter, args=(senders1, output_directories, dir))
                t1.start()
                t2 = threading.Thread(target=multiprocess_filter, args=(senders2, output_directories, dir + 1))
                t2.start()
                t3 = threading.Thread(target=multiprocess_filter, args=(senders3, output_directories, dir + 2))
                t3.start()
                t4 = threading.Thread(target=multiprocess_filter, args=(senders4, output_directories, dir + 3))
                t4.start()
                t1.join()
                t2.join()
                t3.join()
                t4.join()
 
        if (sw == 2):
            folder = get_folder(mail) # выбрать папку
            output_directory = select_and_process(mail, main_directory, email_account, folder)
            senders = get_senders_local(output_directory)
            for sender in senders:
                filter_sender_local(sender, output_directory)
 
        if (sw == 3):
            mail.close()
            mail.logout()
            exit(0)
        # rv, data = mail.select(email_folder)
        # if rv == 'OK':
        #     print("Processing mailbox: ", email_folder)
        #     process_mailbox(mail, output_directory)
        # else:
        #     print("ERROR: Unable to open mailbox ", rv)
 
        mail.close()
        mail.logout()
    except IndexError:
        print('index error')
    except Exception as e:
        print('Error: {}'.format(e))
 
 
if __name__ == "__main__":
    main()