<commit_before>############################################################################
##
## Copyright (C) 2006-2010 University of Utah. All rights reserved.
##
## This file is part of VisTrails.
##
## This file may be used under the terms of the GNU General Public
## License version 2.0 as published by the Free Software Foundation
## and appearing in the file LICENSE.GPL included in the packaging of
## this file.  Please review the following to ensure GNU General Public
## Licensing requirements will be met:
## http://www.opensource.org/licenses/gpl-license.php
##
## If you are unsure which license is appropriate for your use (for
## instance, you are interested in developing a commercial derivative
## of VisTrails), please contact us at vistrails@sci.utah.edu.
##
## This file is provided AS IS with NO WARRANTY OF ANY KIND, INCLUDING THE
## WARRANTY OF DESIGN, MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE.
##
############################################################################
""" The file describes the query tab widget to apply a query/filter to
the current pipeline/version view

QQueryTab
"""

from PyQt4 import QtCore, QtGui
# FIXME broke this as Actions have been changed
#
# from core.vistrail.action import ChangeParameterAction, \
#      ChangeAnnotationAction, AddModulePortAction
from core.vistrail.action import Action
from core.vistrail.location import Location
from core.vistrail.pipeline import Pipeline
from core.vistrail.vistrail import Vistrail
import db.services.io
from core.modules.module_registry import get_module_registry
from gui.method_dropbox import QMethodInputForm
from gui.pipeline_tab import QPipelineTab
from gui.theme import CurrentTheme
from gui.vistrail_controller import VistrailController
import copy

################################################################################

# FIXME: We need new infrastructure for QBE on user-defined types.

class QQueryTab(QPipelineTab):
    """
    QQuery is the similar to the pipeline tab where we can interact
    with pipeline. However, no modules properties is accessibled. Just
    connections. Then we can apply this pipeline to be a query on both
    version and pipeline view

    """
    def __init__(self, parent=None):
        """ QQueryTab(parent: QWidget) -> QQueryTab
        Create an empty vistrail controller for this query tab

        """
        QPipelineTab.__init__(self, parent)
        self.pipelineView.setBackgroundBrush(
            CurrentTheme.QUERY_BACKGROUND_BRUSH)

        self.moduleMethods.vWidget.formType = QFunctionQueryForm

        controller = QueryVistrailController(auto_save=False)
        controller.set_vistrail(Vistrail(), None)
        self.setController(controller)
        controller.change_selected_version(0)
        self.connect(controller,
                     QtCore.SIGNAL('vistrailChanged()'),
                     self.vistrailChanged)

    def vistrailChanged(self):
        """ vistrailChanged() -> None
        Update the pipeline when the vistrail version has changed

        """
        self.controller.current_pipeline.ensure_connection_specs()
        self.emit(QtCore.SIGNAL("queryPipelineChange"),
                  len(self.controller.current_pipeline.modules)>0)

################################################################################

class QFunctionQueryForm(QMethodInputForm):
    def __init__(self, parent=None):
        """ QFunctionQueryForm(parent: QWidget) -> QFunctionQueryForm
        Intialize the widget

        """
        QMethodInputForm.__init__(self, parent)
        self.function = None
        self.fields = []

    def updateFunction(self, function, port_spec):
        """ updateFunction(function: ModuleFunction,
                           port_spec: PortSpec) -> None
        Auto create widgets to describes the function 'function'

        """
        self.function = function
        gLayout = self.layout()
        for pIdx in xrange(len(function.params)):
            p = function.params[pIdx]
            field = QParameterQuery(p)
            self.fields.append(field)
            gLayout.addWidget(field, pIdx, 0)
        self.update()

    def updateMethod(self):
        """ updateMethod() -> None
        Update the method values to vistrail. We only keep a monotonic
        version tree of the query pipeline, we can skip the actions
        here.

        """
        methodBox = self.parent().parent().parent()
        if methodBox.controller:
            paramList = []
            pipeline = methodBox.controller.current_pipeline
            f = pipeline.modules[methodBox.module.id].functions[self.fId]
            p = f.params
            for i in xrange(len(self.fields)):
                p[i].strValue = str(self.fields[i].editor.contents())
                p[i].queryMethod = self.fields[i].selector.getCurrentMethod()

        # Go upstream to update the pipeline
        qtab = self
        while type(qtab)!=QQueryTab and qtab!=None:
            qtab = qtab.parent()
        if qtab:
            qtab.vistrailChanged()

################################################################################

class QParameterQuery(QtGui.QWidget):
    """
    QParameterQuery is a widget containing a line edit and a drop down
    menu allowing users to choose how they want to query a parameter

    """
    def __init__(self, param, parent=None):
        """ QParameterQuery(param: ModuleParam) -> QParameterQuery
        Construct the widget layout

        """
        QtGui.QWidget.__init__(self, parent)
        self.value = param.strValue
        self.type = param.type

        layout = QtGui.QHBoxLayout()
        layout.setSpacing(0)
        layout.setMargin(0)
        self.setLayout(layout)

        self.label = QtGui.QLabel('')
        layout.addWidget(self.label)

        self.selector = QParameterQuerySelector(self.type)
        layout.addWidget(self.selector)

        reg = get_module_registry()
        constant_class = reg.get_module_by_name(param.identifier,
                                                param.type,
                                                param.namespace)
        widget_type = constant_class.get_widget_class()
        self.editor = widget_type(param)
        layout.addWidget(self.editor)

        self.connect(self.selector.operationActionGroup,
                     QtCore.SIGNAL('triggered(QAction*)'),
                     self.operationChanged)
        if self.type =='String':
            self.connect(self.selector.caseActionGroup,
                         QtCore.SIGNAL('triggered(QAction*)'),
                         self.caseChanged)
        self.selector.initAction(param.queryMethod)

    def operationChanged(self, action):
        """ operationChanged(action: QAction) -> None
        Update the text to reflect the operation being used

        """
        self.label.setText(action.text())
        self.updateMethod()

    def caseChanged(self, action):
        """ caseChanged(action: QAction) -> None
        Update the text to reflect the case sensitivity being used

        """
        self.updateMethod()

    def updateMethod(self):
        """ updateMethod() -> None
        Update the method values to vistrail. We only keep a monotonic
        version tree of the query pipeline, we can skip the actions
        here.

        """
        if self.parent():
            self.parent().updateMethod()

################################################################################

class QParameterQuerySelector(QtGui.QToolButton):
    """
    QParameterEditorSelector is a button with a down arrow allowing
    users to select which type of interpolator he/she wants to use

    """
    def __init__(self, pType, parent=None):
        """ QParameterEditorSelector(pType: str, parent: QWidget)
                                     -> QParameterEditorSelector
        Put a stacked widget and a popup button

        """
        QtGui.QToolButton.__init__(self, parent)
        self.type = pType
        self.setAutoRaise(True)
        self.setToolButtonStyle(QtCore.Qt.ToolButtonTextOnly)
        menu = QtGui.QMenu(self)

        self.setPopupMode(QtGui.QToolButton.InstantPopup)

        if pType=='String':
            self.setText(QtCore.QString(QtCore.QChar(0x25bc))) # Down triangle
            self.operationActionGroup = QtGui.QActionGroup(self)
            self.containAction = self.operationActionGroup.addAction('Contain')
            self.containAction.setCheckable(True)

            self.exactAction = self.operationActionGroup.addAction('Exact')
            self.exactAction.setCheckable(True)

            self.regAction = self.operationActionGroup.addAction('Reg Exp')
            self.regAction.setCheckable(True)
            menu.addActions(self.operationActionGroup.actions())
            menu.addSeparator()

            self.caseActionGroup = QtGui.QActionGroup(self)
            self.sensitiveAction = self.caseActionGroup.addAction(
                'Case Sensitive')
            self.sensitiveAction.setCheckable(True)

            self.insensitiveAction = self.caseActionGroup.addAction(
                'Case Insensitive')
            self.insensitiveAction.setCheckable(True)
            menu.addActions(self.caseActionGroup.actions())

        else:
            self.setText('') # Down triangle
            self.operationActionGroup = QtGui.QActionGroup(self)
            self.expAction = self.operationActionGroup.addAction('Expression')
            self.expAction.setCheckable(True)
            self.setEnabled(False)

        self.setMenu(menu)

    def initAction(self, pMethod):
        """ initAction(pMethod: int) -> None
        Select the first choice of selector based on self.type

        """
        if self.type=='String':
            opMap = {0: self.containAction,
                     1: self.exactAction,
                     2: self.regAction}
            caseMap = {0: self.insensitiveAction,
                       1: self.sensitiveAction}
            if pMethod>5:
                pMethod = 0
            opMap[pMethod/2].trigger()
            caseMap[pMethod%2].trigger()
        else:
            self.expAction.trigger()

    def getCurrentMethod(self):
        """ getCurrentMethod() -> int
        Get the current method based on the popup selection

        """
        if self.type=='String':
            opMap = {self.containAction: 0,
                     self.exactAction: 1,
                     self.regAction: 2}
            caseMap = {self.insensitiveAction: 0,
                       self.sensitiveAction: 1}
            method = (opMap[self.operationActionGroup.checkedAction()]*2 +
                      caseMap[self.caseActionGroup.checkedAction()])
            return method
        else:
            return 0

################################################################################

class QueryVistrailController(VistrailController):
    """ QueryVistrailController derives from VistrailController to
    include the copy and paste of query methods

    """

#     def pasteModulesAndConnections(self, str):
#         """ pasteModulesAndConnections(modules: [Module],
#                                        connections: [Connection]) -> version id
#         Paste a list of modules and connections into the current
#         pipeline. Also paste the queryMethod attribute

#         """
#         pipeline = db.services.io.getWorkflowFromXML(str)
#         Pipeline.convert(pipeline)

#         ops = []
#         module_remap = {}
#         for module in pipeline.modules.itervalues():
#             module = copy.copy(module)
#             old_id = module.id
#             if module.location is not None:
#                 loc_id = self.vistrail.idScope.getNewId(Location.vtType)
#                 module.location = Location(id=loc_id,
#                                            x=module.location.x + 10.0,
#                                            y=module.location.y + 10.0,
#                                            )
#             mops = db.services.action.create_copy_op_chain(object=module,
#                                                id_scope=self.vistrail.idScope)
#             module_remap[old_id] = mops[0].db_objectId
#             ops.extend(mops)
#         for connection in pipeline.connections.itervalues():
#             connection = copy.copy(connection)
#             for port in connection.ports:
#                 port.moduleId = module_remap[port.moduleId]
#             ops.extend( \
#                 db.services.action.create_copy_op_chain(object=connection,
#                                                id_scope=self.vistrail.idScope))
#         action = db.services.action.create_action_from_ops(ops)
#         Action.convert(action)

#         self.perform_action(action)

#         self.quiet = True
#         modulesMap ={}
#         modulesToSelect = []
#         actions = []
#         self.previousModuleIds = []
#         for module in modules:
#             name = module.name
#             x = module.center.x + 10.0
#             y = module.center.y + 10.0
#             t = self.addModule(name,x,y)
#             newId = self.vistrail.actionMap[self.currentVersion].module.id
#             self.previousModuleIds.append(newId)
#             modulesMap[module.id]=newId
#             modulesToSelect.append(newId)
#             for fi in range(len(module.functions)):
#                 f = module.functions[fi]
#                 action = ChangeParameterAction()
#                 if f.getNumParams() == 0:
#                     action.addParameter(newId, fi, -1, f.name, "",
#                                         "","", "")
#                 queryMethods = []
#                 for i in range(f.getNumParams()):
#                     p = f.params[i]
#                     if self.currentPipeline.hasAlias(p.alias):
#                         p.alias = ""
#                     action.addParameter(newId, fi, i, f.name, p.name,
#                                         p.strValue, p.type, p.alias)
#                     queryMethods.append(p.queryMethod)
#                 self.performAction(action)
#                 newF = self.currentPipeline.modules[newId].functions[fi]
#                 for i in range(f.getNumParams()):
#                     newF.params[i].queryMethod = queryMethods[i]
#             for (key,value) in module.annotations.items():
#                 action = ChangeAnnotationAction()
#                 action.addAnnotation(newId,key,value)
#                 actions.append(action)
#             if module.registry:
#                 desc = module.registry.getDescriptorByName(module.name)
#                 for (name, spec) in desc.inputPorts.iteritems():
#                     names = [module_registry.registry.getDescriptor(p[0]).name
#                              for p in spec[0]]
#                     action = AddModulePortAction()
#                     action.addModulePort(newId, 'input', name, '('+','.join(names)+')')
#                     actions.append(action)
#                 for (name, spec) in desc.outputPorts.iteritems():
#                     names = [module_registry.registry.getDescriptor(p[0]).name
#                              for p in spec[0]]
#                     action = AddModulePortAction()
#                     action.addModulePort(newId, 'output', name, '('+','.join(names)+')')
#                     actions.append(action)

#         currentAction = self.performBulkActions(actions)

#         for c in connections:
#             conn = copy.copy(c)
#             conn.id = self.currentPipeline.fresh_connection_id()
#             conn.sourceId = modulesMap[conn.sourceId]
#             conn.destinationId = modulesMap[conn.destinationId]
#             currentAction = self.addConnection(conn)
#         self.quiet = False

#         self.currentVersion = currentAction
#         self.invalidate_version_tree()

#     def setPipeline(self, pipeline):
#         """ setPipeline(pipeline) -> None
#         Replace the current pipeline with the given pipeline

#         """
#         # First remove all modules
#         self.deleteModuleList(copy.copy(self.currentPipeline.modules.keys()))

#         # Then paste this new modules and connections
#         self.pasteModulesAndConnections(pipeline.modules.values(),
#                                         pipeline.connections.values())
<commit_msg><bugfix>Fixed issue with query-by-example where parameters couldn't be updated</bugfix><commit_after>############################################################################
##
## Copyright (C) 2006-2010 University of Utah. All rights reserved.
##
## This file is part of VisTrails.
##
## This file may be used under the terms of the GNU General Public
## License version 2.0 as published by the Free Software Foundation
## and appearing in the file LICENSE.GPL included in the packaging of
## this file.  Please review the following to ensure GNU General Public
## Licensing requirements will be met:
## http://www.opensource.org/licenses/gpl-license.php
##
## If you are unsure which license is appropriate for your use (for
## instance, you are interested in developing a commercial derivative
## of VisTrails), please contact us at vistrails@sci.utah.edu.
##
## This file is provided AS IS with NO WARRANTY OF ANY KIND, INCLUDING THE
## WARRANTY OF DESIGN, MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE.
##
############################################################################
""" The file describes the query tab widget to apply a query/filter to
the current pipeline/version view

QQueryTab
"""

from PyQt4 import QtCore, QtGui
# FIXME broke this as Actions have been changed
#
# from core.vistrail.action import ChangeParameterAction, \
#      ChangeAnnotationAction, AddModulePortAction
from core.vistrail.action import Action
from core.vistrail.location import Location
from core.vistrail.pipeline import Pipeline
from core.vistrail.vistrail import Vistrail
import db.services.io
from core.modules.module_registry import get_module_registry
from gui.method_dropbox import QMethodInputForm
from gui.pipeline_tab import QPipelineTab
from gui.theme import CurrentTheme
from gui.vistrail_controller import VistrailController
import copy

################################################################################

# FIXME: We need new infrastructure for QBE on user-defined types.

class QQueryTab(QPipelineTab):
    """
    QQuery is the similar to the pipeline tab where we can interact
    with pipeline. However, no modules properties is accessibled. Just
    connections. Then we can apply this pipeline to be a query on both
    version and pipeline view

    """
    def __init__(self, parent=None):
        """ QQueryTab(parent: QWidget) -> QQueryTab
        Create an empty vistrail controller for this query tab

        """
        QPipelineTab.__init__(self, parent)
        self.pipelineView.setBackgroundBrush(
            CurrentTheme.QUERY_BACKGROUND_BRUSH)

        self.moduleMethods.vWidget.formType = QFunctionQueryForm

        controller = QueryVistrailController(auto_save=False)
        controller.set_vistrail(Vistrail(), None)
        self.setController(controller)
        controller.change_selected_version(0)
        self.connect(controller,
                     QtCore.SIGNAL('vistrailChanged()'),
                     self.vistrailChanged)

    def vistrailChanged(self):
        """ vistrailChanged() -> None
        Update the pipeline when the vistrail version has changed

        """
        self.controller.current_pipeline.ensure_connection_specs()
        self.emit(QtCore.SIGNAL("queryPipelineChange"),
                  len(self.controller.current_pipeline.modules)>0)

################################################################################

class QFunctionQueryForm(QMethodInputForm):
    def __init__(self, parent=None):
        """ QFunctionQueryForm(parent: QWidget) -> QFunctionQueryForm
        Intialize the widget

        """
        QMethodInputForm.__init__(self, parent)
        self.function = None
        self.fields = []

    def updateFunction(self, function, port_spec):
        """ updateFunction(function: ModuleFunction,
                           port_spec: PortSpec) -> None
        Auto create widgets to describes the function 'function'

        """
        self.function = function
        gLayout = self.layout()
        for pIdx in xrange(len(function.params)):
            p = function.params[pIdx]
            field = QParameterQuery(p)
            self.fields.append(field)
            gLayout.addWidget(field, pIdx, 0)
        self.update()

    def updateMethod(self):
        """ updateMethod() -> None
        Update the method values to vistrail. We only keep a monotonic
        version tree of the query pipeline, we can skip the actions
        here.

        """
        methodBox = self.parent().parent().parent().parent()
        if methodBox.controller:
            paramList = []
            pipeline = methodBox.controller.current_pipeline
            f = pipeline.modules[methodBox.module.id].functions[self.fId]
            p = f.params
            for i in xrange(len(self.fields)):
                p[i].strValue = str(self.fields[i].editor.contents())
                p[i].queryMethod = self.fields[i].selector.getCurrentMethod()

        # Go upstream to update the pipeline
        qtab = self
        while type(qtab)!=QQueryTab and qtab!=None:
            qtab = qtab.parent()
        if qtab:
            qtab.vistrailChanged()

################################################################################

class QParameterQuery(QtGui.QWidget):
    """
    QParameterQuery is a widget containing a line edit and a drop down
    menu allowing users to choose how they want to query a parameter

    """
    def __init__(self, param, parent=None):
        """ QParameterQuery(param: ModuleParam) -> QParameterQuery
        Construct the widget layout

        """
        QtGui.QWidget.__init__(self, parent)
        self.value = param.strValue
        self.type = param.type

        layout = QtGui.QHBoxLayout()
        layout.setSpacing(0)
        layout.setMargin(0)
        self.setLayout(layout)

        self.label = QtGui.QLabel('')
        layout.addWidget(self.label)

        self.selector = QParameterQuerySelector(self.type)
        layout.addWidget(self.selector)

        reg = get_module_registry()
        constant_class = reg.get_module_by_name(param.identifier,
                                                param.type,
                                                param.namespace)
        widget_type = constant_class.get_widget_class()
        self.editor = widget_type(param)
        layout.addWidget(self.editor)

        self.connect(self.selector.operationActionGroup,
                     QtCore.SIGNAL('triggered(QAction*)'),
                     self.operationChanged)
        if self.type =='String':
            self.connect(self.selector.caseActionGroup,
                         QtCore.SIGNAL('triggered(QAction*)'),
                         self.caseChanged)
        self.selector.initAction(param.queryMethod)

    def operationChanged(self, action):
        """ operationChanged(action: QAction) -> None
        Update the text to reflect the operation being used

        """
        self.label.setText(action.text())
        self.updateMethod()

    def caseChanged(self, action):
        """ caseChanged(action: QAction) -> None
        Update the text to reflect the case sensitivity being used

        """
        self.updateMethod()

    def updateMethod(self):
        """ updateMethod() -> None
        Update the method values to vistrail. We only keep a monotonic
        version tree of the query pipeline, we can skip the actions
        here.

        """
        if self.parent():
            self.parent().updateMethod()

################################################################################

class QParameterQuerySelector(QtGui.QToolButton):
    """
    QParameterEditorSelector is a button with a down arrow allowing
    users to select which type of interpolator he/she wants to use

    """
    def __init__(self, pType, parent=None):
        """ QParameterEditorSelector(pType: str, parent: QWidget)
                                     -> QParameterEditorSelector
        Put a stacked widget and a popup button

        """
        QtGui.QToolButton.__init__(self, parent)
        self.type = pType
        self.setAutoRaise(True)
        self.setToolButtonStyle(QtCore.Qt.ToolButtonTextOnly)
        menu = QtGui.QMenu(self)

        self.setPopupMode(QtGui.QToolButton.InstantPopup)

        if pType=='String':
            self.setText(QtCore.QString(QtCore.QChar(0x25bc))) # Down triangle
            self.operationActionGroup = QtGui.QActionGroup(self)
            self.containAction = self.operationActionGroup.addAction('Contain')
            self.containAction.setCheckable(True)

            self.exactAction = self.operationActionGroup.addAction('Exact')
            self.exactAction.setCheckable(True)

            self.regAction = self.operationActionGroup.addAction('Reg Exp')
            self.regAction.setCheckable(True)
            menu.addActions(self.operationActionGroup.actions())
            menu.addSeparator()

            self.caseActionGroup = QtGui.QActionGroup(self)
            self.sensitiveAction = self.caseActionGroup.addAction(
                'Case Sensitive')
            self.sensitiveAction.setCheckable(True)

            self.insensitiveAction = self.caseActionGroup.addAction(
                'Case Insensitive')
            self.insensitiveAction.setCheckable(True)
            menu.addActions(self.caseActionGroup.actions())

        else:
            self.setText('') # Down triangle
            self.operationActionGroup = QtGui.QActionGroup(self)
            self.expAction = self.operationActionGroup.addAction('Expression')
            self.expAction.setCheckable(True)
            self.setEnabled(False)

        self.setMenu(menu)

    def initAction(self, pMethod):
        """ initAction(pMethod: int) -> None
        Select the first choice of selector based on self.type

        """
        if self.type=='String':
            opMap = {0: self.containAction,
                     1: self.exactAction,
                     2: self.regAction}
            caseMap = {0: self.insensitiveAction,
                       1: self.sensitiveAction}
            if pMethod>5:
                pMethod = 0
            opMap[pMethod/2].trigger()
            caseMap[pMethod%2].trigger()
        else:
            self.expAction.trigger()

    def getCurrentMethod(self):
        """ getCurrentMethod() -> int
        Get the current method based on the popup selection

        """
        if self.type=='String':
            opMap = {self.containAction: 0,
                     self.exactAction: 1,
                     self.regAction: 2}
            caseMap = {self.insensitiveAction: 0,
                       self.sensitiveAction: 1}
            method = (opMap[self.operationActionGroup.checkedAction()]*2 +
                      caseMap[self.caseActionGroup.checkedAction()])
            return method
        else:
            return 0

################################################################################

class QueryVistrailController(VistrailController):
    """ QueryVistrailController derives from VistrailController to
    include the copy and paste of query methods

    """

#     def pasteModulesAndConnections(self, str):
#         """ pasteModulesAndConnections(modules: [Module],
#                                        connections: [Connection]) -> version id
#         Paste a list of modules and connections into the current
#         pipeline. Also paste the queryMethod attribute

#         """
#         pipeline = db.services.io.getWorkflowFromXML(str)
#         Pipeline.convert(pipeline)

#         ops = []
#         module_remap = {}
#         for module in pipeline.modules.itervalues():
#             module = copy.copy(module)
#             old_id = module.id
#             if module.location is not None:
#                 loc_id = self.vistrail.idScope.getNewId(Location.vtType)
#                 module.location = Location(id=loc_id,
#                                            x=module.location.x + 10.0,
#                                            y=module.location.y + 10.0,
#                                            )
#             mops = db.services.action.create_copy_op_chain(object=module,
#                                                id_scope=self.vistrail.idScope)
#             module_remap[old_id] = mops[0].db_objectId
#             ops.extend(mops)
#         for connection in pipeline.connections.itervalues():
#             connection = copy.copy(connection)
#             for port in connection.ports:
#                 port.moduleId = module_remap[port.moduleId]
#             ops.extend( \
#                 db.services.action.create_copy_op_chain(object=connection,
#                                                id_scope=self.vistrail.idScope))
#         action = db.services.action.create_action_from_ops(ops)
#         Action.convert(action)

#         self.perform_action(action)

#         self.quiet = True
#         modulesMap ={}
#         modulesToSelect = []
#         actions = []
#         self.previousModuleIds = []
#         for module in modules:
#             name = module.name
#             x = module.center.x + 10.0
#             y = module.center.y + 10.0
#             t = self.addModule(name,x,y)
#             newId = self.vistrail.actionMap[self.currentVersion].module.id
#             self.previousModuleIds.append(newId)
#             modulesMap[module.id]=newId
#             modulesToSelect.append(newId)
#             for fi in range(len(module.functions)):
#                 f = module.functions[fi]
#                 action = ChangeParameterAction()
#                 if f.getNumParams() == 0:
#                     action.addParameter(newId, fi, -1, f.name, "",
#                                         "","", "")
#                 queryMethods = []
#                 for i in range(f.getNumParams()):
#                     p = f.params[i]
#                     if self.currentPipeline.hasAlias(p.alias):
#                         p.alias = ""
#                     action.addParameter(newId, fi, i, f.name, p.name,
#                                         p.strValue, p.type, p.alias)
#                     queryMethods.append(p.queryMethod)
#                 self.performAction(action)
#                 newF = self.currentPipeline.modules[newId].functions[fi]
#                 for i in range(f.getNumParams()):
#                     newF.params[i].queryMethod = queryMethods[i]
#             for (key,value) in module.annotations.items():
#                 action = ChangeAnnotationAction()
#                 action.addAnnotation(newId,key,value)
#                 actions.append(action)
#             if module.registry:
#                 desc = module.registry.getDescriptorByName(module.name)
#                 for (name, spec) in desc.inputPorts.iteritems():
#                     names = [module_registry.registry.getDescriptor(p[0]).name
#                              for p in spec[0]]
#                     action = AddModulePortAction()
#                     action.addModulePort(newId, 'input', name, '('+','.join(names)+')')
#                     actions.append(action)
#                 for (name, spec) in desc.outputPorts.iteritems():
#                     names = [module_registry.registry.getDescriptor(p[0]).name
#                              for p in spec[0]]
#                     action = AddModulePortAction()
#                     action.addModulePort(newId, 'output', name, '('+','.join(names)+')')
#                     actions.append(action)

#         currentAction = self.performBulkActions(actions)

#         for c in connections:
#             conn = copy.copy(c)
#             conn.id = self.currentPipeline.fresh_connection_id()
#             conn.sourceId = modulesMap[conn.sourceId]
#             conn.destinationId = modulesMap[conn.destinationId]
#             currentAction = self.addConnection(conn)
#         self.quiet = False

#         self.currentVersion = currentAction
#         self.invalidate_version_tree()

#     def setPipeline(self, pipeline):
#         """ setPipeline(pipeline) -> None
#         Replace the current pipeline with the given pipeline

#         """
#         # First remove all modules
#         self.deleteModuleList(copy.copy(self.currentPipeline.modules.keys()))

#         # Then paste this new modules and connections
#         self.pasteModulesAndConnections(pipeline.modules.values(),
#                                         pipeline.connections.values())
<|endoftext|>


<commit_before># Copyright (C) 2013 Google Inc., authors, and contributors <see AUTHORS file>
# Licensed under http://www.apache.org/licenses/LICENSE-2.0 <see LICENSE file>
# Created By: vraj@reciprocitylabs.com
# Maintained By: vraj@reciprocitylabs.com

from ggrc import db
from .all_models import Directive
from .mixins import Base
from .types import JsonType
from .computed_property import computed_property

class Revision(Base, db.Model):
  __tablename__ = 'revisions'

  resource_id = db.Column(db.Integer, nullable = False)
  resource_type = db.Column(db.String, nullable = False)
  event_id = db.Column(db.Integer, db.ForeignKey('events.id'), nullable = False)
  action = db.Column(db.Enum(u'created', u'modified', u'deleted'), nullable = False)
  content = db.Column(JsonType, nullable=False)

  @staticmethod
  def _extra_table_args(cls):
    return (
        db.Index('revisions_modified_by', 'modified_by_id'),
        )

  _publish_attrs = [
      'resource_id',
      'resource_type',
      'action',
      'content',
      'description',
  ]

  @classmethod
  def eager_query(cls):
    from sqlalchemy import orm

    query = super(Revision, cls).eager_query()
    return query.options(
        orm.subqueryload('modified_by'),
        )

  def __init__(self, obj, modified_by_id, action, content):
    self.resource_id = obj.id
    self.modified_by_id = modified_by_id
    self.resource_type = str(obj.__class__.__name__)
    self.action = action
    self.content = content

  @computed_property
  def description(self):
    link_objects = ['ObjectDocument']
    # TODO: Remove check below if migration can guarantee display_name in content
    if 'display_name' not in self.content:
      return ''
    display_name = self.content['display_name']
    if not display_name:
      result = u"{0} {1}".format(self.resource_type, self.action)
    elif u'<->' in display_name:
      #TODO: Fix too many values to unpack below
      source, destination = display_name.split('<->')[:2]
      if self.resource_type in link_objects:
        if self.action == 'created':
          result = u"{1} linked to {0}".format(source, destination)
        elif self.action == 'deleted':
          result = u"{1} unlinked from {0}".format(source, destination)
        else:
          result = u"{0} {1}".format(display_name, self.action)
      else:
        if self.action == 'created':
          result = u"{1} mapped to {0}".format(source, destination)
        elif self.action == 'deleted':
          result = u"{1} unmapped from {0}".format(source, destination)
        else:
          result = u"{0} {1}".format(display_name, self.action)
    else:
      if 'mapped_directive' in self.content:
        # then this is a special case of combined map/creation
        # should happen only for Section and Control
        mapped_directive = self.content['mapped_directive']
        if self.action == 'created':
          result = u"New {0}, {1}, created and mapped to {2}".format(
              self.resource_type,
              display_name,
              mapped_directive
          )
        elif self.action == 'deleted':
          result = u"{0} unmapped from {1} and deleted".format(
              display_name, mapped_directive)
        else:
          result = u"{0} {1}".format(display_name, self.action)
      else:
        # otherwise, it's a normal creation event
        result = u"{0} {1}".format(display_name, self.action)
    if self.event.action == "IMPORT":
      result += ", via spreadsheet import"
    return result
<commit_msg>Clean up revision model<commit_after># Copyright (C) 2013 Google Inc., authors, and contributors <see AUTHORS file>
# Licensed under http://www.apache.org/licenses/LICENSE-2.0 <see LICENSE file>
# Created By: vraj@reciprocitylabs.com
# Maintained By: vraj@reciprocitylabs.com

"""Defines a Revision model for storing snapshots."""

from ggrc import db
from ggrc.models.computed_property import computed_property
from ggrc.models.mixins import Base
from ggrc.models.types import JsonType

class Revision(Base, db.Model):
  """Revision object holds a JSON snapshot of the object at a time."""

  __tablename__ = 'revisions'

  resource_id = db.Column(db.Integer, nullable=False)
  resource_type = db.Column(db.String, nullable=False)
  event_id = db.Column(db.Integer, db.ForeignKey('events.id'), nullable=False)
  action = db.Column(db.Enum(u'created', u'modified', u'deleted'),
                     nullable=False)
  content = db.Column(JsonType, nullable=False)

  @staticmethod
  def _extra_table_args(_):
    return (db.Index('revisions_modified_by', 'modified_by_id'),)

  _publish_attrs = [
      'resource_id',
      'resource_type',
      'action',
      'content',
      'description',
  ]

  @classmethod
  def eager_query(cls):
    from sqlalchemy import orm

    query = super(Revision, cls).eager_query()
    return query.options(
        orm.subqueryload('modified_by'),
    )

  def __init__(self, obj, modified_by_id, action, content):
    self.resource_id = obj.id
    self.modified_by_id = modified_by_id
    self.resource_type = str(obj.__class__.__name__)
    self.action = action
    self.content = content

  def _description_mapping(self, link_objects):
    """Compute description for revisions with <-> in display name."""
    display_name = self.content['display_name']
    source, destination = display_name.split('<->')[:2]
    mapping_verb = "linked" if self.resource_type in link_objects else "mapped"
    if self.action == 'created':
      result = u"{1} {2} to {0}".format(source, destination, mapping_verb)
    elif self.action == 'deleted':
      result = u"{1} un{2} from {0}".format(source, destination, mapping_verb)
    else:
      result = u"{0} {1}".format(display_name, self.action)
    return result

  @computed_property
  def description(self):
    """Compute a human readable description from action and content."""
    link_objects = ['ObjectDocument']
    if 'display_name' not in self.content:
      return ''
    display_name = self.content['display_name']
    if not display_name:
      result = u"{0} {1}".format(self.resource_type, self.action)
    elif u'<->' in display_name:
      result = self._description_mapping(link_objects)
    else:
      if 'mapped_directive' in self.content:
        # then this is a special case of combined map/creation
        # should happen only for Section and Control
        mapped_directive = self.content['mapped_directive']
        if self.action == 'created':
          result = u"New {0}, {1}, created and mapped to {2}".format(
              self.resource_type,
              display_name,
              mapped_directive
          )
        elif self.action == 'deleted':
          result = u"{0} unmapped from {1} and deleted".format(
              display_name, mapped_directive)
        else:
          result = u"{0} {1}".format(display_name, self.action)
      else:
        # otherwise, it's a normal creation event
        result = u"{0} {1}".format(display_name, self.action)
    if self.event.action == "IMPORT":
      result += ", via spreadsheet import"
    return result
<|endoftext|>


<commit_before>
import os
import re
import json
import time
import types
import socket
import traceback

import classad
htcondor = None

def check_htcondor():
    global htcondor
    if not htcondor:
        os.environ.setdefault('CONDOR_CONFIG', "/etc/condor-ce/condor_config")
        htcondor = __import__("htcondor")
    return htcondor

def _get_pool(environ):
    check_htcondor()
    environ_pool = None
    if environ and 'htcondorce.pool' in environ:
        environ_pool = environ['htcondorce.pool']
    if environ_pool:
        return environ_pool
    if not htcondor:
        return None

    return htcondor.param.get("HTCONDORCE_VIEW_POOL")

def _get_name(environ):
    environ_name = None
    if environ and 'htcondorce.name' in environ:
        environ_name = environ['htcondorce.name']
    if environ_name:
        return environ_name
    check_htcondor()
    if not htcondor:
        return _get_pool(environ)

    return htcondor.param.get("HTCONDORCE_VIEW_NAME")

def get_schedd_objs(environ=None):
    pool = _get_pool(environ)
    if pool:
        name = _get_name(environ)
        coll = htcondor.Collector(pool)
        if name:
            schedds = [coll.locate(htcondor.DaemonTypes.Schedd, name)]
        else:
            schedds = coll.locateAll(htcondor.DaemonTypes.Schedd)
        results = []
        for ad in schedds:
            if not ad.get("Name"):
                continue
            results.append((htcondor.Schedd(ad), ad['Name']))
        return results
    return [(htcondor.Schedd(), socket.getfqdn())]

def get_schedd_ads(environ):
    pool = _get_pool(environ)
    coll = htcondor.Collector(pool)
    if pool:
        name = _get_name(environ)
        if name:
            return [coll.query(htcondor.AdTypes.Schedd, "Name=?=%s" % classad.quote(name))[0]]
        else:
            return coll.query(htcondor.AdTypes.Schedd, "true")
    return [coll.locate(htcondor.DaemonTypes.Schedd)]

def get_schedd_statuses(environ={}):
    ads = get_schedd_ads(environ)
    results = {}
    for ad in ads:
        if 'Name' not in ad:
            continue

        for missing_attr in ['Status', 'IsOK', 'IsWarning', 'IsCritical']:
            if missing_attr in ad:
                continue
            if missing_attr not in htcondor.param:
                continue
            ad[missing_attr] = classad.ExprTree(htcondor.param[missing_attr])

        if 'Status' not in ad:
            results[ad['Name']] = 'Unknown'
        else:
            results[ad['Name']] = ad['Status'].eval()

    return results

def get_schedd_status(environ={}):
    statuses = get_schedd_statuses()
    keys = statuses.keys()
    keys.sort()
    return statuses[keys[0]]

def ad_to_json(ad):
    result = {}
    for key in ad:
        val_expr = ad.lookup(key)
        if classad.ExprTree("%s =?= UNDEFINED" % key).eval(ad):
            result[key] = {"_condor_type": "expr", "expr": val_expr.__repr__()}
        else:
            val = val_expr.eval()
            if isinstance(val, types.ListType) or isinstance(val, types.DictType):
                result[key] = {"_condor_type": "expr", "expr": val_expr.__repr__()}
            else:
                result[key] = val
    return result

def generate_queue_ad(resource_catalog, ce):
    queues = {}
    for entry in resource_catalog:

        queue = entry.get('Transform', {}).get('set_remote_queue', 'default')
        queues[queue] = {'ce': ce,
                         'max_wallclocktime': int(entry.get('MaxWallTime', 1440)),
                         'entry': entry.get('Name', ''),
                         'name': queue,
                         'status': 'Production',
                         'memory': int(entry.get('Memory')),
                         'cpus': int(entry.get('CPUs', 1)),
                         'votag': entry.get('VOTag', ''),
                         'subclusters': entry.get('Subclusters', [])
                        }
    return queues

def agis_data(environ):
    ads = get_schedd_ads(environ)
    results = {"ce_services": {}, "queues": {}, "failed_ces": {}, "resource_groups": {}}
    for ad in ads:
        if 'Name' not in ad:
            continue
        try:
            ce_ad = {
                "endpoint": ad['CollectorHost'],
                "flavour": "HTCONDOR-CE",
                "jobmanager": ad['OSG_BatchSystems'].lower(),
                "name": "%s-CE-HTCondorCE-%s" % (ad['OSG_ResourceGroup'], ad['CollectorHost'].split(':')[0]),
                "site": ad['OSG_ResourceGroup'],
                "status": "Production",
                "type": "CE",
                "version": ad['HTCondorCEVersion']
            }
            rgroup = ad.get('OSG_ResourceGroup')
            if rgroup:
                ce_ad['resource_group'] = rgroup
                results['resource_groups'][rgroup] = {'name': rgroup}
            default_queue_ad = {
                "default": {
                    "ce": ad['OSG_Resource'],
                    "max_cputime": 1440,
                    "max_wallclocktime": 1440,
                    "name": "default",
                    "status": "Production"
                }
            }
            results['ce_services'][ad['OSG_Resource']] = ce_ad
            results['queues'][ad['OSG_Resource']] = default_queue_ad
            queue_ad = generate_queue_ad(ad['OSG_ResourceCatalog'], ad['OSG_Resource'])
            if queue_ad:
                results['queues'][ad['OSG_Resource']] = queue_ad
        except (KeyError, TypeError):
            # No way to log an error, stderr doesn't work, stdout, or logging module
            # So, just add it to the json as "failed_ces"
            error = traceback.format_exc()
            results['failed_ces'][ad['Name']] = error

    return results

def get_spooldir():
    check_htcondor()
    spooldir = htcondor.param.get("HTCONDORCE_VIEW_SPOOL")
    if not spooldir:
        if not os.path.exists("tmp"):
            os.mkdir("tmp")
        spooldir = "tmp"
    return spooldir

<commit_msg>Return full ads in get_schedd_ads() even when only querying the local pool<commit_after>
import os
import re
import json
import time
import types
import socket
import traceback

import classad
htcondor = None

def check_htcondor():
    global htcondor
    if not htcondor:
        os.environ.setdefault('CONDOR_CONFIG', "/etc/condor-ce/condor_config")
        htcondor = __import__("htcondor")
    return htcondor

def _get_pool(environ):
    check_htcondor()
    environ_pool = None
    if environ and 'htcondorce.pool' in environ:
        environ_pool = environ['htcondorce.pool']
    if environ_pool:
        return environ_pool
    if not htcondor:
        return None

    return htcondor.param.get("HTCONDORCE_VIEW_POOL")

def _get_name(environ):
    environ_name = None
    if environ and 'htcondorce.name' in environ:
        environ_name = environ['htcondorce.name']
    if environ_name:
        return environ_name
    check_htcondor()
    if not htcondor:
        return _get_pool(environ)

    return htcondor.param.get("HTCONDORCE_VIEW_NAME")

def get_schedd_objs(environ=None):
    pool = _get_pool(environ)
    if pool:
        name = _get_name(environ)
        coll = htcondor.Collector(pool)
        if name:
            schedds = [coll.locate(htcondor.DaemonTypes.Schedd, name)]
        else:
            schedds = coll.locateAll(htcondor.DaemonTypes.Schedd)
        results = []
        for ad in schedds:
            if not ad.get("Name"):
                continue
            results.append((htcondor.Schedd(ad), ad['Name']))
        return results
    return [(htcondor.Schedd(), socket.getfqdn())]

def get_schedd_ads(environ):
    pool = _get_pool(environ)
    coll = htcondor.Collector(pool)
    if pool:
        name = _get_name(environ)
        if name:
            return [coll.query(htcondor.AdTypes.Schedd, "Name=?=%s" % classad.quote(name))[0]]
        else:
            return coll.query(htcondor.AdTypes.Schedd, "true")

    local_schedd_minimal_ad = coll.locate(htcondor.DaemonTypes.Schedd)
    if not local_schedd_minimal_ad or "Name" not in local_schedd_minimal_ad:
        return []

    local_schedd_ads = coll.query(htcondor.AdTypes.Schedd, "Name=?=%s" % classad.quote(local_schedd_minimal_ad["Name"]))
    return local_schedd_ads

def get_schedd_statuses(environ={}):
    ads = get_schedd_ads(environ)
    results = {}
    for ad in ads:
        if 'Name' not in ad:
            continue

        for missing_attr in ['Status', 'IsOK', 'IsWarning', 'IsCritical']:
            if missing_attr in ad:
                continue
            if missing_attr not in htcondor.param:
                continue
            ad[missing_attr] = classad.ExprTree(htcondor.param[missing_attr])

        if 'Status' not in ad:
            results[ad['Name']] = 'Unknown'
        else:
            results[ad['Name']] = ad['Status'].eval()

    return results

def get_schedd_status(environ={}):
    statuses = get_schedd_statuses()
    keys = statuses.keys()
    keys.sort()
    return statuses[keys[0]]

def ad_to_json(ad):
    result = {}
    for key in ad:
        val_expr = ad.lookup(key)
        if classad.ExprTree("%s =?= UNDEFINED" % key).eval(ad):
            result[key] = {"_condor_type": "expr", "expr": val_expr.__repr__()}
        else:
            val = val_expr.eval()
            if isinstance(val, types.ListType) or isinstance(val, types.DictType):
                result[key] = {"_condor_type": "expr", "expr": val_expr.__repr__()}
            else:
                result[key] = val
    return result

def generate_queue_ad(resource_catalog, ce):
    queues = {}
    for entry in resource_catalog:

        queue = entry.get('Transform', {}).get('set_remote_queue', 'default')
        queues[queue] = {'ce': ce,
                         'max_wallclocktime': int(entry.get('MaxWallTime', 1440)),
                         'entry': entry.get('Name', ''),
                         'name': queue,
                         'status': 'Production',
                         'memory': int(entry.get('Memory')),
                         'cpus': int(entry.get('CPUs', 1)),
                         'votag': entry.get('VOTag', ''),
                         'subclusters': entry.get('Subclusters', [])
                        }
    return queues

def agis_data(environ):
    ads = get_schedd_ads(environ)
    results = {"ce_services": {}, "queues": {}, "failed_ces": {}, "resource_groups": {}}
    for ad in ads:
        if 'Name' not in ad:
            continue
        try:
            ce_ad = {
                "endpoint": ad['CollectorHost'],
                "flavour": "HTCONDOR-CE",
                "jobmanager": ad['OSG_BatchSystems'].lower(),
                "name": "%s-CE-HTCondorCE-%s" % (ad['OSG_ResourceGroup'], ad['CollectorHost'].split(':')[0]),
                "site": ad['OSG_ResourceGroup'],
                "status": "Production",
                "type": "CE",
                "version": ad['HTCondorCEVersion']
            }
            rgroup = ad.get('OSG_ResourceGroup')
            if rgroup:
                ce_ad['resource_group'] = rgroup
                results['resource_groups'][rgroup] = {'name': rgroup}
            default_queue_ad = {
                "default": {
                    "ce": ad['OSG_Resource'],
                    "max_cputime": 1440,
                    "max_wallclocktime": 1440,
                    "name": "default",
                    "status": "Production"
                }
            }
            results['ce_services'][ad['OSG_Resource']] = ce_ad
            results['queues'][ad['OSG_Resource']] = default_queue_ad
            queue_ad = generate_queue_ad(ad['OSG_ResourceCatalog'], ad['OSG_Resource'])
            if queue_ad:
                results['queues'][ad['OSG_Resource']] = queue_ad
        except (KeyError, TypeError):
            # No way to log an error, stderr doesn't work, stdout, or logging module
            # So, just add it to the json as "failed_ces"
            error = traceback.format_exc()
            results['failed_ces'][ad['Name']] = error

    return results

def get_spooldir():
    check_htcondor()
    spooldir = htcondor.param.get("HTCONDORCE_VIEW_SPOOL")
    if not spooldir:
        if not os.path.exists("tmp"):
            os.mkdir("tmp")
        spooldir = "tmp"
    return spooldir

<|endoftext|>


<commit_before># Copyright 2017 Google Inc. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""Single source of truth for dsub's version.

This must remain small and dependency-free so that any dsub module may
import it without creating circular dependencies. Note that this module
is parsed as a text file by setup.py and changes to the format of this
file could break setup.py.

The version should follow formatting requirements specified in PEP-440.
 - https://www.python.org/dev/peps/pep-0440

A typical release sequence will be versioned as:
  0.1.3.dev0 -> 0.1.3 -> 0.1.4.dev0 -> ...
"""

DSUB_VERSION = '0.2.2'
<commit_msg>Bump dsub version to 0.2.3.dev0<commit_after># Copyright 2017 Google Inc. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""Single source of truth for dsub's version.

This must remain small and dependency-free so that any dsub module may
import it without creating circular dependencies. Note that this module
is parsed as a text file by setup.py and changes to the format of this
file could break setup.py.

The version should follow formatting requirements specified in PEP-440.
 - https://www.python.org/dev/peps/pep-0440

A typical release sequence will be versioned as:
  0.1.3.dev0 -> 0.1.3 -> 0.1.4.dev0 -> ...
"""

DSUB_VERSION = '0.2.3.dev0'
<|endoftext|>


<commit_before>"""API for creating and interacting with dtool datasets.

"""

import datetime
import uuid

from pkg_resources import iter_entry_points
from collections import defaultdict

import dtoolcore.utils

__version__ = "3.0.0"

def _generate_storage_broker_lookup():
    """Return dictionary of available storage brokers."""
    storage_broker_lookup = dict()
    for entrypoint in iter_entry_points("dtool.storage_brokers"):
        StorageBroker = entrypoint.load()
        storage_broker_lookup[StorageBroker.key] = StorageBroker
    return storage_broker_lookup

def _get_storage_broker(uri, config_path):
    """Helper function to enable use lookup of appropriate storage brokers."""
    uri = dtoolcore.utils.sanitise_uri(uri)
    storage_broker_lookup = _generate_storage_broker_lookup()
    parsed_uri = dtoolcore.utils.generous_parse_uri(uri)
    StorageBroker = storage_broker_lookup[parsed_uri.scheme]
    return StorageBroker(uri, config_path)

def _admin_metadata_from_uri(uri, config_path):
    """Helper function for getting admin metadata."""
    uri = dtoolcore.utils.sanitise_uri(uri)
    storage_broker = _get_storage_broker(uri, config_path)
    admin_metadata = storage_broker.get_admin_metadata()
    return admin_metadata

def _is_dataset(uri, config_path):
    """Helper function for determining if a URI is a dataset."""
    uri = dtoolcore.utils.sanitise_uri(uri)
    storage_broker = _get_storage_broker(uri, config_path)
    return storage_broker.has_admin_metadata()

def generate_admin_metadata(name, creator_username=None):
    """Return admin metadata as a dictionary."""
    if creator_username is None:
        creator_username = dtoolcore.utils.getuser()
    admin_metadata = {
        "uuid": str(uuid.uuid4()),
        "dtoolcore_version": __version__,
        "name": name,
        "type": "protodataset",
        "creator_username": creator_username,
    }
    return admin_metadata

def _generate_uri(admin_metadata, base_uri):
    """Return dataset URI.

    :param admin_metadata: dataset administrative metadata
    :param base_uri: base URI from which to derive dataset URI
    :returns: dataset URI
    """
    name = admin_metadata["name"]
    uuid = admin_metadata["uuid"]
    # storage_broker_lookup = _generate_storage_broker_lookup()
    # parse_result = urlparse(base_uri)
    # storage = parse_result.scheme
    StorageBroker = _get_storage_broker(base_uri, config_path=None)
    return StorageBroker.generate_uri(name, uuid, base_uri)

def generate_proto_dataset(admin_metadata, base_uri, config_path=None):
    """Return :class:`dtoolcore.ProtoDataSet` instance.

    :param admin_metadata: dataset administrative metadata
    :param base_uri: base URI for proto dataset
    :param config_path: path to dtool configuration file
    """
    uri = _generate_uri(admin_metadata, base_uri)
    return ProtoDataSet(uri, admin_metadata, config_path)

def _copy_create_proto_dataset(
    src_dataset,
    dest_base_uri,
    config_path=None,
    progressbar=None
):
    if progressbar:
        progressbar.label = "Copying dataset"

    admin_metadata = src_dataset._admin_metadata
    admin_metadata["type"] = "protodataset"
    del admin_metadata["frozen_at"]

    proto_dataset = generate_proto_dataset(
        admin_metadata=admin_metadata,
        base_uri=dest_base_uri,
        config_path=config_path
    )

    # Ensure that this bug does not get re-introduced:
    # https://github.com/jic-dtool/dtoolcore/issues/1
    assert proto_dataset._admin_metadata["type"] == "protodataset"
    assert "frozen_at" not in proto_dataset._admin_metadata

    proto_dataset.create()

    return proto_dataset

def _copy_content(src_dataset, dest_proto_dataset, progressbar=None):
    for identifier in src_dataset.identifiers:
        item_properties = src_dataset.item_properties(identifier)
        src_abspath = src_dataset.item_content_abspath(identifier)
        relpath = item_properties["relpath"]
        dest_proto_dataset.put_item(src_abspath, relpath)
        if progressbar:
            progressbar.item_show_func = lambda x: relpath
            progressbar.update(1)

    dest_proto_dataset.put_readme(src_dataset.get_readme_content())

    for overlay_name in src_dataset.list_overlay_names():
        overlay = src_dataset.get_overlay(overlay_name)
        dest_proto_dataset._put_overlay(overlay_name, overlay)

def copy(src_uri, dest_base_uri, config_path=None, progressbar=None):
    """Copy a dataset to another location.

    :param src_uri: URI of dataset to be copied
    :param dest_base_uri: base of URI for copy target
    :param config_path: path to dtool configuration file
    :returns: URI of new dataset
    """
    dataset = DataSet.from_uri(src_uri)

    proto_dataset = _copy_create_proto_dataset(
        dataset,
        dest_base_uri,
        config_path,
        progressbar
    )
    _copy_content(dataset, proto_dataset, progressbar)
    proto_dataset.freeze(progressbar=progressbar)

    return proto_dataset.uri

def copy_resume(src_uri, dest_base_uri, config_path=None, progressbar=None):
    """Copy a dataset to another location.

    :param src_uri: URI of dataset to be copied
    :param dest_base_uri: base of URI for copy target
    :param config_path: path to dtool configuration file
    :returns: URI of new dataset
    """
    dataset = DataSet.from_uri(src_uri)

    # Generate the URI of the destination proto dataset.
    dest_uri = _generate_uri(dataset._admin_metadata, dest_base_uri)
    proto_dataset = ProtoDataSet.from_uri(dest_uri)
    _copy_content(dataset, proto_dataset, progressbar)
    proto_dataset.freeze(progressbar=progressbar)

    return proto_dataset.uri

class DtoolCoreTypeError(TypeError):
    pass

class _BaseDataSet(object):
    """Base class for datasets."""

    def __init__(self, uri, admin_metadata, config_path=None):

        uri = dtoolcore.utils.sanitise_uri(uri)

        self._admin_metadata = admin_metadata
        self._storage_broker = _get_storage_broker(uri, config_path)
        self._uri = uri

    @classmethod
    def _from_uri_with_typecheck(cls, uri, config_path, type_name):
        # Make sure that the URI refers to a dataset.
        if not _is_dataset(uri, config_path):
            raise(DtoolCoreTypeError("{} is not a dataset".format(uri)))

        # Get the admin metadata out of the URI and type check.
        admin_metadata = _admin_metadata_from_uri(uri, config_path)
        if admin_metadata['type'] != type_name:
            raise DtoolCoreTypeError(
                "{} is not a {}".format(uri, cls.__name__))

        # Instantiate and return.
        return cls(uri, admin_metadata, config_path)

    @property
    def uri(self):
        """Return the URI of the dataset."""
        return self._uri

    @property
    def uuid(self):
        """Return the UUID of the dataset."""
        return self._admin_metadata["uuid"]

    @property
    def name(self):
        """Return the name of the dataset."""
        return self._admin_metadata['name']

    def get_readme_content(self):
        """
        Return the content of the README describing the dataset.

        :returns: content of README as a string
        """
        return self._storage_broker.get_readme_content()

    def _put_overlay(self, overlay_name, overlay):
        """Store overlay so that it is accessible by the given name.

        :param overlay_name: name of the overlay
        :param overlay: overlay must be a dictionary where the keys are
                        identifiers in the dataset
        :raises: TypeError if the overlay is not a dictionary,
                 ValueError if identifiers in overlay and dataset do not match
        """
        if not isinstance(overlay, dict):
            raise TypeError("Overlay must be dict")

        if set(self._identifiers()) != set(overlay.keys()):
            raise ValueError("Overlay keys must be dataset identifiers")

        self._storage_broker.put_overlay(overlay_name, overlay)

    def generate_manifest(self, progressbar=None):
        """Return manifest generated from knowledge about contents."""
        items = dict()

        if progressbar:
            progressbar.label = "Generating manifest"

        for handle in self._storage_broker.iter_item_handles():
            key = dtoolcore.utils.generate_identifier(handle)
            value = self._storage_broker.item_properties(handle)
            items[key] = value
            if progressbar:
                progressbar.item_show_func = lambda x: handle
                progressbar.update(1)

        manifest = {
            "items": items,
            "dtoolcore_version": __version__,
            "hash_function": self._storage_broker.hasher.name
        }

        return manifest

class DataSet(_BaseDataSet):
    """
    Class for reading the contents of a dataset.
    """

    def __init__(self, uri, admin_metadata, config_path=None):
        super(DataSet, self).__init__(uri, admin_metadata, config_path)
        self._manifest_cache = None

    def _identifiers(self):
        return self._manifest["items"].keys()

    @classmethod
    def from_uri(cls, uri, config_path=None):
        """
        Return an existing :class:`dtoolcore.DataSet` from a URI.

        :params uri: unique resource identifier where the existing
                     :class:`dtoolcore.DataSet` is stored
        :returns: :class:`dtoolcore.DataSet`
        """
        return cls._from_uri_with_typecheck(uri, config_path, "dataset")

    @property
    def identifiers(self):
        """Return iterable of dataset item identifiers."""
        return self._identifiers()

    @property
    def _manifest(self):
        """Return manifest content."""
        if self._manifest_cache is None:
            self._manifest_cache = self._storage_broker.get_manifest()

        return self._manifest_cache

    def item_properties(self, identifier):
        """Return properties of the item with the given identifier.

        :param identifier: item identifier
        :returns: dictionary of item properties from the manifest
        """
        return self._manifest["items"][identifier]

    def item_content_abspath(self, identifier):
        """Return absolute path at which item content can be accessed.

        :param identifier: item identifier
        :returns: absolute path from which the item content can be accessed
        """
        return self._storage_broker.get_item_abspath(identifier)

    def list_overlay_names(self):
        """Return list of overlay names."""
        return self._storage_broker.list_overlay_names()

    def get_overlay(self, overlay_name):
        """Return overlay as a dictionary.

        :param overlay_name: name of the overlay
        :returns: overlay as a dictionary
        """
        return self._storage_broker.get_overlay(overlay_name)

    def put_overlay(self, overlay_name, overlay):
        """Store overlay so that it is accessible by the given name.

        :param overlay_name: name of the overlay
        :param overlay: overlay must be a dictionary where the keys are
                        identifiers in the dataset
        :raises: TypeError if the overlay is not a dictionary,
                 ValueError if identifiers in overlay and dataset do not match
        """
        self._put_overlay(overlay_name, overlay)

class ProtoDataSet(_BaseDataSet):
    """
    Class for building up a dataset.
    """

    @classmethod
    def from_uri(cls, uri, config_path=None):
        """
        Return an existing :class:`dtoolcore.ProtoDataSet` from a URI.

        :params uri: unique resource identifier where the existing
                     :class:`dtoolcore.ProtoDataSet` is stored
        :returns: :class:`dtoolcore.ProtoDataSet`
        """
        return cls._from_uri_with_typecheck(uri, config_path, "protodataset")

    def _identifiers(self):
        """Return iterable of dataset item identifiers."""
        for handle in self._storage_broker.iter_item_handles():
            yield dtoolcore.utils.generate_identifier(handle)

    def update_name(self, new_name):
        """Update the name of the proto dataset.

        :param new_name: the new name of the proto dataset
        """
        self._admin_metadata['name'] = new_name
        if self._storage_broker.has_admin_metadata():
            self._storage_broker.put_admin_metadata(self._admin_metadata)

    def create(self):
        """Create the required directory structure and admin metadata."""
        self._storage_broker.create_structure()
        self._storage_broker.put_admin_metadata(self._admin_metadata)

    def put_readme(self, content):
        """
        Put content into the README of the dataset.

        The client is responsible for ensuring that the content is valid YAML.

        :param content: string to put into the README
        """
        self._storage_broker.put_readme(content)

    def put_item(self, fpath, relpath):
        """
        Put an item into the dataset.

        :param fpath: path to the item on disk
        :param relpath: relative path name given to the item in the dataset as
                        a handle
        :returns: the handle given to the item
        """
        return self._storage_broker.put_item(fpath, relpath)

    def add_item_metadata(self, handle, key, value):
        """
        Add metadata to a specific item in the :class:`dtoolcore.ProtoDataSet`.

        :param handle: handle representing the relative path of the item in the
                       :class:`dtoolcore.ProtoDataSet`
        :param key: metadata key
        :param value: metadata value
        """
        self._storage_broker.add_item_metadata(handle, key, value)

    def _generate_overlays(self):
        """Return dictionary of overlays generated from added item metadata."""
        overlays = defaultdict(dict)
        for handle in self._storage_broker.iter_item_handles():
            identifier = dtoolcore.utils.generate_identifier(handle)
            item_metadata = self._storage_broker.get_item_metadata(handle)
            for k, v in item_metadata.items():
                overlays[k][identifier] = v

        return overlays

    def freeze(self, progressbar=None):
        """
        Convert :class:`dtoolcore.ProtoDataSet` to :class:`dtoolcore.DataSet`.
        """
        # Call the storage broker pre_freeze hook.
        self._storage_broker.pre_freeze_hook()

        if progressbar:
            progressbar.label = "Freezing dataset"

        # Generate and persist the manifest.
        manifest = self.generate_manifest(progressbar=progressbar)
        self._storage_broker.put_manifest(manifest)

        # Generate and persist overlays from any item metadata that has been
        # added.

        overlays = self._generate_overlays()
        for overlay_name, overlay in overlays.items():
            self._put_overlay(overlay_name, overlay)

        # Change the type of the dataset from "protodataset" to "dataset" and
        # add a "frozen_at" time stamp to the administrative metadata.
        datetime_obj = datetime.datetime.utcnow()
        metadata_update = {
            "type": "dataset",
            "frozen_at": dtoolcore.utils.timestamp(datetime_obj)
        }
        self._admin_metadata.update(metadata_update)
        self._storage_broker.put_admin_metadata(self._admin_metadata)

        # Clean up using the storage broker's post freeze hook.
        self._storage_broker.post_freeze_hook()
<commit_msg>Add some whitepace to make code clearer<commit_after>"""API for creating and interacting with dtool datasets.

"""

import datetime
import uuid

from pkg_resources import iter_entry_points
from collections import defaultdict

import dtoolcore.utils

__version__ = "3.0.0"

def _generate_storage_broker_lookup():
    """Return dictionary of available storage brokers."""
    storage_broker_lookup = dict()
    for entrypoint in iter_entry_points("dtool.storage_brokers"):
        StorageBroker = entrypoint.load()
        storage_broker_lookup[StorageBroker.key] = StorageBroker
    return storage_broker_lookup

def _get_storage_broker(uri, config_path):
    """Helper function to enable use lookup of appropriate storage brokers."""
    uri = dtoolcore.utils.sanitise_uri(uri)
    storage_broker_lookup = _generate_storage_broker_lookup()
    parsed_uri = dtoolcore.utils.generous_parse_uri(uri)
    StorageBroker = storage_broker_lookup[parsed_uri.scheme]
    return StorageBroker(uri, config_path)

def _admin_metadata_from_uri(uri, config_path):
    """Helper function for getting admin metadata."""
    uri = dtoolcore.utils.sanitise_uri(uri)
    storage_broker = _get_storage_broker(uri, config_path)
    admin_metadata = storage_broker.get_admin_metadata()
    return admin_metadata

def _is_dataset(uri, config_path):
    """Helper function for determining if a URI is a dataset."""
    uri = dtoolcore.utils.sanitise_uri(uri)
    storage_broker = _get_storage_broker(uri, config_path)
    return storage_broker.has_admin_metadata()

def generate_admin_metadata(name, creator_username=None):
    """Return admin metadata as a dictionary."""
    if creator_username is None:
        creator_username = dtoolcore.utils.getuser()
    admin_metadata = {
        "uuid": str(uuid.uuid4()),
        "dtoolcore_version": __version__,
        "name": name,
        "type": "protodataset",
        "creator_username": creator_username,
    }
    return admin_metadata

def _generate_uri(admin_metadata, base_uri):
    """Return dataset URI.

    :param admin_metadata: dataset administrative metadata
    :param base_uri: base URI from which to derive dataset URI
    :returns: dataset URI
    """
    name = admin_metadata["name"]
    uuid = admin_metadata["uuid"]
    # storage_broker_lookup = _generate_storage_broker_lookup()
    # parse_result = urlparse(base_uri)
    # storage = parse_result.scheme
    StorageBroker = _get_storage_broker(base_uri, config_path=None)
    return StorageBroker.generate_uri(name, uuid, base_uri)

def generate_proto_dataset(admin_metadata, base_uri, config_path=None):
    """Return :class:`dtoolcore.ProtoDataSet` instance.

    :param admin_metadata: dataset administrative metadata
    :param base_uri: base URI for proto dataset
    :param config_path: path to dtool configuration file
    """
    uri = _generate_uri(admin_metadata, base_uri)
    return ProtoDataSet(uri, admin_metadata, config_path)

def _copy_create_proto_dataset(
    src_dataset,
    dest_base_uri,
    config_path=None,
    progressbar=None
):
    if progressbar:
        progressbar.label = "Copying dataset"

    admin_metadata = src_dataset._admin_metadata
    admin_metadata["type"] = "protodataset"
    del admin_metadata["frozen_at"]

    proto_dataset = generate_proto_dataset(
        admin_metadata=admin_metadata,
        base_uri=dest_base_uri,
        config_path=config_path
    )

    # Ensure that this bug does not get re-introduced:
    # https://github.com/jic-dtool/dtoolcore/issues/1
    assert proto_dataset._admin_metadata["type"] == "protodataset"
    assert "frozen_at" not in proto_dataset._admin_metadata

    proto_dataset.create()

    return proto_dataset

def _copy_content(src_dataset, dest_proto_dataset, progressbar=None):
    for identifier in src_dataset.identifiers:
        item_properties = src_dataset.item_properties(identifier)
        src_abspath = src_dataset.item_content_abspath(identifier)
        relpath = item_properties["relpath"]
        dest_proto_dataset.put_item(src_abspath, relpath)
        if progressbar:
            progressbar.item_show_func = lambda x: relpath
            progressbar.update(1)

    dest_proto_dataset.put_readme(src_dataset.get_readme_content())

    for overlay_name in src_dataset.list_overlay_names():
        overlay = src_dataset.get_overlay(overlay_name)
        dest_proto_dataset._put_overlay(overlay_name, overlay)

def copy(src_uri, dest_base_uri, config_path=None, progressbar=None):
    """Copy a dataset to another location.

    :param src_uri: URI of dataset to be copied
    :param dest_base_uri: base of URI for copy target
    :param config_path: path to dtool configuration file
    :returns: URI of new dataset
    """
    dataset = DataSet.from_uri(src_uri)

    proto_dataset = _copy_create_proto_dataset(
        dataset,
        dest_base_uri,
        config_path,
        progressbar
    )
    _copy_content(dataset, proto_dataset, progressbar)
    proto_dataset.freeze(progressbar=progressbar)

    return proto_dataset.uri

def copy_resume(src_uri, dest_base_uri, config_path=None, progressbar=None):
    """Copy a dataset to another location.

    :param src_uri: URI of dataset to be copied
    :param dest_base_uri: base of URI for copy target
    :param config_path: path to dtool configuration file
    :returns: URI of new dataset
    """
    dataset = DataSet.from_uri(src_uri)

    # Generate the URI of the destination proto dataset.
    dest_uri = _generate_uri(dataset._admin_metadata, dest_base_uri)

    proto_dataset = ProtoDataSet.from_uri(dest_uri)

    _copy_content(dataset, proto_dataset, progressbar)
    proto_dataset.freeze(progressbar=progressbar)

    return proto_dataset.uri

class DtoolCoreTypeError(TypeError):
    pass

class _BaseDataSet(object):
    """Base class for datasets."""

    def __init__(self, uri, admin_metadata, config_path=None):

        uri = dtoolcore.utils.sanitise_uri(uri)

        self._admin_metadata = admin_metadata
        self._storage_broker = _get_storage_broker(uri, config_path)
        self._uri = uri

    @classmethod
    def _from_uri_with_typecheck(cls, uri, config_path, type_name):
        # Make sure that the URI refers to a dataset.
        if not _is_dataset(uri, config_path):
            raise(DtoolCoreTypeError("{} is not a dataset".format(uri)))

        # Get the admin metadata out of the URI and type check.
        admin_metadata = _admin_metadata_from_uri(uri, config_path)
        if admin_metadata['type'] != type_name:
            raise DtoolCoreTypeError(
                "{} is not a {}".format(uri, cls.__name__))

        # Instantiate and return.
        return cls(uri, admin_metadata, config_path)

    @property
    def uri(self):
        """Return the URI of the dataset."""
        return self._uri

    @property
    def uuid(self):
        """Return the UUID of the dataset."""
        return self._admin_metadata["uuid"]

    @property
    def name(self):
        """Return the name of the dataset."""
        return self._admin_metadata['name']

    def get_readme_content(self):
        """
        Return the content of the README describing the dataset.

        :returns: content of README as a string
        """
        return self._storage_broker.get_readme_content()

    def _put_overlay(self, overlay_name, overlay):
        """Store overlay so that it is accessible by the given name.

        :param overlay_name: name of the overlay
        :param overlay: overlay must be a dictionary where the keys are
                        identifiers in the dataset
        :raises: TypeError if the overlay is not a dictionary,
                 ValueError if identifiers in overlay and dataset do not match
        """
        if not isinstance(overlay, dict):
            raise TypeError("Overlay must be dict")

        if set(self._identifiers()) != set(overlay.keys()):
            raise ValueError("Overlay keys must be dataset identifiers")

        self._storage_broker.put_overlay(overlay_name, overlay)

    def generate_manifest(self, progressbar=None):
        """Return manifest generated from knowledge about contents."""
        items = dict()

        if progressbar:
            progressbar.label = "Generating manifest"

        for handle in self._storage_broker.iter_item_handles():
            key = dtoolcore.utils.generate_identifier(handle)
            value = self._storage_broker.item_properties(handle)
            items[key] = value
            if progressbar:
                progressbar.item_show_func = lambda x: handle
                progressbar.update(1)

        manifest = {
            "items": items,
            "dtoolcore_version": __version__,
            "hash_function": self._storage_broker.hasher.name
        }

        return manifest

class DataSet(_BaseDataSet):
    """
    Class for reading the contents of a dataset.
    """

    def __init__(self, uri, admin_metadata, config_path=None):
        super(DataSet, self).__init__(uri, admin_metadata, config_path)
        self._manifest_cache = None

    def _identifiers(self):
        return self._manifest["items"].keys()

    @classmethod
    def from_uri(cls, uri, config_path=None):
        """
        Return an existing :class:`dtoolcore.DataSet` from a URI.

        :params uri: unique resource identifier where the existing
                     :class:`dtoolcore.DataSet` is stored
        :returns: :class:`dtoolcore.DataSet`
        """
        return cls._from_uri_with_typecheck(uri, config_path, "dataset")

    @property
    def identifiers(self):
        """Return iterable of dataset item identifiers."""
        return self._identifiers()

    @property
    def _manifest(self):
        """Return manifest content."""
        if self._manifest_cache is None:
            self._manifest_cache = self._storage_broker.get_manifest()

        return self._manifest_cache

    def item_properties(self, identifier):
        """Return properties of the item with the given identifier.

        :param identifier: item identifier
        :returns: dictionary of item properties from the manifest
        """
        return self._manifest["items"][identifier]

    def item_content_abspath(self, identifier):
        """Return absolute path at which item content can be accessed.

        :param identifier: item identifier
        :returns: absolute path from which the item content can be accessed
        """
        return self._storage_broker.get_item_abspath(identifier)

    def list_overlay_names(self):
        """Return list of overlay names."""
        return self._storage_broker.list_overlay_names()

    def get_overlay(self, overlay_name):
        """Return overlay as a dictionary.

        :param overlay_name: name of the overlay
        :returns: overlay as a dictionary
        """
        return self._storage_broker.get_overlay(overlay_name)

    def put_overlay(self, overlay_name, overlay):
        """Store overlay so that it is accessible by the given name.

        :param overlay_name: name of the overlay
        :param overlay: overlay must be a dictionary where the keys are
                        identifiers in the dataset
        :raises: TypeError if the overlay is not a dictionary,
                 ValueError if identifiers in overlay and dataset do not match
        """
        self._put_overlay(overlay_name, overlay)

class ProtoDataSet(_BaseDataSet):
    """
    Class for building up a dataset.
    """

    @classmethod
    def from_uri(cls, uri, config_path=None):
        """
        Return an existing :class:`dtoolcore.ProtoDataSet` from a URI.

        :params uri: unique resource identifier where the existing
                     :class:`dtoolcore.ProtoDataSet` is stored
        :returns: :class:`dtoolcore.ProtoDataSet`
        """
        return cls._from_uri_with_typecheck(uri, config_path, "protodataset")

    def _identifiers(self):
        """Return iterable of dataset item identifiers."""
        for handle in self._storage_broker.iter_item_handles():
            yield dtoolcore.utils.generate_identifier(handle)

    def update_name(self, new_name):
        """Update the name of the proto dataset.

        :param new_name: the new name of the proto dataset
        """
        self._admin_metadata['name'] = new_name
        if self._storage_broker.has_admin_metadata():
            self._storage_broker.put_admin_metadata(self._admin_metadata)

    def create(self):
        """Create the required directory structure and admin metadata."""
        self._storage_broker.create_structure()
        self._storage_broker.put_admin_metadata(self._admin_metadata)

    def put_readme(self, content):
        """
        Put content into the README of the dataset.

        The client is responsible for ensuring that the content is valid YAML.

        :param content: string to put into the README
        """
        self._storage_broker.put_readme(content)

    def put_item(self, fpath, relpath):
        """
        Put an item into the dataset.

        :param fpath: path to the item on disk
        :param relpath: relative path name given to the item in the dataset as
                        a handle
        :returns: the handle given to the item
        """
        return self._storage_broker.put_item(fpath, relpath)

    def add_item_metadata(self, handle, key, value):
        """
        Add metadata to a specific item in the :class:`dtoolcore.ProtoDataSet`.

        :param handle: handle representing the relative path of the item in the
                       :class:`dtoolcore.ProtoDataSet`
        :param key: metadata key
        :param value: metadata value
        """
        self._storage_broker.add_item_metadata(handle, key, value)

    def _generate_overlays(self):
        """Return dictionary of overlays generated from added item metadata."""
        overlays = defaultdict(dict)
        for handle in self._storage_broker.iter_item_handles():
            identifier = dtoolcore.utils.generate_identifier(handle)
            item_metadata = self._storage_broker.get_item_metadata(handle)
            for k, v in item_metadata.items():
                overlays[k][identifier] = v

        return overlays

    def freeze(self, progressbar=None):
        """
        Convert :class:`dtoolcore.ProtoDataSet` to :class:`dtoolcore.DataSet`.
        """
        # Call the storage broker pre_freeze hook.
        self._storage_broker.pre_freeze_hook()

        if progressbar:
            progressbar.label = "Freezing dataset"

        # Generate and persist the manifest.
        manifest = self.generate_manifest(progressbar=progressbar)
        self._storage_broker.put_manifest(manifest)

        # Generate and persist overlays from any item metadata that has been
        # added.

        overlays = self._generate_overlays()
        for overlay_name, overlay in overlays.items():
            self._put_overlay(overlay_name, overlay)

        # Change the type of the dataset from "protodataset" to "dataset" and
        # add a "frozen_at" time stamp to the administrative metadata.
        datetime_obj = datetime.datetime.utcnow()
        metadata_update = {
            "type": "dataset",
            "frozen_at": dtoolcore.utils.timestamp(datetime_obj)
        }
        self._admin_metadata.update(metadata_update)
        self._storage_broker.put_admin_metadata(self._admin_metadata)

        # Clean up using the storage broker's post freeze hook.
        self._storage_broker.post_freeze_hook()
<|endoftext|>