<commit_before># encoding: utf-8
"""
server.py

Created by David Farrar  on 2011-11-30.
Copyright (c) 2011 Exa Networks. All rights reserved.
"""

from exaproxy.util.log import Logger
from .worker import Client

class ClientManager (object):
	def __init__(self, poller, configuration):
		self.total_sent = 0L
		self.norequest = {}
		self.bysock = {}
		self.byname = {}
		self.buffered = []
		self._nextid = 0
		self.poller = poller
		self.log = Logger('client', configuration.log.client)

	def __contains__(self, item):
		return item in self.byname

	def getnextid(self):
		self._nextid += 1
		return str(self._nextid)

	def newConnection(self, sock, peer, source):
		name = self.getnextid()
		client = Client(name, sock, peer, self.log)

		self.norequest[sock] = client, source
		self.byname[name] = client

		# watch for request data becoming available to read
		self.poller.addReadSocket('opening_client', client.sock)

		self.log.info('new id %s (socket %s) in clients : %s' % (name, sock, sock in self.bysock))
		return peer

	def readRequest(self, sock):
		"""Read only the initial HTTP headers sent by the client"""

		client, source = self.norequest.get(sock, (None, None))
		if client:
			name, peer, request, content = client.readData()
			if request:
				# headers can be read only once
				self.norequest.pop(sock, (None, None))

				# we don't care about new requests from the client
				self.poller.removeReadSocket('opening_client', client.sock)

			elif request is None:
				self.cleanup(sock, client.name)
		else:
			self.log.error('trying to read headers from a client that does not exist %s' % sock)
			name, peer, request, content, source = None, None, None, None, None

		return name, peer, request, content, source

	def readDataBySocket(self, sock):
		client = self.bysock.get(sock, None)
		if client:
			name, peer, request, content = client.readData()
			if request:
				# Parsing of the new request will be handled asynchronously. Ensure that
				# we do not read anything from the client until a request has been sent
				# to the remote webserver.
				# Since we just read a request, we know that the cork is not currently
				# set and so there's no risk of it being erroneously removed.
				self.poller.corkReadSocket('read_client', sock)

			elif request is None:
				self.cleanup(sock, client.name)
		else:
			self.log.error('trying to read from a client that does not exist %s' % sock)
			name, peer, request, content = None, None, None, None

		return name, peer, request, content

	def readDataByName(self, name):
		client = self.byname.get(name, None)
		if client:
			name, peer, request, content = client.readData()
			if request:
				# Parsing of the new request will be handled asynchronously. Ensure that
				# we do not read anything from the client until a request has been sent
				# to the remote webserver.
				# Since we just read a request, we know that the cork is not currently
				# set and so there's no risk of it being erroneously removed.
				self.poller.corkReadSocket('read_client', client.sock)

			elif request is None:
				self.cleanup(client.sock, name)
		else:
			self.log.error('trying to read from a client that does not exist %s' % name)
			name, peer, request, content = None, None, None, None

		return name, peer, request, content

	def sendDataBySocket(self, sock, data):
		client = self.bysock.get(sock, None)
		if client:
			name = client.name
			res = client.writeData(data)

			if res is None:
				# close the client connection
				self.cleanup(sock, client.name)

				buffered, had_buffer,sent = None, None, 0
				result = None
				buffer_change = None
			else:
				buffered, had_buffer, sent = res
				self.total_sent += sent
				result = buffered

			if buffered:
				if sock not in self.buffered:
					self.buffered.append(sock)
					buffer_change = True

					# watch for the socket's send buffer becoming less than full
					self.poller.addWriteSocket('write_client', client.sock)
				else:
					buffer_change = False

			elif had_buffer and sock in self.buffered:
				self.buffered.remove(sock)
				buffer_change = True

				# we no longer care about writing to the client
				self.poller.removeWriteSocket('write_client', client.sock)

			else:
				buffer_change = False
		else:
			result = None
			buffer_change = None
			name = None

		return result, buffer_change, name

	def sendDataByName(self, name, data):
		client = self.byname.get(name, None)
		if client:
			res = client.writeData(data)

			if res is None:
				# close the client connection only if sendDataBySocket is not due to be called
				if client.sock not in self.buffered:
					self.cleanup(client.sock, name)

				buffered, had_buffer, sent = None, None, 0
				#self.total_sent += sent
				result = None
				buffer_change = None
			else:
				buffered, had_buffer, sent = res
				self.total_sent += sent
				result = buffered

			if buffered:
				if client.sock not in self.buffered:
					self.buffered.append(client.sock)
					buffer_change = True

					# watch for the socket's send buffer becoming less than full
					self.poller.addWriteSocket('write_client', client.sock)
				else:
					buffer_change = False

			elif had_buffer and client.sock in self.buffered:
				self.buffered.remove(client.sock)
				buffer_change = True

				# we no longer care about writing to the client
				self.poller.removeWriteSocket('write_client', client.sock)

			else:
				buffer_change = False
		else:
			result = None
			buffer_change = None

		return result, buffer_change

	def startData(self, name, data, remaining):
		client = self.byname.get(name, None)
		if client:
			try:
				command, d = data
			except (ValueError, TypeError):
				self.log.error('invalid command sent to client %s' % name)
				self.cleanup(client.sock, name)
				res = None
			else:
				if client.sock not in self.bysock:
					# Start checking for content sent by the client
					self.bysock[client.sock] = client

					# watch for the client sending new data
					self.poller.addReadSocket('read_client', client.sock)

					# make sure we don't somehow end up with this still here
					self.norequest.pop(client.sock, (None,None))

					# NOTE: always done already in readRequest
					self.poller.removeReadSocket('opening_client', client.sock)
					res = client.startData(command, d)

				else:
					res = client.restartData(command, d)

					# If we are here then we must have prohibited reading from the client
					# and it must otherwise have been in a readable state
					self.poller.uncorkReadSocket('read_client', client.sock)

			if res is not None:
				buffered, had_buffer, sent = res

				# buffered data we read with the HTTP headers
				name, peer, request, content = client.readRelated(remaining)
				if request:
					self.log.error('reading multiple requests')
					self.cleanup(client.sock, name)
					buffered, had_buffer = None, None
					content = None

				elif request is None:
					self.cleanup(client.sock, name)
					buffered, had_buffer = None, None
					content = None

			else:
				# close the client connection only if sendDataBySocket is not due to be called
				if client.sock not in self.buffered:
					self.cleanup(client.sock, name)

				buffered, had_buffer = None, None
				content = None

			if buffered:
				if client.sock not in self.buffered:
					self.buffered.append(client.sock)

					# watch for the socket's send buffer becoming less than full
					self.poller.addWriteSocket('write_client', client.sock)

			elif had_buffer and client.sock in self.buffered:
				self.buffered.remove(client.sock)

				# we no longer care about writing to the client
				self.poller.removeWriteSocket('write_client', client.sock)
		else:
			content = None

		return content

	def corkUploadByName(self, name):
		client = self.byname.get(name, None)
		if client:
			self.poller.corkReadSocket('read_client', client.sock)

	def uncorkUploadByName(self, name):
		client = self.byname.get(name, None)
		if client:
			if client.sock in self.bysock:
				self.poller.uncorkReadSocket('read_client', client.sock)

	def cleanup(self, sock, name):
		self.log.debug('cleanup for socket %s' % sock)
		client = self.bysock.get(sock, None)
		client, source = (client,None) or self.norequest.get(sock, (None,None))
		client = client or self.byname.get(name, None)

		self.bysock.pop(sock, None)
		self.norequest.pop(sock, (None,None))

		if client:
			self.poller.removeWriteSocket('write_client', client.sock)
			self.poller.removeReadSocket('read_client', client.sock)
			self.poller.removeReadSocket('opening_client', client.sock)

			client.shutdown()

		self.byname.pop(name, None)
		if sock in self.buffered:
			self.buffered.remove(sock)

	def stop(self):
		for client in self.bysock.itervalues():
			client.shutdown()

		for client, source in self.norequest.itervalues():
			client.shutdown()

		self.poller.clearRead('read_client')
		self.poller.clearRead('opening_client')
		self.poller.clearWrite('write_client')

		self.bysock = {}
		self.norequest = {}
		self.byname = {}
		self.buffered = []
<commit_msg>Fix test - (client,None) is never False<commit_after># encoding: utf-8
"""
server.py

Created by David Farrar  on 2011-11-30.
Copyright (c) 2011 Exa Networks. All rights reserved.
"""

from exaproxy.util.log import Logger
from .worker import Client

class ClientManager (object):
	def __init__(self, poller, configuration):
		self.total_sent = 0L
		self.norequest = {}
		self.bysock = {}
		self.byname = {}
		self.buffered = []
		self._nextid = 0
		self.poller = poller
		self.log = Logger('client', configuration.log.client)

	def __contains__(self, item):
		return item in self.byname

	def getnextid(self):
		self._nextid += 1
		return str(self._nextid)

	def newConnection(self, sock, peer, source):
		name = self.getnextid()
		client = Client(name, sock, peer, self.log)

		self.norequest[sock] = client, source
		self.byname[name] = client

		# watch for request data becoming available to read
		self.poller.addReadSocket('opening_client', client.sock)

		self.log.info('new id %s (socket %s) in clients : %s' % (name, sock, sock in self.bysock))
		return peer

	def readRequest(self, sock):
		"""Read only the initial HTTP headers sent by the client"""

		client, source = self.norequest.get(sock, (None, None))
		if client:
			name, peer, request, content = client.readData()
			if request:
				# headers can be read only once
				self.norequest.pop(sock, (None, None))

				# we don't care about new requests from the client
				self.poller.removeReadSocket('opening_client', client.sock)

			elif request is None:
				self.cleanup(sock, client.name)
		else:
			self.log.error('trying to read headers from a client that does not exist %s' % sock)
			name, peer, request, content, source = None, None, None, None, None

		return name, peer, request, content, source

	def readDataBySocket(self, sock):
		client = self.bysock.get(sock, None)
		if client:
			name, peer, request, content = client.readData()
			if request:
				# Parsing of the new request will be handled asynchronously. Ensure that
				# we do not read anything from the client until a request has been sent
				# to the remote webserver.
				# Since we just read a request, we know that the cork is not currently
				# set and so there's no risk of it being erroneously removed.
				self.poller.corkReadSocket('read_client', sock)

			elif request is None:
				self.cleanup(sock, client.name)
		else:
			self.log.error('trying to read from a client that does not exist %s' % sock)
			name, peer, request, content = None, None, None, None

		return name, peer, request, content

	def readDataByName(self, name):
		client = self.byname.get(name, None)
		if client:
			name, peer, request, content = client.readData()
			if request:
				# Parsing of the new request will be handled asynchronously. Ensure that
				# we do not read anything from the client until a request has been sent
				# to the remote webserver.
				# Since we just read a request, we know that the cork is not currently
				# set and so there's no risk of it being erroneously removed.
				self.poller.corkReadSocket('read_client', client.sock)

			elif request is None:
				self.cleanup(client.sock, name)
		else:
			self.log.error('trying to read from a client that does not exist %s' % name)
			name, peer, request, content = None, None, None, None

		return name, peer, request, content

	def sendDataBySocket(self, sock, data):
		client = self.bysock.get(sock, None)
		if client:
			name = client.name
			res = client.writeData(data)

			if res is None:
				# close the client connection
				self.cleanup(sock, client.name)

				buffered, had_buffer,sent = None, None, 0
				result = None
				buffer_change = None
			else:
				buffered, had_buffer, sent = res
				self.total_sent += sent
				result = buffered

			if buffered:
				if sock not in self.buffered:
					self.buffered.append(sock)
					buffer_change = True

					# watch for the socket's send buffer becoming less than full
					self.poller.addWriteSocket('write_client', client.sock)
				else:
					buffer_change = False

			elif had_buffer and sock in self.buffered:
				self.buffered.remove(sock)
				buffer_change = True

				# we no longer care about writing to the client
				self.poller.removeWriteSocket('write_client', client.sock)

			else:
				buffer_change = False
		else:
			result = None
			buffer_change = None
			name = None

		return result, buffer_change, name

	def sendDataByName(self, name, data):
		client = self.byname.get(name, None)
		if client:
			res = client.writeData(data)

			if res is None:
				# close the client connection only if sendDataBySocket is not due to be called
				if client.sock not in self.buffered:
					self.cleanup(client.sock, name)

				buffered, had_buffer, sent = None, None, 0
				#self.total_sent += sent
				result = None
				buffer_change = None
			else:
				buffered, had_buffer, sent = res
				self.total_sent += sent
				result = buffered

			if buffered:
				if client.sock not in self.buffered:
					self.buffered.append(client.sock)
					buffer_change = True

					# watch for the socket's send buffer becoming less than full
					self.poller.addWriteSocket('write_client', client.sock)
				else:
					buffer_change = False

			elif had_buffer and client.sock in self.buffered:
				self.buffered.remove(client.sock)
				buffer_change = True

				# we no longer care about writing to the client
				self.poller.removeWriteSocket('write_client', client.sock)

			else:
				buffer_change = False
		else:
			result = None
			buffer_change = None

		return result, buffer_change

	def startData(self, name, data, remaining):
		client = self.byname.get(name, None)
		if client:
			try:
				command, d = data
			except (ValueError, TypeError):
				self.log.error('invalid command sent to client %s' % name)
				self.cleanup(client.sock, name)
				res = None
			else:
				if client.sock not in self.bysock:
					# Start checking for content sent by the client
					self.bysock[client.sock] = client

					# watch for the client sending new data
					self.poller.addReadSocket('read_client', client.sock)

					# make sure we don't somehow end up with this still here
					self.norequest.pop(client.sock, (None,None))

					# NOTE: always done already in readRequest
					self.poller.removeReadSocket('opening_client', client.sock)
					res = client.startData(command, d)

				else:
					res = client.restartData(command, d)

					# If we are here then we must have prohibited reading from the client
					# and it must otherwise have been in a readable state
					self.poller.uncorkReadSocket('read_client', client.sock)

			if res is not None:
				buffered, had_buffer, sent = res

				# buffered data we read with the HTTP headers
				name, peer, request, content = client.readRelated(remaining)
				if request:
					self.log.error('reading multiple requests')
					self.cleanup(client.sock, name)
					buffered, had_buffer = None, None
					content = None

				elif request is None:
					self.cleanup(client.sock, name)
					buffered, had_buffer = None, None
					content = None

			else:
				# close the client connection only if sendDataBySocket is not due to be called
				if client.sock not in self.buffered:
					self.cleanup(client.sock, name)

				buffered, had_buffer = None, None
				content = None

			if buffered:
				if client.sock not in self.buffered:
					self.buffered.append(client.sock)

					# watch for the socket's send buffer becoming less than full
					self.poller.addWriteSocket('write_client', client.sock)

			elif had_buffer and client.sock in self.buffered:
				self.buffered.remove(client.sock)

				# we no longer care about writing to the client
				self.poller.removeWriteSocket('write_client', client.sock)
		else:
			content = None

		return content

	def corkUploadByName(self, name):
		client = self.byname.get(name, None)
		if client:
			self.poller.corkReadSocket('read_client', client.sock)

	def uncorkUploadByName(self, name):
		client = self.byname.get(name, None)
		if client:
			if client.sock in self.bysock:
				self.poller.uncorkReadSocket('read_client', client.sock)

	def cleanup(self, sock, name):
		self.log.debug('cleanup for socket %s' % sock)
		client = self.bysock.get(sock, None)
		client, source = (client,None) if client else self.norequest.get(sock, (None,None))
		client = client or self.byname.get(name, None)

		self.bysock.pop(sock, None)
		self.norequest.pop(sock, (None,None))
		self.byname.pop(name, None)

		if client:
			self.poller.removeWriteSocket('write_client', client.sock)
			self.poller.removeReadSocket('read_client', client.sock)
			self.poller.removeReadSocket('opening_client', client.sock)

			client.shutdown()

		if sock in self.buffered:
			self.buffered.remove(sock)

	def stop(self):
		for client in self.bysock.itervalues():
			client.shutdown()

		for client, source in self.norequest.itervalues():
			client.shutdown()

		self.poller.clearRead('read_client')
		self.poller.clearRead('opening_client')
		self.poller.clearWrite('write_client')

		self.bysock = {}
		self.norequest = {}
		self.byname = {}
		self.buffered = []
<|endoftext|>


<commit_before>from __main__ import vtk, qt, ctk, slicer
import os
import logging

# *************************************************  ********** #
# **************** Longitudinal Quantification **************** #
# ************************************************************* #

class LongitudinalQuantification(slicer.ScriptedLoadableModule.ScriptedLoadableModule):
    def __init__(self, parent):

        slicer.ScriptedLoadableModule.ScriptedLoadableModule.__init__(self, parent)
        parent.title = "Longitudinal Quantification"
        parent.categories = ["Quantification"]
        parent.dependencies = []
        parent.contributors = ["Jean-Baptiste Vimort"]
        parent.helpText = """
            TODO
            """

        parent.acknowledgementText = """
            This work was supported by the National
            Institutes of Dental and Craniofacial Research
            and Biomedical Imaging and Bioengineering of
            the National Institutes of Health under Award
            Number R01DE024450.
            """

        self.parent = parent

# ******************************************************************* #
# **************** Longitudinal Quantification Widget *************** #
# ******************************************************************* #

class LongitudinalQuantificationWidget(slicer.ScriptedLoadableModule.ScriptedLoadableModuleWidget):

    # ************************************************ #
    # ---------------- Initialisation ---------------- #
    # ************************************************ #

    def setup(self):
        print "----- Longitudinal Quantification widget setup -----"

        # ------ Initialisation of Longitudinal quantification and its logic ----- #
        slicer.ScriptedLoadableModule.ScriptedLoadableModuleWidget.setup(self)
        self.logic = LongitudinalQuantificationLogic(self)

        # ---------------------------------------------------------------- #
        # ---------------- Definition of the UI interface ---------------- #
        # ---------------------------------------------------------------- #

        # ------------ Loading of the .ui file ---------- #

        loader = qt.QUiLoader()
        moduleName = 'LongitudinalQuantification'
        scriptedModulesPath = eval('slicer.modules.%s.path' % moduleName.lower())
        scriptedModulesPath = os.path.dirname(scriptedModulesPath)
        path = os.path.join(scriptedModulesPath, 'Resources', 'UI', '%s.ui' %moduleName)

        qfile = qt.QFile(path)
        qfile.open(qt.QFile.ReadOnly)
        widget = loader.load(qfile, self.parent)
        self.layout = self.parent.layout()
        self.widget = widget
        self.layout.addWidget(widget)

        # ------ Scene Collapsible Button ----- #
        self.SceneCollapsibleButton = self.logic.get("SceneCollapsibleButton")
        treeView = self.logic.get("treeView")
        treeView.setMRMLScene(slicer.app.mrmlScene())
        treeView.sceneModel().setHorizontalHeaderLabels(["Models"])
        treeView.sortFilterProxyModel().nodeTypes = ['vtkMRMLModelNode','vtkMRMLMarkupsFiducialNode']
        treeView.header().setVisible(False)
        self.autoChangeLayout = self.logic.get("autoChangeLayout")
        self.computeBox = self.logic.get("computeBox")

        # ------ Preprocessing Collapsible Button ----- #
        self.PreprocessingCollapsibleButton = self.logic.get("PreprocessingCollapsibleButton")
        self.PreprocessingLayout = self.logic.get("PreprocessingLayout")
        self.PreprocessingChoiceComboBox = self.logic.get("PreprocessingChoiceComboBox")

        # ------ Quantification Collapsible Button ----- #
        self.QuantificationCollapsibleButton = self.logic.get("QuantificationCollapsibleButton")
        self.QuantificationLayout = self.logic.get("QuantificationLayout")
        self.QuantificationChoiceComboBox = self.logic.get("QuantificationChoiceComboBox")

        # ------ Analysis Collapsible Button ----- #
        self.AnalysisCollapsibleButton = self.logic.get("AnalysisCollapsibleButton")
        self.AnalysisLayout = self.logic.get("AnalysisLayout")
        self.AnalysisChoiceComboBox = self.logic.get("AnalysisChoiceComboBox")

        # ------------------------------------------------------------------------------ #
        # ---------------- Setup and initialisation of global variables ---------------- #
        # ------------------------------------------------------------------------------ #

        # ------ Initialisation of the other modules is Slicer, if that's not already done ----- #
        if not hasattr(slicer.modules, 'AnglePlanesWidget'):
            slicer.modules.angleplanes.createNewWidgetRepresentation()
        if not hasattr(slicer.modules, 'EasyClipWidget'):
            slicer.modules.easyclip.createNewWidgetRepresentation()
        if not hasattr(slicer.modules, 'MeshStatisticsWidget'):
            slicer.modules.meshstatistics.createNewWidgetRepresentation()
        if not hasattr(slicer.modules, 'PickAndPaintWidget'):
            slicer.modules.pickandpaint.createNewWidgetRepresentation()
        if not hasattr(slicer.modules, 'Q3DCWidget'):
            slicer.modules.q3dc.createNewWidgetRepresentation()
        if not hasattr(slicer.modules, 'SurfaceRegistrationWidget'):
            slicer.modules.surfaceregistration.createNewWidgetRepresentation()

        # ------ Creation of a dictionary that will contain the widgets of all the modules ----- #
        self.ExternalModulesWidgets = dict()
        self.ExternalModulesWidgets["Angle Planes"] = slicer.modules.AnglePlanesWidget.widget
        self.ExternalModulesWidgets["Easy Clip"] = slicer.modules.EasyClipWidget.widget
        self.ExternalModulesWidgets["Mesh Statistics"] = slicer.modules.MeshStatisticsWidget.widget
        self.ExternalModulesWidgets["Model to Model Distance"] = slicer.modules.modeltomodeldistance.widgetRepresentation()
        self.ExternalModulesWidgets["Pick and Paint"] = slicer.modules.PickAndPaintWidget.widget
        self.ExternalModulesWidgets["Q3DC"] = slicer.modules.Q3DCWidget.widget
        self.ExternalModulesWidgets["Shape Population Viewer"] = slicer.modules.launcher.widgetRepresentation()
        self.ExternalModulesWidgets["Surface Registration"] = slicer.modules.SurfaceRegistrationWidget.widget

        # ------ Initialisation of variables to know which module is currently used ----- #
        self.curentQuantificationWidget = dict()
        self.curentQuantificationWidget[self.PreprocessingLayout] = None
        self.curentQuantificationWidget[self.QuantificationLayout] = None
        self.curentQuantificationWidget[self.AnalysisLayout] = None

        # ------ Setup of the external modules ------ #
        # Hiding of the scene tabs and the input tabs in
        # all the external modules to avoid redundancies
        # and make this module as clear and simple as possible

        # AnglePlanes
        slicer.modules.AnglePlanesWidget.SceneCollapsibleButton.hide()
        slicer.modules.AnglePlanesWidget.inputModelLabel.hide()
        slicer.modules.AnglePlanesWidget.inputLandmarksLabel.hide()
        slicer.modules.AnglePlanesWidget.inputModelSelector.hide()
        slicer.modules.AnglePlanesWidget.inputLandmarksSelector.hide()
        slicer.modules.AnglePlanesWidget.loadLandmarksOnSurfacCheckBox.hide()

        #EasyClip
        slicer.modules.EasyClipWidget.SceneCollapsibleButton.hide()

        #MeshStatistics

        #PickAndPaint
        slicer.modules.PickAndPaintWidget.inputModelLabel.hide()
        slicer.modules.PickAndPaintWidget.inputLandmarksLabel.hide()
        slicer.modules.PickAndPaintWidget.inputModelSelector.hide()
        slicer.modules.PickAndPaintWidget.inputLandmarksSelector.hide()
        slicer.modules.PickAndPaintWidget.loadLandmarksOnSurfacCheckBox.hide()

        #Q3DC
        slicer.modules.Q3DCWidget.SceneCollapsibleButton.hide()
        slicer.modules.Q3DCWidget.inputModelLabel.hide()
        slicer.modules.Q3DCWidget.inputLandmarksLabel.hide()
        slicer.modules.Q3DCWidget.inputModelSelector.hide()
        slicer.modules.Q3DCWidget.inputLandmarksSelector.hide()
        slicer.modules.Q3DCWidget.loadLandmarksOnSurfacCheckBox.hide()

        #Surfaceregostration
        slicer.modules.SurfaceRegistrationWidget.InputCollapsibleButton.hide()

        # --------------------------------------------- #
        # ---------------- Connections ---------------- #
        # --------------------------------------------- #

        # ------ Scene Collapsible Button ----- #

        # ------ Preprocessing Collapsible Button ----- #
        self.PreprocessingChoiceComboBox.connect('currentIndexChanged(QString)', self.onPreprocessingSelectionChanged)

        # ------ Quantification Collapsible Button ----- #
        self.QuantificationChoiceComboBox.connect('currentIndexChanged(QString)', self.onQuantificationSelectionChanged)

        # ------ Analysis Collapsible Button ----- #
        self.AnalysisChoiceComboBox.connect('currentIndexChanged(QString)', self.onAnalysisSelectionChanged)

        # ------ Closing of the scene -----#
        slicer.mrmlScene.AddObserver(slicer.mrmlScene.EndCloseEvent, self.onCloseScene)

    # ******************************************* #
    # ---------------- Algorithm ---------------- #
    # ******************************************* #

    # function called each time that the user "enter" in Longitudinal Quantification interface
    def enter(self):
        print "---- Enter Longitudinal Quantification ---- "

    # function called each time that the user "exit" in Longitudinal Quantification interface
    def exit(self):
        print "---- Exit Longitudinal Quantification ---- "

    # function called each time that the scene is closed (if Longitudinal Quantification has been initialized)
    def onCloseScene(self, obj, event):
        print "---- Close Longitudinal Quantification ---- "
        self.PreprocessingChoiceComboBox.setCurrentIndex(0)
        self.QuantificationChoiceComboBox.setCurrentIndex(0)
        self.AnalysisChoiceComboBox.setCurrentIndex(0)

    def onPreprocessingSelectionChanged(self, newModule):
        print "--- onPreprocessingSelectionChanged --- "
        self.moduleChangement(self.PreprocessingLayout, newModule)

    def onQuantificationSelectionChanged(self, newModule):
        print "--- onQuantificationSelectionChanged --- "
        self.moduleChangement(self.QuantificationLayout, newModule)

    def onAnalysisSelectionChanged(self, newModule):
        print "--- onAnalysisSelectionChanged --- "
        self.moduleChangement(self.AnalysisLayout, newModule)

    def moduleChangement(self, layout, newModule):
        if self.curentQuantificationWidget[layout]:
            layout.removeWidget(self.curentQuantificationWidget[layout])
            self.curentQuantificationWidget[layout].hide()
        if newModule == "None":
            self.curentQuantificationWidget[layout] = None
            return
        self.curentQuantificationWidget[layout] = self.ExternalModulesWidgets[newModule]
        layout.addWidget(self.curentQuantificationWidget[layout])
        self.curentQuantificationWidget[layout].show()

# ******************************************************************* #
# **************** Longitudinal Quantification Logic **************** #
# ******************************************************************* #

class LongitudinalQuantificationLogic(slicer.ScriptedLoadableModule.ScriptedLoadableModuleLogic):

    # ************************************************ #
    # ---------------- Initialisation ---------------- #
    # ************************************************ #

    def __init__(self, interface):
        print "----- Longitudinal Quantification logic init -----"
        self.interface = interface

    # ******************************************* #
    # ---------------- Algorithm ---------------- #
    # ******************************************* #

    # ----------- Connection of the User Interface ----------- #
    # This function will look for an object with the given name in the UI and return it.
    def get(self, objectName):
        return self.findWidget(self.interface.widget, objectName)

    # This function will recursively look into all the object of the UI and compare it to
    # the given name, if it never find it will return "None"
    def findWidget(self, widget, objectName):
        if widget.objectName == objectName:
            return widget
        else:
            for w in widget.children():
                resulting_widget = self.findWidget(w, objectName)
                if resulting_widget:
                    return resulting_widget
            return None

# ****************************************************************** #
# **************** Longitudinal Quantification Test **************** #
# ****************************************************************** #

class LongitudinalQuantificationTest(slicer.ScriptedLoadableModule.ScriptedLoadableModuleTest):
    def __init__(self):
        print "----- Longitudinal Quantification test init -----"

    def setUp(self):
        print "----- Longitudinal Quantification test setup -----"
        # reset the state - clear scene
        slicer.mrmlScene.Clear(0)

    def runTest(self):
        # run all tests needed
        self.delayDisplay("Clear the scene")
        self.setUp()
        self.delayDisplay("Download and load datas")
        self.downloaddata()

    def downloaddata(self):
        import urllib
        downloads = (
            ('http://slicer.kitware.com/midas3/download?items=213632', '01.vtk', slicer.util.loadModel),
            ('http://slicer.kitware.com/midas3/download?items=213633', '02.vtk', slicer.util.loadModel),
        )
        for url, name, loader in downloads:
            filePath = slicer.app.temporaryPath + '/' + name
            print filePath
            if not os.path.exists(filePath) or os.stat(filePath).st_size == 0:
                logging.info('Requesting download %s from %s...\n' % (name, url))
                urllib.urlretrieve(url, filePath)
            if loader:
                logging.info('Loading %s...' % (name,))
                loader(filePath)

        layoutManager = slicer.app.layoutManager()
        threeDWidget = layoutManager.threeDWidget(0)
        threeDView = threeDWidget.threeDView()
        threeDView.resetFocalPoint()<commit_msg>STYLE: Factorisation of the previus commit<commit_after>from __main__ import vtk, qt, ctk, slicer
import os
import logging

# *************************************************  ********** #
# **************** Longitudinal Quantification **************** #
# ************************************************************* #

class LongitudinalQuantification(slicer.ScriptedLoadableModule.ScriptedLoadableModule):
    def __init__(self, parent):

        slicer.ScriptedLoadableModule.ScriptedLoadableModule.__init__(self, parent)
        parent.title = "Longitudinal Quantification"
        parent.categories = ["Quantification"]
        parent.dependencies = []
        parent.contributors = ["Jean-Baptiste Vimort"]
        parent.helpText = """
            TODO
            """

        parent.acknowledgementText = """
            This work was supported by the National
            Institutes of Dental and Craniofacial Research
            and Biomedical Imaging and Bioengineering of
            the National Institutes of Health under Award
            Number R01DE024450.
            """

        self.parent = parent

# ******************************************************************* #
# **************** Longitudinal Quantification Widget *************** #
# ******************************************************************* #

class LongitudinalQuantificationWidget(slicer.ScriptedLoadableModule.ScriptedLoadableModuleWidget):

    # ************************************************ #
    # ---------------- Initialisation ---------------- #
    # ************************************************ #

    def setup(self):
        print "----- Longitudinal Quantification widget setup -----"

        # ------ Initialisation of Longitudinal quantification and its logic ----- #
        slicer.ScriptedLoadableModule.ScriptedLoadableModuleWidget.setup(self)
        self.logic = LongitudinalQuantificationLogic(self)

        # ---------------------------------------------------------------- #
        # ---------------- Definition of the UI interface ---------------- #
        # ---------------------------------------------------------------- #

        # ------------ Loading of the .ui file ---------- #

        loader = qt.QUiLoader()
        moduleName = 'LongitudinalQuantification'
        scriptedModulesPath = eval('slicer.modules.%s.path' % moduleName.lower())
        scriptedModulesPath = os.path.dirname(scriptedModulesPath)
        path = os.path.join(scriptedModulesPath, 'Resources', 'UI', '%s.ui' %moduleName)

        qfile = qt.QFile(path)
        qfile.open(qt.QFile.ReadOnly)
        widget = loader.load(qfile, self.parent)
        self.layout = self.parent.layout()
        self.widget = widget
        self.layout.addWidget(widget)

        # ------ Scene Collapsible Button ----- #
        self.SceneCollapsibleButton = self.logic.get("SceneCollapsibleButton")
        treeView = self.logic.get("treeView")
        treeView.setMRMLScene(slicer.app.mrmlScene())
        treeView.sceneModel().setHorizontalHeaderLabels(["Models"])
        treeView.sortFilterProxyModel().nodeTypes = ['vtkMRMLModelNode','vtkMRMLMarkupsFiducialNode']
        treeView.header().setVisible(False)
        self.autoChangeLayout = self.logic.get("autoChangeLayout")
        self.computeBox = self.logic.get("computeBox")

        # ------ Preprocessing Collapsible Button ----- #
        self.PreprocessingCollapsibleButton = self.logic.get("PreprocessingCollapsibleButton")
        self.PreprocessingLayout = self.logic.get("PreprocessingLayout")
        self.PreprocessingChoiceComboBox = self.logic.get("PreprocessingChoiceComboBox")

        # ------ Quantification Collapsible Button ----- #
        self.QuantificationCollapsibleButton = self.logic.get("QuantificationCollapsibleButton")
        self.QuantificationLayout = self.logic.get("QuantificationLayout")
        self.QuantificationChoiceComboBox = self.logic.get("QuantificationChoiceComboBox")

        # ------ Analysis Collapsible Button ----- #
        self.AnalysisCollapsibleButton = self.logic.get("AnalysisCollapsibleButton")
        self.AnalysisLayout = self.logic.get("AnalysisLayout")
        self.AnalysisChoiceComboBox = self.logic.get("AnalysisChoiceComboBox")

        # ------------------------------------------------------------------------------ #
        # ---------------- Setup and initialisation of global variables ---------------- #
        # ------------------------------------------------------------------------------ #

        # ------ Initialisation of the other modules is Slicer, if that's not already done ----- #
        if not hasattr(slicer.modules, 'AnglePlanesWidget'):
            slicer.modules.angleplanes.createNewWidgetRepresentation()
        if not hasattr(slicer.modules, 'EasyClipWidget'):
            slicer.modules.easyclip.createNewWidgetRepresentation()
        if not hasattr(slicer.modules, 'MeshStatisticsWidget'):
            slicer.modules.meshstatistics.createNewWidgetRepresentation()
        if not hasattr(slicer.modules, 'PickAndPaintWidget'):
            slicer.modules.pickandpaint.createNewWidgetRepresentation()
        if not hasattr(slicer.modules, 'Q3DCWidget'):
            slicer.modules.q3dc.createNewWidgetRepresentation()
        if not hasattr(slicer.modules, 'SurfaceRegistrationWidget'):
            slicer.modules.surfaceregistration.createNewWidgetRepresentation()

        # ------ Creation of a dictionary that will contain the pythons modules ----- #
        self.ExternalPythonModules = dict()
        self.ExternalPythonModules["Angle Planes"] = slicer.modules.AnglePlanesWidget
        self.ExternalPythonModules["Easy Clip"] = slicer.modules.EasyClipWidget
        self.ExternalPythonModules["Mesh Statistics"] = slicer.modules.MeshStatisticsWidget
        self.ExternalPythonModules["Pick and Paint"] = slicer.modules.PickAndPaintWidget
        self.ExternalPythonModules["Q3DC"] = slicer.modules.Q3DCWidget
        self.ExternalPythonModules["Surface Registration"] = slicer.modules.SurfaceRegistrationWidget

        # ------ Creation of a dictionary that will contain the CLI modules ----- #
        self.ExternalCLIModules = dict()
        self.ExternalCLIModules["Model to Model Distance"] = slicer.modules.modeltomodeldistance
        self.ExternalCLIModules["Shape Population Viewer"] = slicer.modules.launcher

        # ------ Creation of a dictionary that will contain the widgets of all the modules ----- #
        self.ExternalModulesWidgets = dict()
        for key, value in self.ExternalPythonModules.iteritems():
            self.ExternalModulesWidgets[key] = value.widget
        for key, value in self.ExternalCLIModules.iteritems():
            self.ExternalModulesWidgets[key] = value.widgetRepresentation()

        # ------ Initialisation of variables to know which module is currently used ----- #
        self.curentQuantificationWidget = dict()
        self.curentQuantificationWidget[self.PreprocessingLayout] = None
        self.curentQuantificationWidget[self.QuantificationLayout] = None
        self.curentQuantificationWidget[self.AnalysisLayout] = None

        # ------ Setup of the external modules ------ #
        # Hiding of the scene tabs and the input tabs in
        # all the external modules to avoid redundancies
        # and make this module as clear and simple as possible

        for key, value in self.ExternalPythonModules.iteritems():
            if hasattr(value, 'SceneCollapsibleButton'):
                value.SceneCollapsibleButton.hide()
            if hasattr(value, 'inputModelLabel'):
                value.inputModelLabel.hide()
                value.inputLandmarksLabel.hide()
                value.inputModelSelector.hide()
                value.inputLandmarksSelector.hide()
                value.loadLandmarksOnSurfacCheckBox.hide()
            elif hasattr(value, 'InputCollapsibleButton'):
                value.InputCollapsibleButton.hide()

        # --------------------------------------------- #
        # ---------------- Connections ---------------- #
        # --------------------------------------------- #

        # ------ Scene Collapsible Button ----- #

        # ------ Preprocessing Collapsible Button ----- #
        self.PreprocessingChoiceComboBox.connect('currentIndexChanged(QString)', self.onPreprocessingSelectionChanged)

        # ------ Quantification Collapsible Button ----- #
        self.QuantificationChoiceComboBox.connect('currentIndexChanged(QString)', self.onQuantificationSelectionChanged)

        # ------ Analysis Collapsible Button ----- #
        self.AnalysisChoiceComboBox.connect('currentIndexChanged(QString)', self.onAnalysisSelectionChanged)

        # ------ Closing of the scene -----#
        slicer.mrmlScene.AddObserver(slicer.mrmlScene.EndCloseEvent, self.onCloseScene)

    # ******************************************* #
    # ---------------- Algorithm ---------------- #
    # ******************************************* #

    # function called each time that the user "enter" in Longitudinal Quantification interface
    def enter(self):
        print "---- Enter Longitudinal Quantification ---- "

    # function called each time that the user "exit" in Longitudinal Quantification interface
    def exit(self):
        print "---- Exit Longitudinal Quantification ---- "

    # function called each time that the scene is closed (if Longitudinal Quantification has been initialized)
    def onCloseScene(self, obj, event):
        print "---- Close Longitudinal Quantification ---- "
        self.PreprocessingChoiceComboBox.setCurrentIndex(0)
        self.QuantificationChoiceComboBox.setCurrentIndex(0)
        self.AnalysisChoiceComboBox.setCurrentIndex(0)

    def onPreprocessingSelectionChanged(self, newModule):
        print "--- onPreprocessingSelectionChanged --- "
        self.moduleChangement(self.PreprocessingLayout, newModule)

    def onQuantificationSelectionChanged(self, newModule):
        print "--- onQuantificationSelectionChanged --- "
        self.moduleChangement(self.QuantificationLayout, newModule)

    def onAnalysisSelectionChanged(self, newModule):
        print "--- onAnalysisSelectionChanged --- "
        self.moduleChangement(self.AnalysisLayout, newModule)

    def moduleChangement(self, layout, newModule):
        if self.curentQuantificationWidget[layout]:
            layout.removeWidget(self.curentQuantificationWidget[layout])
            self.curentQuantificationWidget[layout].hide()
        if newModule == "None":
            self.curentQuantificationWidget[layout] = None
            return
        self.curentQuantificationWidget[layout] = self.ExternalModulesWidgets[newModule]
        layout.addWidget(self.curentQuantificationWidget[layout])
        self.curentQuantificationWidget[layout].show()

# ******************************************************************* #
# **************** Longitudinal Quantification Logic **************** #
# ******************************************************************* #

class LongitudinalQuantificationLogic(slicer.ScriptedLoadableModule.ScriptedLoadableModuleLogic):

    # ************************************************ #
    # ---------------- Initialisation ---------------- #
    # ************************************************ #

    def __init__(self, interface):
        print "----- Longitudinal Quantification logic init -----"
        self.interface = interface

    # ******************************************* #
    # ---------------- Algorithm ---------------- #
    # ******************************************* #

    # ----------- Connection of the User Interface ----------- #
    # This function will look for an object with the given name in the UI and return it.
    def get(self, objectName):
        return self.findWidget(self.interface.widget, objectName)

    # This function will recursively look into all the object of the UI and compare it to
    # the given name, if it never find it will return "None"
    def findWidget(self, widget, objectName):
        if widget.objectName == objectName:
            return widget
        else:
            for w in widget.children():
                resulting_widget = self.findWidget(w, objectName)
                if resulting_widget:
                    return resulting_widget
            return None

# ****************************************************************** #
# **************** Longitudinal Quantification Test **************** #
# ****************************************************************** #

class LongitudinalQuantificationTest(slicer.ScriptedLoadableModule.ScriptedLoadableModuleTest):
    def __init__(self):
        print "----- Longitudinal Quantification test init -----"

    def setUp(self):
        print "----- Longitudinal Quantification test setup -----"
        # reset the state - clear scene
        slicer.mrmlScene.Clear(0)

    def runTest(self):
        # run all tests needed
        self.delayDisplay("Clear the scene")
        self.setUp()
        self.delayDisplay("Download and load datas")
        self.downloaddata()

    def downloaddata(self):
        import urllib
        downloads = (
            ('http://slicer.kitware.com/midas3/download?items=213632', '01.vtk', slicer.util.loadModel),
            ('http://slicer.kitware.com/midas3/download?items=213633', '02.vtk', slicer.util.loadModel),
        )
        for url, name, loader in downloads:
            filePath = slicer.app.temporaryPath + '/' + name
            print filePath
            if not os.path.exists(filePath) or os.stat(filePath).st_size == 0:
                logging.info('Requesting download %s from %s...\n' % (name, url))
                urllib.urlretrieve(url, filePath)
            if loader:
                logging.info('Loading %s...' % (name,))
                loader(filePath)

        layoutManager = slicer.app.layoutManager()
        threeDWidget = layoutManager.threeDWidget(0)
        threeDView = threeDWidget.threeDView()
        threeDView.resetFocalPoint()<|endoftext|>


<commit_before>import sys
import os
import uuid
import random
import string
import atexit
from pandaharvester.harvestercore.job_spec import JobSpec
from pandaharvester.harvestercore.file_spec import FileSpec
from pandaharvester.harvestercore.queue_config_mapper import QueueConfigMapper
from pandaharvester.harvestercore.plugin_factory import PluginFactory

file_prefix = 'panda.sgotest.'

def exit_func():
    for f in os.listdir('.'):
        if f.startswith(file_prefix):
            os.remove(f)

atexit.register(exit_func)

queueName = sys.argv[1]
queueConfigMapper = QueueConfigMapper()
queueConfig = queueConfigMapper.get_queue(queueName)

fileSpec = FileSpec()
fileSpec.fileType = 'output'
fileSpec.lfn = file_prefix + uuid.uuid4().hex + '.gz'
fileSpec.fileAttributes = {'guid': str(uuid.uuid4())}
fileSpec.checksum = '0d439274'
assFileSpec = FileSpec()
assFileSpec.lfn = file_prefix + uuid.uuid4().hex
assFileSpec.fileType = 'es_output'
assFileSpec.fsize = random.randint(10, 100)
assFileSpec.path = os.getcwd() + '/' + assFileSpec.lfn
oFile = open(assFileSpec.lfn, 'w')
oFile.write(''.join(random.choice(string.ascii_uppercase + string.digits) for _ in range(assFileSpec.fsize)))
oFile.close()
fileSpec.add_associated_file(assFileSpec)
jobSpec = JobSpec()
jobSpec.jobParams = {'outFiles': fileSpec.lfn + ',log',
                     'scopeOut': 'panda',
                     'scopeLog': 'panda',
                     'logFile': 'log',
                     'realDatasets': 'panda.' + fileSpec.lfn,
                     'ddmEndPointOut': 'BNL-OSG2_DATADISK',
                     }
jobSpec.add_out_file(fileSpec)

pluginFactory = PluginFactory()

# get stage-out plugin
stagerCore = pluginFactory.get_plugin(queueConfig.stager)
print ("plugin={0}".format(stagerCore.__class__.__name__))

print ("testing zip")
tmpStat, tmpOut = stagerCore.zip_output(jobSpec)
if tmpStat:
    print (" OK")
else:
    print (" NG {0}".format(tmpOut))

print ()

print ("testing stage-out")
transferID = None
tmpStat, tmpOut = stagerCore.trigger_stage_out(jobSpec)
if tmpStat:
    if fileSpec.fileAttributes is None and 'transferID' in fileSpec.fileAttributes:
        transferID = fileSpec.fileAttributes['transferID']
    print (" OK transferID={0}".format(transferID))
else:
    print (" NG {0}".format(tmpOut))
    sys.exit(1)

print ()

print ("checking status for transferID={0}".format(transferID))
tmpStat, tmpOut = stagerCore.check_stage_out_status(jobSpec)
if tmpStat:
    print (" OK")
else:
    print (" NG {0}".format(tmpOut))
<commit_msg>Change "checksum" to "chksum"<commit_after>import sys
import os
import uuid
import random
import string
import atexit
from pandaharvester.harvestercore.job_spec import JobSpec
from pandaharvester.harvestercore.file_spec import FileSpec
from pandaharvester.harvestercore.queue_config_mapper import QueueConfigMapper
from pandaharvester.harvestercore.plugin_factory import PluginFactory

file_prefix = 'panda.sgotest.'

def exit_func():
    for f in os.listdir('.'):
        if f.startswith(file_prefix):
            os.remove(f)

atexit.register(exit_func)

queueName = sys.argv[1]
queueConfigMapper = QueueConfigMapper()
queueConfig = queueConfigMapper.get_queue(queueName)

fileSpec = FileSpec()
fileSpec.fileType = 'output'
fileSpec.lfn = file_prefix + uuid.uuid4().hex + '.gz'
fileSpec.fileAttributes = {'guid': str(uuid.uuid4())}
fileSpec.chksum = '0d439274'
assFileSpec = FileSpec()
assFileSpec.lfn = file_prefix + uuid.uuid4().hex
assFileSpec.fileType = 'es_output'
assFileSpec.fsize = random.randint(10, 100)
assFileSpec.path = os.getcwd() + '/' + assFileSpec.lfn
oFile = open(assFileSpec.lfn, 'w')
oFile.write(''.join(random.choice(string.ascii_uppercase + string.digits) for _ in range(assFileSpec.fsize)))
oFile.close()
fileSpec.add_associated_file(assFileSpec)
jobSpec = JobSpec()
jobSpec.jobParams = {'outFiles': fileSpec.lfn + ',log',
                     'scopeOut': 'panda',
                     'scopeLog': 'panda',
                     'logFile': 'log',
                     'realDatasets': 'panda.' + fileSpec.lfn,
                     'ddmEndPointOut': 'BNL-OSG2_DATADISK',
                     }
jobSpec.add_out_file(fileSpec)

pluginFactory = PluginFactory()

# get stage-out plugin
stagerCore = pluginFactory.get_plugin(queueConfig.stager)
print ("plugin={0}".format(stagerCore.__class__.__name__))

print ("testing zip")
tmpStat, tmpOut = stagerCore.zip_output(jobSpec)
if tmpStat:
    print (" OK")
else:
    print (" NG {0}".format(tmpOut))

print ()

print ("testing stage-out")
transferID = None
tmpStat, tmpOut = stagerCore.trigger_stage_out(jobSpec)
if tmpStat:
    if fileSpec.fileAttributes is None and 'transferID' in fileSpec.fileAttributes:
        transferID = fileSpec.fileAttributes['transferID']
    print (" OK transferID={0}".format(transferID))
else:
    print (" NG {0}".format(tmpOut))
    sys.exit(1)

print ()

print ("checking status for transferID={0}".format(transferID))
tmpStat, tmpOut = stagerCore.check_stage_out_status(jobSpec)
if tmpStat:
    print (" OK")
else:
    print (" NG {0}".format(tmpOut))
<|endoftext|>


<commit_before>#!/usr/bin/env python3

import argparse
import itertools
import re
import sys
from collections import defaultdict
from collections import namedtuple
from numbers import Integral
from typing import Callable, Dict, Iterable, Iterator, List, Mapping, Tuple

import pandas as pd
from sklearn.linear_model import LogisticRegression

import game_utterances
import session_data as sd
import utterances

CrossValidationDataFrames = namedtuple("CrossValidationDataFrames", ("training", "testing"))

def __create_regex_disjunction(regexes: Iterable[str]):
	return "(?:" + ")|(?:".join(regexes) + ")"

DEPENDENT_VARIABLE_COL_NAME_PATTERN = re.compile(__create_regex_disjunction(
	("SHAPE_(?:\\S+)", "EDGE_COUNT", "SIZE", "RED", "GREEN", "BLUE", "HUE", "POSITION_X", "POSITION_Y")))
INDEPENDENT_VARIABLE_COL_NAME = "REFERENT"

DYAD_ID_COL_NAME = "DYAD"
OUT_OF_VOCABULARY_TOKEN_LABEL = "__OUT_OF_VOCABULARY__"
'''
See T. Shore and G. Skantze. 2017. "Enhancing reference resolution in dialogue using participant feedback." Grounding Language Understanding 2017
'''
TOKEN_CLASS_SMOOTHING_FREQ_CUTOFF = 4

CATEGORICAL_VAR_COL_NAMES = ("SHAPE", "SUBMITTER")
# NOTE: For some reason, "pandas.get_dummies(..., columns=[col_name_1,...])" works with list objects but not with tuples
CATEGORICAL_DEPENDENT_VAR_COL_NAMES = ["SHAPE"]
assert all(x in CATEGORICAL_VAR_COL_NAMES for x in CATEGORICAL_DEPENDENT_VAR_COL_NAMES)

class CachingSessionDataFrameFactory(object):
	def __init__(self, session_data_frame_factory: Callable[[sd.SessionData], pd.DataFrame]):
		self.session_data_frame_factory = session_data_frame_factory
		self.cache = {}

	def __call__(self, infile: str, session: sd.SessionData) -> pd.DataFrame:
		try:
			session_data_frame = self.cache[infile]
		except KeyError:
			session_data_frame = self.session_data_frame_factory(session)
			session_data_frame[DYAD_ID_COL_NAME] = infile
			self.cache[infile] = session_data_frame
		return session_data_frame

class CrossValidationData(object):
	def __init__(self, testing_data: Tuple[str, sd.SessionData], training_data: Mapping[str, sd.SessionData]):
		self.testing_data = testing_data
		self.training_data = training_data

	@property
	def __key(self):
		return self.testing_data, self.training_data

	def __eq__(self, other):
		return (self is other or (isinstance(other, type(self))
								  and self.__key == other.__key))

	def __hash__(self):
		return hash(self.__key)

	def __ne__(self, other):
		return not (self == other)

	def __repr__(self):
		return self.__class__.__name__ + str(self.__dict__)

class CrossValidationDataFrameFactory(object):
	@staticmethod
	def __categoricize_data(training_feature_df: pd.DataFrame, testing_feature_df: pd.DataFrame):
		for col_name in CATEGORICAL_VAR_COL_NAMES:
			unique_vals = tuple(sorted(frozenset(
				itertools.chain(training_feature_df[col_name].unique(), testing_feature_df[col_name].unique()))))
			training_feature_df[col_name] = pd.Categorical(training_feature_df[col_name], categories=unique_vals,
														   ordered=False)
			testing_feature_df[col_name] = pd.Categorical(testing_feature_df[col_name], categories=unique_vals,
														  ordered=False)

	def __init__(self, session_data_frame_factory: Callable[[str, sd.SessionData], pd.DataFrame]):
		self.session_data_frame_factory = session_data_frame_factory

	def __call__(self, named_session_data=Iterable[Tuple[str, sd.SessionData]]) -> Iterator[CrossValidationDataFrames]:
		for testing_session_name, testing_session_data in named_session_data:
			training_sessions = dict(
				(infile, training_session_data) for (infile, training_session_data) in named_session_data if
				testing_session_data != training_session_data)
			cross_validation_set = CrossValidationData((testing_session_name, testing_session_data),
													   training_sessions)
			yield self.__create_cross_validation_data_frames(cross_validation_set)

	def __create_cross_validation_data_frames(self,
											  cross_validation_data: CrossValidationData) -> CrossValidationDataFrames:
		training_feature_df = pd.concat(self.session_data_frame_factory(infile, session) for (infile, session) in
										cross_validation_data.training_data.items())
		# noinspection PyUnresolvedReferences
		# print("Training data shape: {}".format(training_feature_df.shape), file=sys.stderr)
		# print(training_feature_df)
		testing_feature_df = self.session_data_frame_factory(*cross_validation_data.testing_data)
		# print("Testing data shape: {}".format(testing_feature_df.shape), file=sys.stderr)
		# print(testing_feature_df)

		# noinspection PyTypeChecker
		self.__categoricize_data(training_feature_df, testing_feature_df)
		dummified_training_feature_df = pd.get_dummies(training_feature_df, columns=CATEGORICAL_DEPENDENT_VAR_COL_NAMES)
		dummified_testing_feature_df = pd.get_dummies(testing_feature_df, columns=CATEGORICAL_DEPENDENT_VAR_COL_NAMES)

		return CrossValidationDataFrames(dummified_training_feature_df, dummified_testing_feature_df)

def __create_argparser() -> argparse.ArgumentParser:
	result = argparse.ArgumentParser(
		description="Cross-validation of reference resolution for tangram sessions.")
	result.add_argument("inpaths", metavar="INPATH", nargs='+',
						help="The directories to process.")
	return result

def create_token_type_row_idx_dict(df: pd.DataFrame) -> Dict[str, List[Integral]]:
	result = defaultdict(list)
	for row in df.itertuples():
		# noinspection PyProtectedMember
		row_dict = row._asdict()
		idx = row.Index
		utts = row_dict[game_utterances.SessionGameRoundUtteranceFactory.UTTERANCE_SEQUENCE_COL_NAME]
		for utt in utts:
			tokens = utt.content
			for token in tokens:
				result[token].append(idx)
	return result

def smooth(df: pd.DataFrame) -> Iterator[Tuple[str, List[Integral]]]:
	token_type_row_idxs = create_token_type_row_idx_dict(df)

	# smoothed_token_classes = set()
	unsmoothed_token_type_row_idxs = token_type_row_idxs.items()
	for token_type, row_idxs in unsmoothed_token_type_row_idxs:
		smoothed_token_type = OUT_OF_VOCABULARY_TOKEN_LABEL if len(
			token_type_row_idxs) < TOKEN_CLASS_SMOOTHING_FREQ_CUTOFF else token_type
		yield smoothed_token_type, row_idxs
	# smoothed_token_classes.add(token_type)
	# del token_type_row_idxs[token_type]
	# token_type_row_idxs[OUT_OF_VOCABULARY_TOKEN_LABEL].extend(row_idxs)

	# print("Token class(es) used for smoothing: {}; {} data point(s) used as out-of-vocabulary instance(s).".format(
	#	smoothed_token_classes, len(token_type_row_idxs[OUT_OF_VOCABULARY_TOKEN_LABEL])), file=sys.stderr)
	# return token_type_row_idxs

def __cross_validate(cross_validation_df: CrossValidationDataFrames):
	training_df = cross_validation_df.training
	dependent_var_cols = tuple(
		col for col in training_df.columns if DEPENDENT_VARIABLE_COL_NAME_PATTERN.match(col))
	token_type_row_idxs = smooth(training_df)
	for token_class, training_inst_idxs in token_type_row_idxs:
		training_insts = training_df.loc[training_inst_idxs]
		training_x = training_insts.loc[:, dependent_var_cols]
		training_y = training_insts.loc[:, INDEPENDENT_VARIABLE_COL_NAME]
		model = LogisticRegression()
		model.fit(training_x, training_y)

	testing_df = cross_validation_df.testing
	testing_y = training_df[INDEPENDENT_VARIABLE_COL_NAME]

def __main(args):
	inpaths = args.inpaths
	print("Looking for session data underneath {}.".format(inpaths), file=sys.stderr)
	infile_session_data = tuple(sorted(sd.walk_session_data(inpaths), key=lambda item: item[0]))

	cross_validation_data_frame_factory = CrossValidationDataFrameFactory(
		CachingSessionDataFrameFactory(game_utterances.SessionGameRoundUtteranceFactory(
			utterances.TokenSequenceFactory())))
	print("Creating cross-validation datasets from {} session(s).".format(len(infile_session_data)), file=sys.stderr)
	# NOTE: This must be lazily-generated lest a dataframe be created in-memory for each cross-validation fold
	for cross_validation_df in cross_validation_data_frame_factory(infile_session_data):
		__cross_validate(cross_validation_df)

if __name__ == "__main__":
	__main(__create_argparser().parse_args())
<commit_msg>now ensuring that all token classes with less than 4 instances are used for smoothing, not just the last encountered one<commit_after>#!/usr/bin/env python3

import argparse
import itertools
import re
import sys
from collections import defaultdict
from collections import namedtuple
from numbers import Integral
from typing import Callable, Dict, Iterable, Iterator, List, Mapping, Tuple

import pandas as pd
from sklearn.linear_model import LogisticRegression

import game_utterances
import session_data as sd
import utterances

CrossValidationDataFrames = namedtuple("CrossValidationDataFrames", ("training", "testing"))

def __create_regex_disjunction(regexes: Iterable[str]):
	return "(?:" + ")|(?:".join(regexes) + ")"

DEPENDENT_VARIABLE_COL_NAME_PATTERN = re.compile(__create_regex_disjunction(
	("SHAPE_(?:\\S+)", "EDGE_COUNT", "SIZE", "RED", "GREEN", "BLUE", "HUE", "POSITION_X", "POSITION_Y")))
INDEPENDENT_VARIABLE_COL_NAME = "REFERENT"

DYAD_ID_COL_NAME = "DYAD"
OUT_OF_VOCABULARY_TOKEN_LABEL = "__OUT_OF_VOCABULARY__"
'''
See T. Shore and G. Skantze. 2017. "Enhancing reference resolution in dialogue using participant feedback." Grounding Language Understanding 2017
'''
TOKEN_CLASS_SMOOTHING_FREQ_CUTOFF = 4

CATEGORICAL_VAR_COL_NAMES = ("SHAPE", "SUBMITTER")
# NOTE: For some reason, "pandas.get_dummies(..., columns=[col_name_1,...])" works with list objects but not with tuples
CATEGORICAL_DEPENDENT_VAR_COL_NAMES = ["SHAPE"]
assert all(x in CATEGORICAL_VAR_COL_NAMES for x in CATEGORICAL_DEPENDENT_VAR_COL_NAMES)

class CachingSessionDataFrameFactory(object):
	def __init__(self, session_data_frame_factory: Callable[[sd.SessionData], pd.DataFrame]):
		self.session_data_frame_factory = session_data_frame_factory
		self.cache = {}

	def __call__(self, infile: str, session: sd.SessionData) -> pd.DataFrame:
		try:
			session_data_frame = self.cache[infile]
		except KeyError:
			session_data_frame = self.session_data_frame_factory(session)
			session_data_frame[DYAD_ID_COL_NAME] = infile
			self.cache[infile] = session_data_frame
		return session_data_frame

class CrossValidationData(object):
	def __init__(self, testing_data: Tuple[str, sd.SessionData], training_data: Mapping[str, sd.SessionData]):
		self.testing_data = testing_data
		self.training_data = training_data

	@property
	def __key(self):
		return self.testing_data, self.training_data

	def __eq__(self, other):
		return (self is other or (isinstance(other, type(self))
								  and self.__key == other.__key))

	def __hash__(self):
		return hash(self.__key)

	def __ne__(self, other):
		return not (self == other)

	def __repr__(self):
		return self.__class__.__name__ + str(self.__dict__)

class CrossValidationDataFrameFactory(object):
	@staticmethod
	def __categoricize_data(training_feature_df: pd.DataFrame, testing_feature_df: pd.DataFrame):
		for col_name in CATEGORICAL_VAR_COL_NAMES:
			unique_vals = tuple(sorted(frozenset(
				itertools.chain(training_feature_df[col_name].unique(), testing_feature_df[col_name].unique()))))
			training_feature_df[col_name] = pd.Categorical(training_feature_df[col_name], categories=unique_vals,
														   ordered=False)
			testing_feature_df[col_name] = pd.Categorical(testing_feature_df[col_name], categories=unique_vals,
														  ordered=False)

	def __init__(self, session_data_frame_factory: Callable[[str, sd.SessionData], pd.DataFrame]):
		self.session_data_frame_factory = session_data_frame_factory

	def __call__(self, named_session_data=Iterable[Tuple[str, sd.SessionData]]) -> Iterator[CrossValidationDataFrames]:
		for testing_session_name, testing_session_data in named_session_data:
			training_sessions = dict(
				(infile, training_session_data) for (infile, training_session_data) in named_session_data if
				testing_session_data != training_session_data)
			cross_validation_set = CrossValidationData((testing_session_name, testing_session_data),
													   training_sessions)
			yield self.__create_cross_validation_data_frames(cross_validation_set)

	def __create_cross_validation_data_frames(self,
											  cross_validation_data: CrossValidationData) -> CrossValidationDataFrames:
		training_feature_df = pd.concat(self.session_data_frame_factory(infile, session) for (infile, session) in
										cross_validation_data.training_data.items())
		# noinspection PyUnresolvedReferences
		# print("Training data shape: {}".format(training_feature_df.shape), file=sys.stderr)
		# print(training_feature_df)
		testing_feature_df = self.session_data_frame_factory(*cross_validation_data.testing_data)
		# print("Testing data shape: {}".format(testing_feature_df.shape), file=sys.stderr)
		# print(testing_feature_df)

		# noinspection PyTypeChecker
		self.__categoricize_data(training_feature_df, testing_feature_df)
		dummified_training_feature_df = pd.get_dummies(training_feature_df, columns=CATEGORICAL_DEPENDENT_VAR_COL_NAMES)
		dummified_testing_feature_df = pd.get_dummies(testing_feature_df, columns=CATEGORICAL_DEPENDENT_VAR_COL_NAMES)

		return CrossValidationDataFrames(dummified_training_feature_df, dummified_testing_feature_df)

def __create_argparser() -> argparse.ArgumentParser:
	result = argparse.ArgumentParser(
		description="Cross-validation of reference resolution for tangram sessions.")
	result.add_argument("inpaths", metavar="INPATH", nargs='+',
						help="The directories to process.")
	return result

def create_token_type_row_idx_dict(df: pd.DataFrame) -> Dict[str, List[Integral]]:
	result = defaultdict(list)
	for row in df.itertuples():
		# noinspection PyProtectedMember
		row_dict = row._asdict()
		idx = row.Index
		utts = row_dict[game_utterances.SessionGameRoundUtteranceFactory.UTTERANCE_SEQUENCE_COL_NAME]
		for utt in utts:
			tokens = utt.content
			for token in tokens:
				result[token].append(idx)
	return result

def smooth(df: pd.DataFrame) -> Iterator[Tuple[str, List[Integral]]]:
	token_type_row_idxs = create_token_type_row_idx_dict(df)

	# smoothed_token_classes = set()
	unsmoothed_token_type_row_idxs = token_type_row_idxs.items()
	smoothed_token_types = set()
	smoothed_row_idxs = []
	result = []
	for token_type, row_idxs in unsmoothed_token_type_row_idxs:
		if len(row_idxs) < TOKEN_CLASS_SMOOTHING_FREQ_CUTOFF:
			smoothed_token_types.add(token_type)
			smoothed_row_idxs.extend(row_idxs)
		else:
			result.append((token_type, row_idxs))

	result.append((OUT_OF_VOCABULARY_TOKEN_LABEL, smoothed_row_idxs))
	print("Token class(es) used for smoothing: {}; {} data point(s) used as out-of-vocabulary instance(s).".format(
		smoothed_token_types, len(smoothed_row_idxs), file=sys.stderr))
	return result

def __cross_validate(cross_validation_df: CrossValidationDataFrames):
	training_df = cross_validation_df.training
	dependent_var_cols = tuple(
		col for col in training_df.columns if DEPENDENT_VARIABLE_COL_NAME_PATTERN.match(col))
	token_type_row_idxs = smooth(training_df)
	for token_class, training_inst_idxs in token_type_row_idxs:
		training_insts = training_df.loc[training_inst_idxs]
		training_x = training_insts.loc[:, dependent_var_cols]
		training_y = training_insts.loc[:, INDEPENDENT_VARIABLE_COL_NAME]
		model = LogisticRegression()
		model.fit(training_x, training_y)

	testing_df = cross_validation_df.testing
	testing_y = training_df[INDEPENDENT_VARIABLE_COL_NAME]

def __main(args):
	inpaths = args.inpaths
	print("Looking for session data underneath {}.".format(inpaths), file=sys.stderr)
	infile_session_data = tuple(sorted(sd.walk_session_data(inpaths), key=lambda item: item[0]))

	cross_validation_data_frame_factory = CrossValidationDataFrameFactory(
		CachingSessionDataFrameFactory(game_utterances.SessionGameRoundUtteranceFactory(
			utterances.TokenSequenceFactory())))
	print("Creating cross-validation datasets from {} session(s).".format(len(infile_session_data)), file=sys.stderr)
	# NOTE: This must be lazily-generated lest a dataframe be created in-memory for each cross-validation fold
	for cross_validation_df in cross_validation_data_frame_factory(infile_session_data):
		__cross_validate(cross_validation_df)

if __name__ == "__main__":
	__main(__create_argparser().parse_args())
<|endoftext|>


<commit_before># Copyright 2012 The Chromium Authors. All rights reserved.
# Use of this source code is governed by a BSD-style license that can be
# found in the LICENSE file.

"""Start and stop Web Page Replay.

Of the public module names, the following one is key:
  ReplayServer: a class to start/stop Web Page Replay.
"""

import logging
import os
import re
import signal
import subprocess
import sys
import time
import urllib

_CHROME_SRC_DIR = os.path.abspath(os.path.join(
    os.path.dirname(__file__), os.pardir, os.pardir, os.pardir, os.pardir))
REPLAY_DIR = os.path.join(
    _CHROME_SRC_DIR, 'third_party', 'webpagereplay')
LOG_PATH = os.path.join(
    _CHROME_SRC_DIR, 'webpagereplay_logs', 'logs.txt')

# Chrome options to make it work with Web Page Replay.
def GetChromeFlags(replay_host, http_port, https_port):
  assert replay_host and http_port and https_port, 'All arguments required'
  return [
      '--host-resolver-rules=MAP * %s,EXCLUDE localhost' % replay_host,
      '--testing-fixed-http-port=%s' % http_port,
      '--testing-fixed-https-port=%s' % https_port,
      '--ignore-certificate-errors',
      ]

# Signal masks on Linux are inherited from parent processes.  If anything
# invoking us accidentally masks SIGINT (e.g. by putting a process in the
# background from a shell script), sending a SIGINT to the child will fail
# to terminate it.  Running this signal handler before execing should fix that
# problem.
def ResetInterruptHandler():
  signal.signal(signal.SIGINT, signal.SIG_DFL)

class ReplayError(Exception):
  """Catch-all exception for the module."""
  pass

class ReplayNotFoundError(ReplayError):
  def __init__(self, label, path):
    super(ReplayNotFoundError, self).__init__()
    self.args = (label, path)

  def __str__(self):
    label, path = self.args
    return 'Path does not exist for %s: %s' % (label, path)

class ReplayNotStartedError(ReplayError):
  pass

class ReplayServer(object):
  """Start and Stop Web Page Replay.

  Web Page Replay is a proxy that can record and "replay" web pages with
  simulated network characteristics -- without having to edit the pages
  by hand. With WPR, tests can use "real" web content, and catch
  performance issues that may result from introducing network delays and
  bandwidth throttling.

  Example:
     with ReplayServer(archive_path):
       self.NavigateToURL(start_url)
       self.WaitUntil(...)

  Environment Variables (for development):
    WPR_ARCHIVE_PATH: path to alternate archive file (e.g. '/tmp/foo.wpr').
    WPR_RECORD: if set, puts Web Page Replay in record mode instead of replay.
    WPR_REPLAY_DIR: path to alternate Web Page Replay source.
  """

  def __init__(self, archive_path, replay_host, dns_port, http_port, https_port,
               replay_options=None, replay_dir=None,
               log_path=None):
    """Initialize ReplayServer.

    Args:
      archive_path: a path to a specific WPR archive (required).
      replay_host: the hostname to serve traffic.
      dns_port: an integer port on which to serve DNS traffic. May be zero
          to let the OS choose an available port. If None DNS forwarding is
          disabled.
      http_port: an integer port on which to serve HTTP traffic. May be zero
          to let the OS choose an available port.
      https_port: an integer port on which to serve HTTPS traffic. May be zero
          to let the OS choose an available port.
      replay_options: an iterable of options strings to forward to replay.py.
      replay_dir: directory that has replay.py and related modules.
      log_path: a path to a log file.
    """
    self.archive_path = os.environ.get('WPR_ARCHIVE_PATH', archive_path)
    self.replay_options = list(replay_options or ())
    self.replay_dir = os.environ.get('WPR_REPLAY_DIR', replay_dir or REPLAY_DIR)
    self.log_path = log_path or LOG_PATH
    self.dns_port = dns_port
    self.http_port = http_port
    self.https_port = https_port
    self._replay_host = replay_host

    if 'WPR_RECORD' in os.environ and '--record' not in self.replay_options:
      self.replay_options.append('--record')
    self.is_record_mode = '--record' in self.replay_options
    self._AddDefaultReplayOptions()

    self.replay_py = os.path.join(self.replay_dir, 'replay.py')

    if self.is_record_mode:
      self._CheckPath('archive directory', os.path.dirname(self.archive_path))
    elif not os.path.exists(self.archive_path):
      self._CheckPath('archive file', self.archive_path)
    self._CheckPath('replay script', self.replay_py)

    self.replay_process = None

  def _AddDefaultReplayOptions(self):
    """Set WPR command-line options. Can be overridden if needed."""
    self.replay_options = [
        '--host', str(self._replay_host),
        '--port', str(self.http_port),
        '--ssl_port', str(self.https_port),
        '--use_closest_match',
        '--no-dns_forwarding',
        '--log_level', 'warning'
        ] + self.replay_options
    if self.dns_port is not None:
      self.replay_options.extend(['--dns_port', str(self.dns_port)])

  def _CheckPath(self, label, path):
    if not os.path.exists(path):
      raise ReplayNotFoundError(label, path)

  def _OpenLogFile(self):
    log_dir = os.path.dirname(self.log_path)
    if not os.path.exists(log_dir):
      os.makedirs(log_dir)
    return open(self.log_path, 'w')

  def WaitForStart(self, timeout):
    """Checks to see if the server is up and running."""
    port_re = re.compile(
        '.*?(?P<protocol>[A-Z]+) server started on (?P<host>.*):(?P<port>\d+)')

    start_time = time.time()
    elapsed_time = 0
    while elapsed_time < timeout:
      if self.replay_process.poll() is not None:
        break  # The process has exited.

      # Read the ports from the WPR log.
      if not self.http_port or not self.https_port or not self.dns_port:
        with open(self.log_path) as f:
          for line in f.readlines():
            m = port_re.match(line.strip())
            if m:
              if not self.http_port and m.group('protocol') == 'HTTP':
                self.http_port = int(m.group('port'))
              elif not self.https_port and m.group('protocol') == 'HTTPS':
                self.https_port = int(m.group('port'))
              elif not self.dns_port and m.group('protocol') == 'DNS':
                self.dns_port = int(m.group('port'))

      # Try to connect to the WPR ports.
      if self.http_port and self.https_port:
        try:
          up_url = '%s://%s:%s/web-page-replay-generate-200'
          http_up_url = up_url % ('http', self._replay_host, self.http_port)
          https_up_url = up_url % ('https', self._replay_host, self.https_port)
          if (200 == urllib.urlopen(http_up_url, None, {}).getcode() and
              200 == urllib.urlopen(https_up_url, None, {}).getcode()):
            return True
        except IOError:
          pass

      poll_interval = min(max(elapsed_time / 10., .1), 5)
      time.sleep(poll_interval)
      elapsed_time = time.time() - start_time

    return False

  def StartServer(self):
    """Start Web Page Replay and verify that it started.

    Raises:
      ReplayNotStartedError: if Replay start-up fails.
    """
    cmd_line = [sys.executable, self.replay_py]
    cmd_line.extend(self.replay_options)
    cmd_line.append(self.archive_path)

    logging.debug('Starting Web-Page-Replay: %s', cmd_line)
    with self._OpenLogFile() as log_fh:
      kwargs = {'stdout': log_fh, 'stderr': subprocess.STDOUT}
      if sys.platform.startswith('linux') or sys.platform == 'darwin':
        kwargs['preexec_fn'] = ResetInterruptHandler
      self.replay_process = subprocess.Popen(cmd_line, **kwargs)

    if not self.WaitForStart(30):
      with open(self.log_path) as f:
        log = f.read()
      raise ReplayNotStartedError(
          'Web Page Replay failed to start. Log output:\n%s' % log)

  def StopServer(self):
    """Stop Web Page Replay."""
    if self.replay_process:
      logging.debug('Trying to stop Web-Page-Replay gracefully')
      try:
        urllib.urlopen('http://%s:%s/web-page-replay-command-exit' % (
            self._replay_host, self.http_port), None, {}).close()
      except IOError:
        # IOError is possible because the server might exit without response.
        pass

      start_time = time.time()
      while time.time() - start_time < 10:  # Timeout after 10 seconds.
        if self.replay_process.poll() is not None:
          break
        time.sleep(1)
      else:
        try:
          # Use a SIGINT so that it can do graceful cleanup.
          self.replay_process.send_signal(signal.SIGINT)
        except:  # pylint: disable=W0702
          # On Windows, we are left with no other option than terminate().
          if 'no-dns_forwarding' not in self.replay_options:
            logging.warning('DNS configuration might not be restored!')
          try:
            self.replay_process.terminate()
          except:  # pylint: disable=W0702
            pass
        self.replay_process.wait()

  def __enter__(self):
    """Add support for with-statement."""
    self.StartServer()
    return self

  def __exit__(self, unused_exc_type, unused_exc_val, unused_exc_tb):
    """Add support for with-statement."""
    self.StopServer()
<commit_msg>[Telemetry] Use WaitFor in webpagereplay.py.<commit_after># Copyright 2012 The Chromium Authors. All rights reserved.
# Use of this source code is governed by a BSD-style license that can be
# found in the LICENSE file.

"""Start and stop Web Page Replay.

Of the public module names, the following one is key:
  ReplayServer: a class to start/stop Web Page Replay.
"""

import logging
import os
import re
import signal
import subprocess
import sys
import urllib

from telemetry.core import util

_CHROME_SRC_DIR = os.path.abspath(os.path.join(
    os.path.dirname(__file__), os.pardir, os.pardir, os.pardir, os.pardir))
REPLAY_DIR = os.path.join(
    _CHROME_SRC_DIR, 'third_party', 'webpagereplay')
LOG_PATH = os.path.join(
    _CHROME_SRC_DIR, 'webpagereplay_logs', 'logs.txt')

# Chrome options to make it work with Web Page Replay.
def GetChromeFlags(replay_host, http_port, https_port):
  assert replay_host and http_port and https_port, 'All arguments required'
  return [
      '--host-resolver-rules=MAP * %s,EXCLUDE localhost' % replay_host,
      '--testing-fixed-http-port=%s' % http_port,
      '--testing-fixed-https-port=%s' % https_port,
      '--ignore-certificate-errors',
      ]

# Signal masks on Linux are inherited from parent processes.  If anything
# invoking us accidentally masks SIGINT (e.g. by putting a process in the
# background from a shell script), sending a SIGINT to the child will fail
# to terminate it.  Running this signal handler before execing should fix that
# problem.
def ResetInterruptHandler():
  signal.signal(signal.SIGINT, signal.SIG_DFL)

class ReplayError(Exception):
  """Catch-all exception for the module."""
  pass

class ReplayNotFoundError(ReplayError):
  def __init__(self, label, path):
    super(ReplayNotFoundError, self).__init__()
    self.args = (label, path)

  def __str__(self):
    label, path = self.args
    return 'Path does not exist for %s: %s' % (label, path)

class ReplayNotStartedError(ReplayError):
  pass

class ReplayServer(object):
  """Start and Stop Web Page Replay.

  Web Page Replay is a proxy that can record and "replay" web pages with
  simulated network characteristics -- without having to edit the pages
  by hand. With WPR, tests can use "real" web content, and catch
  performance issues that may result from introducing network delays and
  bandwidth throttling.

  Example:
     with ReplayServer(archive_path):
       self.NavigateToURL(start_url)
       self.WaitUntil(...)

  Environment Variables (for development):
    WPR_ARCHIVE_PATH: path to alternate archive file (e.g. '/tmp/foo.wpr').
    WPR_RECORD: if set, puts Web Page Replay in record mode instead of replay.
    WPR_REPLAY_DIR: path to alternate Web Page Replay source.
  """

  def __init__(self, archive_path, replay_host, dns_port, http_port, https_port,
               replay_options=None, replay_dir=None,
               log_path=None):
    """Initialize ReplayServer.

    Args:
      archive_path: a path to a specific WPR archive (required).
      replay_host: the hostname to serve traffic.
      dns_port: an integer port on which to serve DNS traffic. May be zero
          to let the OS choose an available port. If None DNS forwarding is
          disabled.
      http_port: an integer port on which to serve HTTP traffic. May be zero
          to let the OS choose an available port.
      https_port: an integer port on which to serve HTTPS traffic. May be zero
          to let the OS choose an available port.
      replay_options: an iterable of options strings to forward to replay.py.
      replay_dir: directory that has replay.py and related modules.
      log_path: a path to a log file.
    """
    self.archive_path = os.environ.get('WPR_ARCHIVE_PATH', archive_path)
    self.replay_options = list(replay_options or ())
    self.replay_dir = os.environ.get('WPR_REPLAY_DIR', replay_dir or REPLAY_DIR)
    self.log_path = log_path or LOG_PATH
    self.dns_port = dns_port
    self.http_port = http_port
    self.https_port = https_port
    self._replay_host = replay_host

    if 'WPR_RECORD' in os.environ and '--record' not in self.replay_options:
      self.replay_options.append('--record')
    self.is_record_mode = '--record' in self.replay_options
    self._AddDefaultReplayOptions()

    self.replay_py = os.path.join(self.replay_dir, 'replay.py')

    if self.is_record_mode:
      self._CheckPath('archive directory', os.path.dirname(self.archive_path))
    elif not os.path.exists(self.archive_path):
      self._CheckPath('archive file', self.archive_path)
    self._CheckPath('replay script', self.replay_py)

    self.replay_process = None

  def _AddDefaultReplayOptions(self):
    """Set WPR command-line options. Can be overridden if needed."""
    self.replay_options = [
        '--host', str(self._replay_host),
        '--port', str(self.http_port),
        '--ssl_port', str(self.https_port),
        '--use_closest_match',
        '--no-dns_forwarding',
        '--log_level', 'warning'
        ] + self.replay_options
    if self.dns_port is not None:
      self.replay_options.extend(['--dns_port', str(self.dns_port)])

  def _CheckPath(self, label, path):
    if not os.path.exists(path):
      raise ReplayNotFoundError(label, path)

  def _OpenLogFile(self):
    log_dir = os.path.dirname(self.log_path)
    if not os.path.exists(log_dir):
      os.makedirs(log_dir)
    return open(self.log_path, 'w')

  def IsStarted(self):
    """Checks to see if the server is up and running."""
    port_re = re.compile(
        '.*?(?P<protocol>[A-Z]+) server started on (?P<host>.*):(?P<port>\d+)')

    if self.replay_process.poll() is not None:
      return False

    # Read the ports from the WPR log.
    if not self.http_port or not self.https_port or not self.dns_port:
      with open(self.log_path) as f:
        for line in f.readlines():
          m = port_re.match(line.strip())
          if m:
            if not self.http_port and m.group('protocol') == 'HTTP':
              self.http_port = int(m.group('port'))
            elif not self.https_port and m.group('protocol') == 'HTTPS':
              self.https_port = int(m.group('port'))
            elif not self.dns_port and m.group('protocol') == 'DNS':
              self.dns_port = int(m.group('port'))

    # Try to connect to the WPR ports.
    if self.http_port and self.https_port:
      try:
        up_url = '%s://%s:%s/web-page-replay-generate-200'
        http_up_url = up_url % ('http', self._replay_host, self.http_port)
        https_up_url = up_url % ('https', self._replay_host, self.https_port)
        if (200 == urllib.urlopen(http_up_url, None, {}).getcode() and
            200 == urllib.urlopen(https_up_url, None, {}).getcode()):
          return True
      except IOError:
        pass
    return False

  def StartServer(self):
    """Start Web Page Replay and verify that it started.

    Raises:
      ReplayNotStartedError: if Replay start-up fails.
    """
    cmd_line = [sys.executable, self.replay_py]
    cmd_line.extend(self.replay_options)
    cmd_line.append(self.archive_path)

    logging.debug('Starting Web-Page-Replay: %s', cmd_line)
    with self._OpenLogFile() as log_fh:
      kwargs = {'stdout': log_fh, 'stderr': subprocess.STDOUT}
      if sys.platform.startswith('linux') or sys.platform == 'darwin':
        kwargs['preexec_fn'] = ResetInterruptHandler
      self.replay_process = subprocess.Popen(cmd_line, **kwargs)

    try:
      util.WaitFor(self.IsStarted, 30)
    except util.TimeoutException:
      with open(self.log_path) as f:
        log = f.read()
      raise ReplayNotStartedError(
          'Web Page Replay failed to start. Log output:\n%s' % log)

  def StopServer(self):
    """Stop Web Page Replay."""
    if not self.replay_process:
      return

    logging.debug('Trying to stop Web-Page-Replay gracefully')
    try:
      urllib.urlopen('http://%s:%s/web-page-replay-command-exit' % (
          self._replay_host, self.http_port), None, {}).close()
    except IOError:
      # IOError is possible because the server might exit without response.
      pass

    try:
      util.WaitFor(lambda: self.replay_process.poll() is not None, 10)
    except util.TimeoutException:
      try:
        # Use a SIGINT so that it can do graceful cleanup.
        self.replay_process.send_signal(signal.SIGINT)
      except:  # pylint: disable=W0702
        # On Windows, we are left with no other option than terminate().
        if 'no-dns_forwarding' not in self.replay_options:
          logging.warning('DNS configuration might not be restored!')
        try:
          self.replay_process.terminate()
        except:  # pylint: disable=W0702
          pass
      self.replay_process.wait()

  def __enter__(self):
    """Add support for with-statement."""
    self.StartServer()
    return self

  def __exit__(self, unused_exc_type, unused_exc_val, unused_exc_tb):
    """Add support for with-statement."""
    self.StopServer()
<|endoftext|>