nbviewer_output_cleaner.py

import sys
import os


OUTPUTS_FOLDER_NAME = 'outputs'
OUTPUT_FILE_FORMATS = ['.html']
ASK_BEFORE_DELETE = True
DEBUG = False


class bcolors:
    HEADER = '\033[95m'
    OKBLUE = '\033[94m'
    OKCYAN = '\033[96m'
    OKGREEN = '\033[92m'
    WARNING = '\033[93m'
    FAIL = '\033[91m'
    ENDC = '\033[0m'
    BOLD = '\033[1m'
    UNDERLINE = '\033[4m'


def print_and_exit(message):
    print(message)
    exit()


def delete_tenant_outputs(data_folder, tenantId):
    tenant_folder = os.path.join(data_folder, tenantId)
    if not os.path.isdir(tenant_folder):
        pass

    print('Deleting tenant outputs for [{}]'.format(tenantId))

    notebooks = os.listdir(tenant_folder)
    for notebook_id in notebooks:
        output_folder = os.path.join(
            tenant_folder, notebook_id, OUTPUTS_FOLDER_NAME)
        if not os.path.isdir(output_folder):
            continue

        output_count = 0
        output_files = os.listdir(output_folder)

        for output_file in output_files:
            output_filep_path = os.path.join(output_folder, output_file)

            if os.path.isfile(output_filep_path) and str(output_file).endswith(tuple(OUTPUT_FILE_FORMATS)):
                try:
                    os.remove(output_filep_path)
                    output_count = output_count + 1
                    if DEBUG:
                        print('Deleted {}'.format(output_filep_path))
                except Exception as e:
                    print('Error while deleting output file {}'.format(e))

        print('{}{} output file deleted for tenant/notebook  [{}]/[{}]{}'.format(
            bcolors.OKGREEN, output_count, tenantId, notebook_id, bcolors.ENDC))


def check_and_delete_tenants(data_folder):
    print('{}Output delete operation starting [{}]{}'.format(
        bcolors.WARNING, data_folder, bcolors.ENDC))
    if not os.path.isdir(data_folder):
        print_and_exit('Output parameter is not folder!')

    tenant_folders = os.listdir(data_folder)

    for tenantId in tenant_folders:
        if not os.path.isdir(os.path.join(data_folder, tenantId)):
            continue
        if ASK_BEFORE_DELETE:
            print('{}{}Are you sure to delete tenant outputs for {} (Press N to skip this tenant){}'.format(
                bcolors.WARNING, bcolors.BOLD, tenantId, bcolors.ENDC))
            user_response = ''
            if sys.version_info[0] < 3:
                user_response = raw_input()
            else:
                user_response = input()
            if user_response and user_response.lower() == 'n':
                continue

        delete_tenant_outputs(data_folder, tenantId)

    print('Output delete operation completed.')


if __name__ == "__main__":
    if len(sys.argv) <= 1:
        print('Output folder parameter is missing!')
        exit()

    check_and_delete_tenants(sys.argv[1])