Added main function exit code which is the number of warnings during script execution

[imap-restore-mail] / restore_mail_inject.py

diff --git a/restore_mail_inject.py b/restore_mail_inject.py
index c67e6ed..11b9471 100644 (file)
--- a/restore_mail_inject.py
+++ b/restore_mail_inject.py
@@ -12,17 +12,19 @@ import logging
 import argparse
 from mail_iterator import MailIterator
 from file_iterator import FileIterator
+from warnings_handler import WarningsHandler
 
 LOG_FILENAME = "restore_mail_inject.log"
 LOG_FILE_LEVEL = logging.DEBUG
 LOG_SHELL_LEVEL = logging.INFO
+LOG_UNCLEAN_EXIT_LEVEL = logging.WARNING
 
 def main():
     """Main function."""
 
     # prepare configuration
     args = configure_args()
-    prepare_logger()
+    warnings_handler = prepare_logger()
     logging.info("The module restore_mail_inject.py started with user %s, folder %s and source %s." %
           (args.user, args.folder, args.srcdir))
 
@@ -61,8 +63,8 @@ def main():
     for acl_mailbox in storage.acl_mailboxes:
         session.add_acls(acl_mailbox, mailbox_list, args.ouser, args.user)
 
-    logging.info("Finished injecting mails. Exiting.")
-    return
+    logging.info("Finished injecting mails. Exiting with code %s." % warnings_handler.detected_problems)
+    return warnings_handler.detected_problems
 
 def configure_args():
     """Configure arguments and return them."""
@@ -102,12 +104,18 @@ def prepare_logger():
                         level=LOG_FILE_LEVEL)
 
     # add a handler for a console output
+    default_logger = logging.getLogger('')
     console = logging.StreamHandler()
     console.setLevel(LOG_SHELL_LEVEL)
     formatter = logging.Formatter('%(asctime)s - %(name)s - %(levelname)s - %(message)s')
     console.setFormatter(formatter)
-    logging.getLogger('').addHandler(console)
-    return
+    default_logger.addHandler(console)
+    
+    # add a handler for warnings counting
+    warnings_handler = WarningsHandler()
+    warnings_handler.setLevel(LOG_UNCLEAN_EXIT_LEVEL)
+    default_logger.addHandler(warnings_handler)
+    return warnings_handler
 
 if __name__ == "__main__":
     main()