createbag.py

"""
GUI tool to create a Bag from a filesystem folder.
"""

import ConfigParser
import sys
import os
import shutil
import getpass
import bagit
import platform
if platform.system() != 'Darwin' and platform.system() != 'Windows':
    # We don't use Gtk on OSX.
    from gi.repository import Gtk
elif platform.system() == 'Windows':
    from PyQt4 import QtGui, QtCore
elif platform.system() == 'Darwin':
    # Sets up Cocoadialog for error message popup on OSX.
    class cocoaPopup:

        # Change CD_BASE to reflect the location of Cocoadialog on your system.
        CD_BASE = "~/.createbag/"
        CD_PATH = os.path.join(CD_BASE, "CocoaDialog.app/Contents/MacOS/CocoaDialog")

        def __init__(self, title, message, button):
            template = "%s msgbox --title '%s' --text '%s' --button1 '%s'"
            self.pipe = os.popen(template % (cocoaPopup.CD_PATH, title, message, button), "w")
    
        def cocoaError():    
            if __name__ == "__main__":
                popup = cocoaPopup("Error","Sorry, you can't create a bag here -- you may want to change the config file so that bags are always created in a different output directory, rather than in situ.","OK")
                if popup == "1":
                    popup.close()
                    sys.exit()
                    
        def cocoaSuccess(bag_dir):    
            if __name__ == "__main__":
                popup = cocoaPopup("Success!","Bag created at %s" % bag_dir,"OK")
                if popup == "1":
                    popup.close()

# Under Linux and Windows, config file location is first command-line parameter.
# Under OSX, the input directory is the first, the config file is the second.
if platform.system() != 'Darwin':
    if len(sys.argv) > 1:
        config_file = sys.argv[1]
    else:
        config_file = './config.cfg'
else:
    if len(sys.argv) > 2:
        config_file = sys.argv[2]
    else:
        config_file = './config.cfg'

config = ConfigParser.ConfigParser()
config.optionxform = str
config.read(config_file)

# Get custom tags from config file.
bagit_tags = {}
if config.has_option('CustomTags', 'customtags'):
    tags = config.options('CustomTags')
    for tag in tags:
        bagit_tags[tag] = config.get('CustomTags', tag)

# Get shortcuts from config file.
if config.has_option('Shortcuts', 'shortcuts'):
    filechooser_shortcuts = [shortcut.strip() for shortcut in config.get('Shortcuts', 'shortcuts').split(',')]
else:
    filechooser_shortcuts = []

# Get checksum algorithms from config file.
bagit_checksum_algorithms = []
if config.has_option('Checksums', 'algorithms'):
    checksums_string = config.get('Checksums', 'algorithms', 'md5')
    bagit_checksum_algorithms = [algo.strip() for algo in checksums_string.split(',')]
else:
    bagit_checksum_algorithms = ['md5']

def directory_check(chosen_folder):
    """Prevent the utility from creating a Bag in its own directory."""    
    if config.has_option('Output', 'create_bag_in'):
        relativized_picker_path = os.path.relpath(chosen_folder, '/')
        bag_dir = os.path.join(config.get('Output', 'create_bag_in'), relativized_picker_path)
        if os.path.dirname(os.path.realpath(__file__)) == bag_dir:
            if platform.system() == 'Darwin':
                cocoaError()
            elif platform.system() == 'Windows':
                QtChooserWindow.qt_error(ex)
                return
            else:
                FolderChooserWindow.GtkError(win)
    else:
        if os.path.dirname(os.path.realpath(__file__)) == chosen_folder:
            if platform.system() == 'Darwin':
                cocoaError()
            elif platform.system() == 'Windows':
                QtChooserWindow.qt_error(ex)
                return
            else:
                FolderChooserWindow.GtkError(win)

def make_bag(chosen_folder):
    """Create a Bag from the files in chosen_folder, and return the directory
       the Bag was created in.
    """                    
    if config.has_option('AutogeneratedTags', 'add_source_directory_tag'):
        bagit_tags['Source-Directory'] = chosen_folder

    if config.has_option('AutogeneratedTags', 'add_source_user_id_tag'):
        bagit_tags['Source-User'] = getpass.getuser()                

    # If the 'create_bag_in' config option is set, create the Bag from a
    # copy of the selected folder.
    if config.has_option('Output', 'create_bag_in'):
        bag_source_dir = os.path.basename(chosen_folder)
        bag_dir = os.path.join(config.get('Output', 'create_bag_in'), bag_source_dir)

        # In this case, we want to stop if the output folder exists
        if os.path.exists(bag_dir):
            if platform.system() == 'Darwin':
                cocoaError()
            elif platform.system() == 'Windows':
                QtChooserWindow.qt_error(ex)
                return
            else:
                FolderChooserWindow.GtkError(win)

        try:
            shutil.rmtree(bag_dir, True)
            shutil.copytree(chosen_folder, bag_dir)
        except (IOError, os.error) as shutilerror:
            if platform.system() == 'Darwin':
                cocoaError()
            elif platform.system() == 'Windows':
                QtChooserWindow.qt_error(ex)
                return
            else:
                FolderChooserWindow.GtkError(win)

    # If 'create_bag_in' is not set, create the Bag in the selected directory.
    else:
        bag_dir = chosen_folder

    # Create the Bag.
    try:
        bag = bagit.make_bag(bag_dir, bagit_tags, 1, bagit_checksum_algorithms)
    except (bagit.BagError, Exception) as e:
            if platform.system() == 'Darwin':
                cocoaError()
            elif platform.system() == 'Windows':
                QtChooserWindow.qt_error(ex)
                return
            else:
                FolderChooserWindow.GtkError(win)

    return bag_dir

# Linux/Gtk-specific code (will work on Windows but not easily)
if platform.system() != 'Darwin' and platform.system() != 'Windows':
    class FolderChooserWindow(Gtk.Window):

        def __init__(self):
            Gtk.Window.__init__(self, title = config.get('UILabels', 'main_window_title', 'Create a Bag'))
            self.set_border_width(10)
            self.move(200, 200)

            box = Gtk.Box(spacing=6)
            self.add(box)
            self.spinner = Gtk.Spinner()
       
            choose_folder_button = Gtk.Button("Choose a folder to create Bag from")
            choose_folder_button.connect("clicked", self.on_folder_clicked)
            box.add(choose_folder_button)

            quit_button = Gtk.Button("Quit")
            quit_button.connect("clicked", Gtk.main_quit)
            box.add(quit_button)

        def GtkError(self):
            not_allowed_message = "\n\nYou are not allowed to run the program on that directory."
            error_dialog = Gtk.MessageDialog(self, 0, Gtk.MessageType.ERROR,
                Gtk.ButtonsType.OK, "Sorry...")
            error_dialog.format_secondary_text(not_allowed_message)
            error_dialog.run()
            error_dialog.destroy()
            raise SystemExit

        def on_folder_clicked(self, widget):
            folder_picker_dialog = Gtk.FileChooserDialog(
                config.get('UILabels', 'file_chooser_window_title', 'Create a Bag - Choose a folder to create Bag from'),
                self, Gtk.FileChooserAction.SELECT_FOLDER)
            folder_picker_dialog.set_default_size(800, 400)
            folder_picker_dialog.set_create_folders(False)
            folder_picker_dialog.add_button("Create Bag", -5)
            folder_picker_dialog.add_button("Cancel", -6)
            for filechooser_shortcut in filechooser_shortcuts:
                folder_picker_dialog.add_shortcut_folder(filechooser_shortcut)

            response = folder_picker_dialog.run()

            if response == -5:
                directory_check(folder_picker_dialog.get_filename())
                bag_dir = make_bag(folder_picker_dialog.get_filename())
                folder_picker_dialog.destroy()

                if (bag_dir):
                    confirmation_dialog = Gtk.MessageDialog(self, 0, Gtk.MessageType.INFO,
                        Gtk.ButtonsType.OK, "Bag created")
                    confirmation_dialog.format_secondary_text(
                        "The Bag for folder %s has been created." % bag_dir)
                    confirmation_dialog.run()
                    confirmation_dialog.destroy()
            if response == -6:
                folder_picker_dialog.destroy()

    win = FolderChooserWindow()
    win.connect("delete-event", Gtk.main_quit)
    win.show_all()
    Gtk.main()

    
# Windows/Qt-specific code (can also work on Linux but Gtk is nicer)
elif platform.system() == 'Windows':

    class QtChooserWindow(QtGui.QDialog):
    
        def __init__(self, parent=None):
            super(QtChooserWindow, self).__init__(parent)        
            if parent is None:
                self.initUI()
        
        def initUI(self):
            choose_folder_button = QtGui.QPushButton("Choose a folder to create Bag from", self)
            choose_folder_button.clicked.connect(self.showDialog)
            choose_folder_button.resize(choose_folder_button.sizeHint())
            choose_folder_button.move(20, 30)

            quit_button = QtGui.QPushButton("Quit", self)
            quit_button.clicked.connect(QtCore.QCoreApplication.instance().quit)
            quit_button.resize(quit_button.sizeHint())
            quit_button.move(250, 30) 

            self.resize(345, 80)
            self.center()
            self.setWindowTitle('Create a Bag')
    
            self.show()
        
        def center(self):        
            qr = self.frameGeometry()
            cp = QtGui.QDesktopWidget().availableGeometry().center()
            qr.moveCenter(cp)
            self.move(qr.topLeft())

        def showDialog(self):
            fname = QtGui.QFileDialog.getExistingDirectory(self, 'Create a Bag - Choose a folder to create Bag from', '/home')

            directory_check(str(fname))
            bag_dir = make_bag(str(fname))

            if (bag_dir):
                self.qt_confirmation(bag_dir)

        def qt_confirmation(self, bag_dir):
            confirmation_window = QtChooserWindow(self)

            confirmation_string = "The Bag for folder " + bag_dir + " has been created."
            confirmation_message = QtGui.QLabel(confirmation_string, confirmation_window)
            confirmation_message.move(20, 30)

            confirmation_window.resize(500, 80)
            confirmation_window.center()
            confirmation_window.setWindowTitle('Bag created')            

            confirmation_window.show()

        def qt_error(self):
            error_window = QtChooserWindow(self)

            error_message = QtGui.QLabel("You are not allowed to run the program on that directory.", error_window)
            error_message.move(20, 30)

            error_window.resize(360, 80)
            error_window.center()
            error_window.setWindowTitle('Sorry')            

            error_window.show()            


    app = QtGui.QApplication(sys.argv)
    ex = QtChooserWindow()
    sys.exit(app.exec_())


# OSX-specific code.
elif platform.system() == 'Darwin':

    directory_check(sys.argv[1])
    bag_dir = make_bag(sys.argv[1])
    cocoaSuccess(bag_dir)