<commit_before># coding: utf-8
"""ExposureTask: a class representing an ongoing exposure with the detector"""
import datetime
import enum
import logging
import os
import pickle
import time
from typing import Optional

import h5py
import numpy as np
from PyQt5 import QtCore
from PyQt5.QtCore import pyqtSignal as Signal

from ....dataclasses.exposure import Exposure
from ....dataclasses.header import Header
from ....dataclasses.sample import Sample
from ....devices.detector.pilatus.frontend import PilatusDetector
from ....devices.device.frontend import DeviceFrontend

logger = logging.getLogger(__name__)
logger.setLevel(logging.INFO)

class ExposureState(enum.Enum):
    Initializing = enum.auto()  # the detector is being initialized by the Exposer
    Pending = enum.auto()  # another image is being collected, wait for our turn
    Running = enum.auto()  # this image is being collected
    WaitingForImage = enum.auto()  # collection done, trying to load the image file from disk
    Finished = enum.auto()  # finished successfully, image have been loaded
    TimedOut = enum.auto()  # finished unsuccessfully, image file could not be found
    Stopped = enum.auto()  # finished, stopped on external request

class ExposureTask(QtCore.QObject):
    """Representation of a scheduled/running exposure

    The life of an exposure is the following, reflected by the `status` variable:
        - the exposure task is created in the "Initializing" state. This precedes the issuing of the start exposure
          command to the detector.
        - After the detector has acknowledged the exposure command, the exposure task goes into the "Running" state. If
          multiple exposures are requested, the first one will be "Running", the subsequent ones are "Pending".
        - "Pending" tasks become "Running" when all previous exposures have been done and the waiting times inbetween
          them are also elapsed. In this sense, the i-th exposure (i=1..) spends `exptime`*(i-1) + `expdelay`*(i-1) in
          the "Pending" state and finishes after i*`exptime` + (i-1)*`expdelay`.
        - When the exposure time has elapsed (on the clock of the computer, i.e. no interaction is needed with the
          detector hardware), the state turns to "WaitingForImage". In this state the program tries periodically to load
          the image file from the disk.
        - If the image file is found in the preferred timeout interval, the status gets to "Finished". This is success.
        - If the image file is not found even after the timeout has been elapsed, the status will be "TimedOut".

    The following signals are defined:

        - `exposurestarted`: emitted at the actual start of the exposure (estimated using the clock)
        - `exposureended`: emitted at the actual end of the exposure
        - `finished`: emitted when the image has been loaded / the image loading timeout has been spent

    :ivar prefix: file name prefix, e.g. 'crd', 'tst', 'scn', 'tra' etc.
    :type prefix: str
    :ivar fsn: file sequence number of this exposure
    :type fsn: int
    :ivar command_issue_time: the timestamp (a la time.monotonic()) when the detector was started
    :type command_issue_time: float
    :ivar command_ack_time: the timestamp (a la time.monotonic()) when the detector acknowledged the start command
    :type command_ack_time: float
    :ivar exptime: time of a single exposure (sec)
    :type exptime: float
    :ivar expdelay: delay between two subsequent exposures in a multi-exposure series (sec)
    :type expdelay: float
    :ivar imagetimeout: time before giving up on reading a finished image (sec)
    :type imagetimeout: float
    :ivar index: the 0-based index of this image in a multi-frame exposure
    :type index: int
    :ivar maskoverride: use a different mask
    :type maskoverride: str or None
    """
    # all times are from time.monotonic()
    prefix: str
    fsn: int
    firstfsn: int
    command_ack_time: Optional[float] = None  # when the detector backend process replied to the expose command
    exptime: float  # exposure time
    expdelay: float  # delay between exposures
    imagetimeout: float = 2.0  # timeout for waiting for the image.
    index: int  # 0-based index of this image in a multiple exposure sequence
    maskoverride: Optional[str] = None
    status: ExposureState = ExposureState.Initializing
    finished = Signal(bool, object)
    exposurestarted = Signal()
    exposureended = Signal()
    pendingtimer: Optional[int] = None
    runningtimer: Optional[int] = None
    waitforimagetimer: Optional[int] = None
    imageloadtimer: Optional[int] = None
    imageloadperiod: float = 0.1
    instrument: "Instrument"
    detector: PilatusDetector
    h5: Optional[h5py.File] = None

    def __init__(self, instrument: "Instrument", detector: PilatusDetector, prefix: str, fsn: int, index: int,
                 exptime: float, expdelay: float, maskoverride: Optional[str] = None, writenexus: bool = False):
        super().__init__()
        self.prefix = prefix
        self.fsn = fsn
        self.index = index
        self.exptime = exptime
        self.expdelay = expdelay
        self.exptime = exptime
        self.maskoverride = maskoverride
        self.status = ExposureState.Initializing
        self.instrument = instrument
        self.detector = detector
        self.writenexus = writenexus
        if writenexus:
            targetdir = os.path.join(self.instrument.io.getSubDir('nexus'), prefix)
            os.makedirs(targetdir, exist_ok=True)
            self.h5 = h5py.File(
                os.path.join(targetdir, self.instrument.io.formatFileName(prefix, fsn, '.nxs')), 'w',
                libver='latest')
            grp = self.h5.require_group(f'{prefix}_{fsn:05d}')
            grp.attrs['NX_class'] = 'NXentry'
            self.instrument.toNeXus(grp)

    @property
    def starttime(self):
        """expected start time of this image"""
        return self.command_ack_time + self.index * (self.exptime + self.expdelay)

    @property
    def endtime(self):
        """expected end time of this image"""
        return self.starttime + self.exptime

    def timerEvent(self, event: QtCore.QTimerEvent) -> None:
        if (event.timerId() == self.pendingtimer) and (self.status == ExposureState.Pending):
            # the exposure is supposedly started
            self.onExposureStarted()
            self.killTimer(self.pendingtimer)
            self.pendingtimer = None
        elif (event.timerId() == self.runningtimer) and (self.status == ExposureState.Running):
            # the exposure is supposedly ended
            self.onExposureFinished()
            self.killTimer(self.runningtimer)
            self.runningtimer = None
        elif (event.timerId() == self.waitforimagetimer) and (self.status == ExposureState.WaitingForImage):
            # image waiting timeout spent
            self.killTimer(self.waitforimagetimer)
            self.waitforimagetimer = None
            self.status = ExposureState.TimedOut
            self.finished.emit(False, None)
        elif (event.timerId() == self.imageloadtimer) and (self.status == ExposureState.WaitingForImage):
            # try to load the image
            try:
                image = self.instrument.io.loadCBF(self.prefix, self.fsn, check_local=False)
            except FileNotFoundError:
                # the file is not there (yet?). Wait until timeout
                pass
            else:
                # the file has been found. Kill the timers
                self.killTimer(self.waitforimagetimer)
                self.killTimer(self.imageloadtimer)
                self.waitforimagetimer = None
                self.imageloadtimer = None
                self.status = ExposureState.Finished
                # construct the metadata
                header = self.createHeader()
                self.instrument.io.imageReceived(self.prefix, self.fsn)
                try:
                    mask = self.instrument.io.loadMask(header.maskname)
                except FileNotFoundError:
                    logger.warning(f'Invalid mask file "{header.maskname}". You might have to create a mask yourself.')
                    mask = np.ones_like(image, dtype=np.uint8)
                uncertainty = np.empty_like(image)
                uncertainty[image > 0] = image[image > 0] ** 0.5
                uncertainty[image <= 0] = 1
                self.writeNeXus(image, uncertainty)
                # emit the raw image.
                exposure = Exposure(image, header, uncertainty, mask)
                if self.prefix == self.instrument.config['path']['prefixes']['crd']:
                    self.instrument.datareduction.submit(exposure)
                self.finished.emit(True, exposure)

    def onDetectorExposureStarted(self, starttime):
        """This method should be called when the detector has acknowledged the exposure command.
        In this method we initialize the timers.
        """
        self.command_ack_time = starttime
        assert self.status == ExposureState.Initializing
        if not self.index:
            # this is the first exposure
            self.status = ExposureState.Pending
            self.onExposureStarted()
        else:
            # pending time: exptime * index + expdelay * index
            self.status = ExposureState.Pending
            self.pendingtimer = self.startTimer(
                int(1000 * (self.exptime + self.expdelay) * self.index),
                QtCore.Qt.PreciseTimer)
        self.runningtimer = self.startTimer(
            int(1000 * ((self.exptime + self.expdelay) * self.index + self.exptime)),
            QtCore.Qt.PreciseTimer)
        self.waitforimagetimer = self.startTimer(
            int(1000 * ((self.exptime + self.expdelay) * self.index + self.exptime + self.imagetimeout)),
            QtCore.Qt.PreciseTimer)

    def onExposureStarted(self):
        assert self.status == ExposureState.Pending
        self.status = ExposureState.Running
        self.exposurestarted.emit()

    def onExposureFinished(self):
        assert self.status == ExposureState.Running
        self.imageloadtimer = self.startTimer(int(1000 * self.imageloadperiod), QtCore.Qt.PreciseTimer)
        self.status = ExposureState.WaitingForImage
        self.exposureended.emit()

    def stopExposure(self):
        """Stop the ongoing exposure

        This method does not talk to the detector: its job is to stop all the timers, to avoid waiting
        """
        if self.pendingtimer:
            self.killTimer(self.pendingtimer)
            self.pendingtimer = None
        if self.runningtimer:
            self.killTimer(self.runningtimer)
            self.runningtimer = None
        if self.waitforimagetimer:
            self.killTimer(self.waitforimagetimer)
            self.waitforimagetimer = None
        if self.imageloadtimer:
            self.killTimer(self.imageloadtimer)
            self.imageloadtimer = None
        self.status = ExposureState.Stopped
        self.finished.emit(False, None)

    def createHeader(self) -> Header:
        sample: Optional[Sample] = self.instrument.samplestore.currentSample()
        data = {
            'fsn': self.fsn,
            'filename': os.path.abspath(
                os.path.join(
                    self.instrument.config['path']['directories']['param'], self.prefix,
                    self.instrument.io.formatFileName(self.prefix, self.fsn, '.pickle'))),
            'exposure': {
                'fsn': self.fsn,
                'prefix': self.prefix,
                'exptime': self.exptime,
                'monitor': self.exptime,
                'startdate': datetime.datetime.fromtimestamp(time.time() - time.monotonic() + self.starttime),
                'date': datetime.datetime.now(),
                'enddate': datetime.datetime.now(),
            },
            'geometry': self.instrument.geometry.getHeaderEntry(),
            'sample': sample.todict() if sample is not None else {},
            'motors': self.instrument.motors.getHeaderEntry(),
            'devices': {dev.name: dev.toDict() for dev in self.instrument.devicemanager if dev.isOnline()},
            'environment': {},
            'accounting': {'projectid': self.instrument.projects.project().projectid,
                           'operator': self.instrument.auth.username(),
                           'projectname': self.instrument.projects.project().title,
                           'proposer': self.instrument.projects.project().proposer},
        }
        # devices
        for dev in self.instrument.devicemanager:
            if not dev.isOnline():
                continue
            data['devices'][dev.name]['devicetype'] = dev.devicetype.value
            data['devices'][dev.name]['deviceclass'] = dev.devicename
        # environment
        try:
            vac = self.instrument.devicemanager.vacuum()
            data['environment']['vacuum_pressure'] = vac.pressure()
        except (KeyError, IndexError, DeviceFrontend.DeviceError):
            data['environment']['vacuum_pressure'] = 0.0001
        try:
            temp = self.instrument.devicemanager.temperature()
            data['environment']['temperature'] = temp.temperature()
        except (KeyError, IndexError, DeviceFrontend.DeviceError):
            data['environment']['temperature'] = np.nan
        # sensors
        data['sensors'] = {s.name: s.__getstate__() for s in self.instrument.sensors}
        # adjust truedistance
        if sample is not None:
            data['geometry']['truedistance'] = data['geometry']['dist_sample_det'] - data['sample']['distminus.val']
            data['geometry']['truedistance.err'] = (data['geometry']['dist_sample_det.err'] ** 2 + data['sample'][
                'distminus.err'] ** 2) ** 0.5
        if (sample is not None) and (sample.maskoverride is not None):
            data['geometry']['mask'] = sample.maskoverride
        if self.maskoverride is not None:
            # global mask override takes precedence
            data['geometry']['mask'] = self.maskoverride
        # Save the pickle file
        folder, filename = os.path.split(data['filename'])
        os.makedirs(folder, exist_ok=True)
        with open(data['filename'], 'wb') as f:
            pickle.dump(data, f)
        return Header(datadict=data)

    def writeNeXus(self, img: np.ndarray, unc: np.ndarray):
        if self.h5 is None:
            return
        grp = self.h5.require_group(f'{self.prefix}_{self.fsn:05d}')
        grp.attrs['NX_class'] = 'NXentry'
        # the instrument state has already been written, no need to call self.instrument.toNeXus() again.
        instgroup = grp[[g for g in grp if grp[g].attrs['NX_class'] == 'NXinstrument'][0]]
        detectorgroup: h5py.Group = instgroup[[g for g in instgroup if instgroup[g].attrs['NX_class'] == 'NXdetector'][0]]
        detectorgroup.create_dataset('data', data=img)
        detectorgroup.create_dataset('data_errors', data=unc)
        self.h5.close()
        self.h5 = None
<commit_msg>Bug fixing in NeXus writing<commit_after># coding: utf-8
"""ExposureTask: a class representing an ongoing exposure with the detector"""
import datetime
import enum
import logging
import os
import pickle
import time
from typing import Optional

import h5py
import numpy as np
from PyQt5 import QtCore
from PyQt5.QtCore import pyqtSignal as Signal

from ....dataclasses.exposure import Exposure
from ....dataclasses.header import Header
from ....dataclasses.sample import Sample
from ....devices.detector.pilatus.frontend import PilatusDetector
from ....devices.device.frontend import DeviceFrontend

logger = logging.getLogger(__name__)
logger.setLevel(logging.INFO)

class ExposureState(enum.Enum):
    Initializing = enum.auto()  # the detector is being initialized by the Exposer
    Pending = enum.auto()  # another image is being collected, wait for our turn
    Running = enum.auto()  # this image is being collected
    WaitingForImage = enum.auto()  # collection done, trying to load the image file from disk
    Finished = enum.auto()  # finished successfully, image have been loaded
    TimedOut = enum.auto()  # finished unsuccessfully, image file could not be found
    Stopped = enum.auto()  # finished, stopped on external request

class ExposureTask(QtCore.QObject):
    """Representation of a scheduled/running exposure

    The life of an exposure is the following, reflected by the `status` variable:
        - the exposure task is created in the "Initializing" state. This precedes the issuing of the start exposure
          command to the detector.
        - After the detector has acknowledged the exposure command, the exposure task goes into the "Running" state. If
          multiple exposures are requested, the first one will be "Running", the subsequent ones are "Pending".
        - "Pending" tasks become "Running" when all previous exposures have been done and the waiting times inbetween
          them are also elapsed. In this sense, the i-th exposure (i=1..) spends `exptime`*(i-1) + `expdelay`*(i-1) in
          the "Pending" state and finishes after i*`exptime` + (i-1)*`expdelay`.
        - When the exposure time has elapsed (on the clock of the computer, i.e. no interaction is needed with the
          detector hardware), the state turns to "WaitingForImage". In this state the program tries periodically to load
          the image file from the disk.
        - If the image file is found in the preferred timeout interval, the status gets to "Finished". This is success.
        - If the image file is not found even after the timeout has been elapsed, the status will be "TimedOut".

    The following signals are defined:

        - `exposurestarted`: emitted at the actual start of the exposure (estimated using the clock)
        - `exposureended`: emitted at the actual end of the exposure
        - `finished`: emitted when the image has been loaded / the image loading timeout has been spent

    :ivar prefix: file name prefix, e.g. 'crd', 'tst', 'scn', 'tra' etc.
    :type prefix: str
    :ivar fsn: file sequence number of this exposure
    :type fsn: int
    :ivar command_issue_time: the timestamp (a la time.monotonic()) when the detector was started
    :type command_issue_time: float
    :ivar command_ack_time: the timestamp (a la time.monotonic()) when the detector acknowledged the start command
    :type command_ack_time: float
    :ivar exptime: time of a single exposure (sec)
    :type exptime: float
    :ivar expdelay: delay between two subsequent exposures in a multi-exposure series (sec)
    :type expdelay: float
    :ivar imagetimeout: time before giving up on reading a finished image (sec)
    :type imagetimeout: float
    :ivar index: the 0-based index of this image in a multi-frame exposure
    :type index: int
    :ivar maskoverride: use a different mask
    :type maskoverride: str or None
    """
    # all times are from time.monotonic()
    prefix: str
    fsn: int
    firstfsn: int
    command_ack_time: Optional[float] = None  # when the detector backend process replied to the expose command
    exptime: float  # exposure time
    expdelay: float  # delay between exposures
    imagetimeout: float = 2.0  # timeout for waiting for the image.
    index: int  # 0-based index of this image in a multiple exposure sequence
    maskoverride: Optional[str] = None
    status: ExposureState = ExposureState.Initializing
    finished = Signal(bool, object)
    exposurestarted = Signal()
    exposureended = Signal()
    pendingtimer: Optional[int] = None
    runningtimer: Optional[int] = None
    waitforimagetimer: Optional[int] = None
    imageloadtimer: Optional[int] = None
    imageloadperiod: float = 0.1
    instrument: "Instrument"
    detector: PilatusDetector
    h5: Optional[h5py.File] = None

    def __init__(self, instrument: "Instrument", detector: PilatusDetector, prefix: str, fsn: int, index: int,
                 exptime: float, expdelay: float, maskoverride: Optional[str] = None, writenexus: bool = False):
        super().__init__()
        self.prefix = prefix
        self.fsn = fsn
        self.index = index
        self.exptime = exptime
        self.expdelay = expdelay
        self.exptime = exptime
        self.maskoverride = maskoverride
        self.status = ExposureState.Initializing
        self.instrument = instrument
        self.detector = detector
        self.writenexus = writenexus
        if writenexus:
            targetdir = os.path.join(self.instrument.io.getSubDir('nexus'), prefix)
            os.makedirs(targetdir, exist_ok=True)
            self.h5 = h5py.File(
                os.path.join(targetdir, self.instrument.io.formatFileName(prefix, fsn, '.nxs')), 'w',
                libver='latest')
            grp = self.h5.require_group(f'{prefix}_{fsn:05d}')
            grp.attrs['NX_class'] = 'NXentry'
            self.instrument.toNeXus(grp)

    @property
    def starttime(self):
        """expected start time of this image"""
        return self.command_ack_time + self.index * (self.exptime + self.expdelay)

    @property
    def endtime(self):
        """expected end time of this image"""
        return self.starttime + self.exptime

    def timerEvent(self, event: QtCore.QTimerEvent) -> None:
        if (event.timerId() == self.pendingtimer) and (self.status == ExposureState.Pending):
            # the exposure is supposedly started
            self.onExposureStarted()
            self.killTimer(self.pendingtimer)
            self.pendingtimer = None
        elif (event.timerId() == self.runningtimer) and (self.status == ExposureState.Running):
            # the exposure is supposedly ended
            self.onExposureFinished()
            self.killTimer(self.runningtimer)
            self.runningtimer = None
        elif (event.timerId() == self.waitforimagetimer) and (self.status == ExposureState.WaitingForImage):
            # image waiting timeout spent
            self.killTimer(self.waitforimagetimer)
            self.waitforimagetimer = None
            self.status = ExposureState.TimedOut
            self.finished.emit(False, None)
        elif (event.timerId() == self.imageloadtimer) and (self.status == ExposureState.WaitingForImage):
            # try to load the image
            try:
                image = self.instrument.io.loadCBF(self.prefix, self.fsn, check_local=False)
            except FileNotFoundError:
                # the file is not there (yet?). Wait until timeout
                pass
            else:
                # the file has been found. Kill the timers
                self.killTimer(self.waitforimagetimer)
                self.killTimer(self.imageloadtimer)
                self.waitforimagetimer = None
                self.imageloadtimer = None
                self.status = ExposureState.Finished
                # construct the metadata
                header = self.createHeader()
                self.instrument.io.imageReceived(self.prefix, self.fsn)
                try:
                    mask = self.instrument.io.loadMask(header.maskname)
                except FileNotFoundError:
                    logger.warning(f'Invalid mask file "{header.maskname}". You might have to create a mask yourself.')
                    mask = np.ones_like(image, dtype=np.uint8)
                uncertainty = np.empty_like(image)
                uncertainty[image > 0] = image[image > 0] ** 0.5
                uncertainty[image <= 0] = 1
                self.writeNeXus(image, uncertainty)
                # emit the raw image.
                exposure = Exposure(image, header, uncertainty, mask)
                if self.prefix == self.instrument.config['path']['prefixes']['crd']:
                    self.instrument.datareduction.submit(exposure)
                self.finished.emit(True, exposure)

    def onDetectorExposureStarted(self, starttime):
        """This method should be called when the detector has acknowledged the exposure command.
        In this method we initialize the timers.
        """
        self.command_ack_time = starttime
        assert self.status == ExposureState.Initializing
        if not self.index:
            # this is the first exposure
            self.status = ExposureState.Pending
            self.onExposureStarted()
        else:
            # pending time: exptime * index + expdelay * index
            self.status = ExposureState.Pending
            self.pendingtimer = self.startTimer(
                int(1000 * (self.exptime + self.expdelay) * self.index),
                QtCore.Qt.PreciseTimer)
        self.runningtimer = self.startTimer(
            int(1000 * ((self.exptime + self.expdelay) * self.index + self.exptime)),
            QtCore.Qt.PreciseTimer)
        self.waitforimagetimer = self.startTimer(
            int(1000 * ((self.exptime + self.expdelay) * self.index + self.exptime + self.imagetimeout)),
            QtCore.Qt.PreciseTimer)

    def onExposureStarted(self):
        assert self.status == ExposureState.Pending
        self.status = ExposureState.Running
        self.exposurestarted.emit()

    def onExposureFinished(self):
        assert self.status == ExposureState.Running
        self.imageloadtimer = self.startTimer(int(1000 * self.imageloadperiod), QtCore.Qt.PreciseTimer)
        self.status = ExposureState.WaitingForImage
        self.exposureended.emit()

    def stopExposure(self):
        """Stop the ongoing exposure

        This method does not talk to the detector: its job is to stop all the timers, to avoid waiting
        """
        if self.pendingtimer:
            self.killTimer(self.pendingtimer)
            self.pendingtimer = None
        if self.runningtimer:
            self.killTimer(self.runningtimer)
            self.runningtimer = None
        if self.waitforimagetimer:
            self.killTimer(self.waitforimagetimer)
            self.waitforimagetimer = None
        if self.imageloadtimer:
            self.killTimer(self.imageloadtimer)
            self.imageloadtimer = None
        self.status = ExposureState.Stopped
        self.finished.emit(False, None)

    def createHeader(self) -> Header:
        sample: Optional[Sample] = self.instrument.samplestore.currentSample()
        data = {
            'fsn': self.fsn,
            'filename': os.path.abspath(
                os.path.join(
                    self.instrument.config['path']['directories']['param'], self.prefix,
                    self.instrument.io.formatFileName(self.prefix, self.fsn, '.pickle'))),
            'exposure': {
                'fsn': self.fsn,
                'prefix': self.prefix,
                'exptime': self.exptime,
                'monitor': self.exptime,
                'startdate': datetime.datetime.fromtimestamp(time.time() - time.monotonic() + self.starttime),
                'date': datetime.datetime.now(),
                'enddate': datetime.datetime.now(),
            },
            'geometry': self.instrument.geometry.getHeaderEntry(),
            'sample': sample.todict() if sample is not None else {},
            'motors': self.instrument.motors.getHeaderEntry(),
            'devices': {dev.name: dev.toDict() for dev in self.instrument.devicemanager if dev.isOnline()},
            'environment': {},
            'accounting': {'projectid': self.instrument.projects.project().projectid,
                           'operator': self.instrument.auth.username(),
                           'projectname': self.instrument.projects.project().title,
                           'proposer': self.instrument.projects.project().proposer},
        }
        # devices
        for dev in self.instrument.devicemanager:
            if not dev.isOnline():
                continue
            data['devices'][dev.name]['devicetype'] = dev.devicetype.value
            data['devices'][dev.name]['deviceclass'] = dev.devicename
        # environment
        try:
            vac = self.instrument.devicemanager.vacuum()
            data['environment']['vacuum_pressure'] = vac.pressure()
        except (KeyError, IndexError, DeviceFrontend.DeviceError):
            data['environment']['vacuum_pressure'] = 0.0001
        try:
            temp = self.instrument.devicemanager.temperature()
            data['environment']['temperature'] = temp.temperature()
        except (KeyError, IndexError, DeviceFrontend.DeviceError):
            data['environment']['temperature'] = np.nan
        # sensors
        data['sensors'] = {s.name: s.__getstate__() for s in self.instrument.sensors}
        # adjust truedistance
        if sample is not None:
            data['geometry']['truedistance'] = data['geometry']['dist_sample_det'] - data['sample']['distminus.val']
            data['geometry']['truedistance.err'] = (data['geometry']['dist_sample_det.err'] ** 2 + data['sample'][
                'distminus.err'] ** 2) ** 0.5
        if (sample is not None) and (sample.maskoverride is not None):
            data['geometry']['mask'] = sample.maskoverride
        if self.maskoverride is not None:
            # global mask override takes precedence
            data['geometry']['mask'] = self.maskoverride
        # Save the pickle file
        folder, filename = os.path.split(data['filename'])
        os.makedirs(folder, exist_ok=True)
        with open(data['filename'], 'wb') as f:
            pickle.dump(data, f)
        return Header(datadict=data)

    def writeNeXus(self, img: np.ndarray, unc: np.ndarray):
        if self.h5 is None:
            return
        grp = self.h5.require_group(f'{self.prefix}_{self.fsn:05d}')
        grp.attrs['NX_class'] = 'NXentry'
        # the instrument state has already been written, no need to call self.instrument.toNeXus() again.
        instgroup = grp[[g for g in grp if ('NX_class' in grp[g].attrs) and (grp[g].attrs['NX_class'] == 'NXinstrument')][0]]
        detectorgroup: h5py.Group = instgroup[[g for g in instgroup if ('NX_class' in instgroup[g].attrs) and (instgroup[g].attrs['NX_class'] == 'NXdetector')][0]]
        detectorgroup.create_dataset('data', data=img)
        detectorgroup.create_dataset('data_errors', data=unc)
        self.h5.close()
        self.h5 = None
<|endoftext|>


<commit_before>from locust import HttpLocust, TaskSet, task
from bs4 import BeautifulSoup
import random
import json
lat = 47.614848
lng = -122.3359059

class UserBehavior(TaskSet):
    def on_start(self):
        pass

    @task(1)
    def index(self):
        self.client.get("/")

    # @task(1)
    # def create(self):
    #     print "create"
    #     self.response = self.client.get("/create")

    @task(10)
    class SubTaskCreate(TaskSet):
        def on_start(self):
            self.response = self.client.get("/create")
            self.lat = lat
            self.lng = lng
            self.response = self.client.post(
                url="/scene_options_ajax",
                json="True",
                data={'lat': self.lat, 'lng': self.lng}
                )

        @task(5)
        def move_map(self):
            self.lat = random.uniform(-1, 1) + self.lat
            self.lng = random.uniform(-1, 1) + self.lng
            print self.lat, self.lng
            self.response = self.client.post(
                url="/scene_options_ajax",
                json="True",
                data={'lat': self.lat, 'lng': self.lng}
                )

        @task(1)
        def click_link(self):
            json_data = json.loads(self.response.text)
            random_num = random.randint(0, len(json_data['scenes_date']))
            random_url = json_data["scenes_date"][random_num]["download_url"]
            scene_id = json_data["scenes_date"][random_num]["entityid"]
            print random_url

            @task(1)
            class scene_page(self):
                def on_start(self):
                    self.response = self.client.get(random_url)

                def preview(self):
                    self.response = self.client.post(
                                        url="request_p/{}".format(scene_id),
                                        data={'band_combo': "432"}
                                        )
                    soup = BeautifulSoup(self.response.text)

        # @task(1)
        # def select_scene(self):
        #     # Get url
        #     soup = BeautifulSoup(self.client.get(""))

        #     self.client.get()

        #     @task
        #     def render_preview(self):
        #         self.client.get()

        #     @task
        #     def render_full(self):
        #         self.client.get()

class WebsiteUser(HttpLocust):
    task_set = UserBehavior
    min_wait = 1000
    max_wait = 5000
<commit_msg>scene_page inherits TaskSet. Add task decorator to preview method.<commit_after>from locust import HttpLocust, TaskSet, task
from bs4 import BeautifulSoup
import random
import json
lat = 47.614848
lng = -122.3359059

class UserBehavior(TaskSet):
    def on_start(self):
        pass

    @task(1)
    def index(self):
        self.client.get("/")

    # @task(1)
    # def create(self):
    #     print "create"
    #     self.response = self.client.get("/create")

    @task(10)
    class SubTaskCreate(TaskSet):
        def on_start(self):
            self.response = self.client.get("/create")
            self.lat = lat
            self.lng = lng
            self.response = self.client.post(
                url="/scene_options_ajax",
                json="True",
                data={'lat': self.lat, 'lng': self.lng}
                )

        @task(5)
        def move_map(self):
            self.lat = random.uniform(-1, 1) + self.lat
            self.lng = random.uniform(-1, 1) + self.lng
            print self.lat, self.lng
            self.response = self.client.post(
                url="/scene_options_ajax",
                json="True",
                data={'lat': self.lat, 'lng': self.lng}
                )

        @task(1)
        def click_link(self):
            json_data = json.loads(self.response.text)
            random_num = random.randint(0, len(json_data['scenes_date']))
            random_url = json_data["scenes_date"][random_num]["download_url"]
            scene_id = json_data["scenes_date"][random_num]["entityid"]
            print random_url

            @task(1)
            class scene_page(TaskSet):
                def on_start(self):
                    self.response = self.client.get(random_url)

                @task
                def preview(self):
                    self.response = self.client.post(
                                        url="request_p/{}".format(scene_id),
                                        data={'band_combo': "432"}
                                        )
                    # soup = BeautifulSoup(self.response.text)

        # @task(1)
        # def select_scene(self):
        #     # Get url
        #     soup = BeautifulSoup(self.client.get(""))

        #     self.client.get()

        #     @task
        #     def render_preview(self):
        #         self.client.get()

        #     @task
        #     def render_full(self):
        #         self.client.get()

class WebsiteUser(HttpLocust):
    task_set = UserBehavior
    min_wait = 1000
    max_wait = 5000
<|endoftext|>


<commit_before>"""
Template tags for working with lists of model instances which represent
trees.
"""
from django import template
from django.db.models import get_model
from django.db.models.fields import FieldDoesNotExist
from django.utils.encoding import force_unicode
from django.utils.safestring import mark_safe
from django.utils.translation import ugettext as _

from mptt.utils import tree_item_iterator, drilldown_tree_for_node

register = template.Library()

### ITERATIVE TAGS

class FullTreeForModelNode(template.Node):
    def __init__(self, model, context_var):
        self.model = model
        self.context_var = context_var

    def render(self, context):
        cls = get_model(*self.model.split('.'))
        if cls is None:
            raise template.TemplateSyntaxError(
                _('full_tree_for_model tag was given an invalid model: %s') % self.model
            )
        context[self.context_var] = cls._tree_manager.all()
        return ''

class DrilldownTreeForNodeNode(template.Node):
    def __init__(self, node, context_var, foreign_key=None, count_attr=None,
                 cumulative=False):
        self.node = template.Variable(node)
        self.context_var = context_var
        self.foreign_key = foreign_key
        self.count_attr = count_attr
        self.cumulative = cumulative

    def render(self, context):
        # Let any VariableDoesNotExist raised bubble up
        args = [self.node.resolve(context)]

        if self.foreign_key is not None:
            app_label, model_name, fk_attr = self.foreign_key.split('.')
            cls = get_model(app_label, model_name)
            if cls is None:
                raise template.TemplateSyntaxError(
                    _('drilldown_tree_for_node tag was given an invalid model: %s') % \
                    '.'.join([app_label, model_name])
                )
            try:
                cls._meta.get_field(fk_attr)
            except FieldDoesNotExist:
                raise template.TemplateSyntaxError(
                    _('drilldown_tree_for_node tag was given an invalid model field: %s') % fk_attr
                )
            args.extend([cls, fk_attr, self.count_attr, self.cumulative])

        context[self.context_var] = drilldown_tree_for_node(*args)
        return ''

@register.tag
def full_tree_for_model(parser, token):
    """
    Populates a template variable with a ``QuerySet`` containing the
    full tree for a given model.

    Usage::

       {% full_tree_for_model [model] as [varname] %}

    The model is specified in ``[appname].[modelname]`` format.

    Example::

       {% full_tree_for_model tests.Genre as genres %}

    """
    bits = token.contents.split()
    if len(bits) != 4:
        raise template.TemplateSyntaxError(_('%s tag requires three arguments') % bits[0])
    if bits[2] != 'as':
        raise template.TemplateSyntaxError(_("second argument to %s tag must be 'as'") % bits[0])
    return FullTreeForModelNode(bits[1], bits[3])

@register.tag('drilldown_tree_for_node')
def do_drilldown_tree_for_node(parser, token):
    """
    Populates a template variable with the drilldown tree for a given
    node, optionally counting the number of items associated with its
    children.

    A drilldown tree consists of a node's ancestors, itself and its
    immediate children. For example, a drilldown tree for a book
    category "Personal Finance" might look something like::

       Books
          Business, Finance & Law
             Personal Finance
                Budgeting (220)
                Financial Planning (670)

    Usage::

       {% drilldown_tree_for_node [node] as [varname] %}

    Extended usage::

       {% drilldown_tree_for_node [node] as [varname] count [foreign_key] in [count_attr] %}
       {% drilldown_tree_for_node [node] as [varname] cumulative count [foreign_key] in [count_attr] %}

    The foreign key is specified in ``[appname].[modelname].[fieldname]``
    format, where ``fieldname`` is the name of a field in the specified
    model which relates it to the given node's model.

    When this form is used, a ``count_attr`` attribute on each child of
    the given node in the drilldown tree will contain a count of the
    number of items associated with it through the given foreign key.

    If cumulative is also specified, this count will be for items
    related to the child node and all of its descendants.

    Examples::

       {% drilldown_tree_for_node genre as drilldown %}
       {% drilldown_tree_for_node genre as drilldown count tests.Game.genre in game_count %}
       {% drilldown_tree_for_node genre as drilldown cumulative count tests.Game.genre in game_count %}

    """
    bits = token.contents.split()
    len_bits = len(bits)
    if len_bits not in (4, 8, 9):
        raise template.TemplateSyntaxError(
            _('%s tag requires either three, seven or eight arguments') % bits[0])
    if bits[2] != 'as':
        raise template.TemplateSyntaxError(
            _("second argument to %s tag must be 'as'") % bits[0])
    if len_bits == 8:
        if bits[4] != 'count':
            raise template.TemplateSyntaxError(
                _("if seven arguments are given, fourth argument to %s tag must be 'with'") % bits[0])
        if bits[6] != 'in':
            raise template.TemplateSyntaxError(
                _("if seven arguments are given, sixth argument to %s tag must be 'in'") % bits[0])
        return DrilldownTreeForNodeNode(bits[1], bits[3], bits[5], bits[7])
    elif len_bits == 9:
        if bits[4] != 'cumulative':
            raise template.TemplateSyntaxError(
                _("if eight arguments are given, fourth argument to %s tag must be 'cumulative'") % bits[0])
        if bits[5] != 'count':
            raise template.TemplateSyntaxError(
                _("if eight arguments are given, fifth argument to %s tag must be 'count'") % bits[0])
        if bits[7] != 'in':
            raise template.TemplateSyntaxError(
                _("if eight arguments are given, seventh argument to %s tag must be 'in'") % bits[0])
        return DrilldownTreeForNodeNode(bits[1], bits[3], bits[6], bits[8], cumulative=True)
    else:
        return DrilldownTreeForNodeNode(bits[1], bits[3])

@register.filter
def tree_info(items, features=None):
    """
    Given a list of tree items, produces doubles of a tree item and a
    ``dict`` containing information about the tree structure around the
    item, with the following contents:

       new_level
          ``True`` if the current item is the start of a new level in
          the tree, ``False`` otherwise.

       closed_levels
          A list of levels which end after the current item. This will
          be an empty list if the next item is at the same level as the
          current item.

    Using this filter with unpacking in a ``{% for %}`` tag, you should
    have enough information about the tree structure to create a
    hierarchical representation of the tree.

    Example::

       {% for genre,structure in genres|tree_info %}
       {% if tree.new_level %}

{% else %}

{% endif %}
       {{ genre.name }}
       {% for level in tree.closed_levels %}

{% endfor %}
       {% endfor %}

    """
    kwargs = {}
    if features:
        feature_names = features.split(',')
        if 'ancestors' in feature_names:
            kwargs['ancestors'] = True
    return tree_item_iterator(items, **kwargs)

@register.filter
def tree_path(items, separator=' :: '):
    """
    Creates a tree path represented by a list of ``items`` by joining
    the items with a ``separator``.

    Each path item will be coerced to unicode, so a list of model
    instances may be given if required.

    Example::

       {{ some_list|tree_path }}
       {{ some_node.get_ancestors|tree_path:" > " }}

    """
    return separator.join([force_unicode(i) for i in items])

### RECURSIVE TAGS

@register.filter
def cache_tree_children(queryset):
    """
    Takes a list/queryset of model objects in MPTT left (depth-first) order,
    caches the children on each node, as well as the parent of each child node,
    allowing up and down traversal through the tree without the need for
    further queries. This makes it possible to have a recursively included
    template without worrying about database queries.

    Returns a list of top-level nodes. If a single tree was provided in its
    entirety, the list will of course consist of just the tree's root node.

    """

    current_path = []
    top_nodes = []

    # If ``queryset`` is QuerySet-like, set ordering to depth-first
    if hasattr(queryset, 'order_by'):
        mptt_opts = queryset.model._mptt_meta
        tree_id_attr = mptt_opts.tree_id_attr
        left_attr = mptt_opts.left_attr
        queryset = queryset.order_by(tree_id_attr, left_attr)

    if queryset:
        root_level = None
        parent_attr = None
        for obj in queryset:
            # First iteration, so set the parent attribute name
            if parent_attr is None:
                parent_attr = obj._mptt_meta.parent_attr

            # Get the current mptt node level
            node_level = obj.get_level()

            if root_level is None:
                # First iteration, so set the root level to the top node level
                root_level = node_level

            if node_level < root_level:
                # ``queryset`` was a list or other iterable (unable to order),
                # and was provided in an order other than depth-first
                raise ValueError(
                    _(u'Node %s not in depth-first order') % (type(queryset),)
                )

            # Set up the attribute on the node that will store cached children,
            # which is used by ``MPTTModel.get_children``
            obj._cached_children = []

            # Remove nodes not in the current branch
            while len(current_path) > node_level - root_level:
                current_path.pop(-1)

            if node_level == root_level:
                # Add the root to the list of top nodes, which will be returned
                top_nodes.append(obj)
            else:
                # Cache the parent on the current node, and attach the current
                # node to the parent's list of children
                _parent = current_path[-1]
                setattr(obj, parent_attr, _parent)
                _parent._cached_children.append(obj)

            # Add the current node to end of the current path - the last node
            # in the current path is the parent for the next iteration, unless
            # the next iteration is higher up the tree (a new branch), in which
            # case the paths below it (e.g., this one) will be removed from the
            # current path during the next iteration
            current_path.append(obj)

    return top_nodes

class RecurseTreeNode(template.Node):
    def __init__(self, template_nodes, queryset_var):
        self.template_nodes = template_nodes
        self.queryset_var = queryset_var

    def _render_node(self, context, node):
        bits = []
        context.push()
        for child in node.get_children():
            bits.append(self._render_node(context, child))
        context['node'] = node
        context['children'] = mark_safe(u''.join(bits))
        rendered = self.template_nodes.render(context)
        context.pop()
        return rendered

    def render(self, context):
        queryset = self.queryset_var.resolve(context)
        roots = cache_tree_children(queryset)
        bits = [self._render_node(context, node) for node in roots]
        return ''.join(bits)

@register.tag
def recursetree(parser, token):
    """
    Iterates over the nodes in the tree, and renders the contained block for each node.
    This tag will recursively render children into the template variable {{ children }}.
    Only one database query is required (children are cached for the whole tree)

    Usage:

                {% recursetree nodes %}

                        {{ node.name }}
                        {% if not node.is_leaf_node %}

                                {{ children }}

                        {% endif %}

                {% endrecursetree %}

    """
    bits = token.contents.split()
    if len(bits) != 2:
        raise template.TemplateSyntaxError(_('%s tag requires a queryset') % bits[0])

    queryset_var = template.Variable(bits[1])

    template_nodes = parser.parse(('endrecursetree',))
    parser.delete_first_token()

    return RecurseTreeNode(template_nodes, queryset_var)
<commit_msg>Minor optimization<commit_after>"""
Template tags for working with lists of model instances which represent
trees.
"""
from django import template
from django.db.models import get_model
from django.db.models.fields import FieldDoesNotExist
from django.utils.encoding import force_unicode
from django.utils.safestring import mark_safe
from django.utils.translation import ugettext as _

from mptt.utils import tree_item_iterator, drilldown_tree_for_node

register = template.Library()

### ITERATIVE TAGS

class FullTreeForModelNode(template.Node):
    def __init__(self, model, context_var):
        self.model = model
        self.context_var = context_var

    def render(self, context):
        cls = get_model(*self.model.split('.'))
        if cls is None:
            raise template.TemplateSyntaxError(
                _('full_tree_for_model tag was given an invalid model: %s') % self.model
            )
        context[self.context_var] = cls._tree_manager.all()
        return ''

class DrilldownTreeForNodeNode(template.Node):
    def __init__(self, node, context_var, foreign_key=None, count_attr=None,
                 cumulative=False):
        self.node = template.Variable(node)
        self.context_var = context_var
        self.foreign_key = foreign_key
        self.count_attr = count_attr
        self.cumulative = cumulative

    def render(self, context):
        # Let any VariableDoesNotExist raised bubble up
        args = [self.node.resolve(context)]

        if self.foreign_key is not None:
            app_label, model_name, fk_attr = self.foreign_key.split('.')
            cls = get_model(app_label, model_name)
            if cls is None:
                raise template.TemplateSyntaxError(
                    _('drilldown_tree_for_node tag was given an invalid model: %s') % \
                    '.'.join([app_label, model_name])
                )
            try:
                cls._meta.get_field(fk_attr)
            except FieldDoesNotExist:
                raise template.TemplateSyntaxError(
                    _('drilldown_tree_for_node tag was given an invalid model field: %s') % fk_attr
                )
            args.extend([cls, fk_attr, self.count_attr, self.cumulative])

        context[self.context_var] = drilldown_tree_for_node(*args)
        return ''

@register.tag
def full_tree_for_model(parser, token):
    """
    Populates a template variable with a ``QuerySet`` containing the
    full tree for a given model.

    Usage::

       {% full_tree_for_model [model] as [varname] %}

    The model is specified in ``[appname].[modelname]`` format.

    Example::

       {% full_tree_for_model tests.Genre as genres %}

    """
    bits = token.contents.split()
    if len(bits) != 4:
        raise template.TemplateSyntaxError(_('%s tag requires three arguments') % bits[0])
    if bits[2] != 'as':
        raise template.TemplateSyntaxError(_("second argument to %s tag must be 'as'") % bits[0])
    return FullTreeForModelNode(bits[1], bits[3])

@register.tag('drilldown_tree_for_node')
def do_drilldown_tree_for_node(parser, token):
    """
    Populates a template variable with the drilldown tree for a given
    node, optionally counting the number of items associated with its
    children.

    A drilldown tree consists of a node's ancestors, itself and its
    immediate children. For example, a drilldown tree for a book
    category "Personal Finance" might look something like::

       Books
          Business, Finance & Law
             Personal Finance
                Budgeting (220)
                Financial Planning (670)

    Usage::

       {% drilldown_tree_for_node [node] as [varname] %}

    Extended usage::

       {% drilldown_tree_for_node [node] as [varname] count [foreign_key] in [count_attr] %}
       {% drilldown_tree_for_node [node] as [varname] cumulative count [foreign_key] in [count_attr] %}

    The foreign key is specified in ``[appname].[modelname].[fieldname]``
    format, where ``fieldname`` is the name of a field in the specified
    model which relates it to the given node's model.

    When this form is used, a ``count_attr`` attribute on each child of
    the given node in the drilldown tree will contain a count of the
    number of items associated with it through the given foreign key.

    If cumulative is also specified, this count will be for items
    related to the child node and all of its descendants.

    Examples::

       {% drilldown_tree_for_node genre as drilldown %}
       {% drilldown_tree_for_node genre as drilldown count tests.Game.genre in game_count %}
       {% drilldown_tree_for_node genre as drilldown cumulative count tests.Game.genre in game_count %}

    """
    bits = token.contents.split()
    len_bits = len(bits)
    if len_bits not in (4, 8, 9):
        raise template.TemplateSyntaxError(
            _('%s tag requires either three, seven or eight arguments') % bits[0])
    if bits[2] != 'as':
        raise template.TemplateSyntaxError(
            _("second argument to %s tag must be 'as'") % bits[0])
    if len_bits == 8:
        if bits[4] != 'count':
            raise template.TemplateSyntaxError(
                _("if seven arguments are given, fourth argument to %s tag must be 'with'") % bits[0])
        if bits[6] != 'in':
            raise template.TemplateSyntaxError(
                _("if seven arguments are given, sixth argument to %s tag must be 'in'") % bits[0])
        return DrilldownTreeForNodeNode(bits[1], bits[3], bits[5], bits[7])
    elif len_bits == 9:
        if bits[4] != 'cumulative':
            raise template.TemplateSyntaxError(
                _("if eight arguments are given, fourth argument to %s tag must be 'cumulative'") % bits[0])
        if bits[5] != 'count':
            raise template.TemplateSyntaxError(
                _("if eight arguments are given, fifth argument to %s tag must be 'count'") % bits[0])
        if bits[7] != 'in':
            raise template.TemplateSyntaxError(
                _("if eight arguments are given, seventh argument to %s tag must be 'in'") % bits[0])
        return DrilldownTreeForNodeNode(bits[1], bits[3], bits[6], bits[8], cumulative=True)
    else:
        return DrilldownTreeForNodeNode(bits[1], bits[3])

@register.filter
def tree_info(items, features=None):
    """
    Given a list of tree items, produces doubles of a tree item and a
    ``dict`` containing information about the tree structure around the
    item, with the following contents:

       new_level
          ``True`` if the current item is the start of a new level in
          the tree, ``False`` otherwise.

       closed_levels
          A list of levels which end after the current item. This will
          be an empty list if the next item is at the same level as the
          current item.

    Using this filter with unpacking in a ``{% for %}`` tag, you should
    have enough information about the tree structure to create a
    hierarchical representation of the tree.

    Example::

       {% for genre,structure in genres|tree_info %}
       {% if tree.new_level %}

{% else %}

{% endif %}
       {{ genre.name }}
       {% for level in tree.closed_levels %}

{% endfor %}
       {% endfor %}

    """
    kwargs = {}
    if features:
        feature_names = features.split(',')
        if 'ancestors' in feature_names:
            kwargs['ancestors'] = True
    return tree_item_iterator(items, **kwargs)

@register.filter
def tree_path(items, separator=' :: '):
    """
    Creates a tree path represented by a list of ``items`` by joining
    the items with a ``separator``.

    Each path item will be coerced to unicode, so a list of model
    instances may be given if required.

    Example::

       {{ some_list|tree_path }}
       {{ some_node.get_ancestors|tree_path:" > " }}

    """
    return separator.join([force_unicode(i) for i in items])

### RECURSIVE TAGS

@register.filter
def cache_tree_children(queryset):
    """
    Takes a list/queryset of model objects in MPTT left (depth-first) order,
    caches the children on each node, as well as the parent of each child node,
    allowing up and down traversal through the tree without the need for
    further queries. This makes it possible to have a recursively included
    template without worrying about database queries.

    Returns a list of top-level nodes. If a single tree was provided in its
    entirety, the list will of course consist of just the tree's root node.

    """

    current_path = []
    top_nodes = []

    # If ``queryset`` is QuerySet-like, set ordering to depth-first
    if hasattr(queryset, 'order_by'):
        mptt_opts = queryset.model._mptt_meta
        tree_id_attr = mptt_opts.tree_id_attr
        left_attr = mptt_opts.left_attr
        queryset = queryset.order_by(tree_id_attr, left_attr)

    if queryset:
        # Get the model's parent-attribute name
        parent_attr = queryset[0]._mptt_meta.parent_attr
        root_level = None
        for obj in queryset:
            # Get the current mptt node level
            node_level = obj.get_level()

            if root_level is None:
                # First iteration, so set the root level to the top node level
                root_level = node_level

            if node_level < root_level:
                # ``queryset`` was a list or other iterable (unable to order),
                # and was provided in an order other than depth-first
                raise ValueError(
                    _(u'Node %s not in depth-first order') % (type(queryset),)
                )

            # Set up the attribute on the node that will store cached children,
            # which is used by ``MPTTModel.get_children``
            obj._cached_children = []

            # Remove nodes not in the current branch
            while len(current_path) > node_level - root_level:
                current_path.pop(-1)

            if node_level == root_level:
                # Add the root to the list of top nodes, which will be returned
                top_nodes.append(obj)
            else:
                # Cache the parent on the current node, and attach the current
                # node to the parent's list of children
                _parent = current_path[-1]
                setattr(obj, parent_attr, _parent)
                _parent._cached_children.append(obj)

            # Add the current node to end of the current path - the last node
            # in the current path is the parent for the next iteration, unless
            # the next iteration is higher up the tree (a new branch), in which
            # case the paths below it (e.g., this one) will be removed from the
            # current path during the next iteration
            current_path.append(obj)

    return top_nodes

class RecurseTreeNode(template.Node):
    def __init__(self, template_nodes, queryset_var):
        self.template_nodes = template_nodes
        self.queryset_var = queryset_var

    def _render_node(self, context, node):
        bits = []
        context.push()
        for child in node.get_children():
            bits.append(self._render_node(context, child))
        context['node'] = node
        context['children'] = mark_safe(u''.join(bits))
        rendered = self.template_nodes.render(context)
        context.pop()
        return rendered

    def render(self, context):
        queryset = self.queryset_var.resolve(context)
        roots = cache_tree_children(queryset)
        bits = [self._render_node(context, node) for node in roots]
        return ''.join(bits)

@register.tag
def recursetree(parser, token):
    """
    Iterates over the nodes in the tree, and renders the contained block for each node.
    This tag will recursively render children into the template variable {{ children }}.
    Only one database query is required (children are cached for the whole tree)

    Usage:

                {% recursetree nodes %}

                        {{ node.name }}
                        {% if not node.is_leaf_node %}

                                {{ children }}

                        {% endif %}

                {% endrecursetree %}

    """
    bits = token.contents.split()
    if len(bits) != 2:
        raise template.TemplateSyntaxError(_('%s tag requires a queryset') % bits[0])

    queryset_var = template.Variable(bits[1])

    template_nodes = parser.parse(('endrecursetree',))
    parser.delete_first_token()

    return RecurseTreeNode(template_nodes, queryset_var)
<|endoftext|>


<commit_before># Copyright 2015-2015 Amazon.com, Inc. or its affiliates. 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. A copy of the License is located at
#
#     http://aws.amazon.com/apache2.0/
#
# or in the "license" file accompanying this file. This file 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.

import os
import string
import random
from locust import HttpLocust, TaskSet, task

class MyTaskSet(TaskSet):
    @task(1000)
    def index(self):
        response = self.client.get("/")

    @task(3)
    def discover(self):
        self.client.get("/discover/")

    @task(2)
    def projects(self):
        self.client.get("/projects/unitarian-church-of-staten-island/")

    @task(1)
    def blog(self):
        self.client.get("/blog/")
        self.client.get("/welcome-to-the-new-blocpower/")

    @task(1)
    def about(self):
        self.client.get("/about-us/")

    # This task will 15 times for every 1000 runs of the above task
    # @task(15)
    # def about(self):
    #     self.client.get("/blog")

    # This task will run once for every 1000 runs of the above task
    # @task(1)
    # def about(self):
    #     id = id_generator()
    #     self.client.post("/signup", {"email": "example@example.com", "name": "Test"})

class MyLocust(HttpLocust):
    host = os.getenv('TARGET_URL', "http://loadtest.marketplace.blocpower.us")
    task_set = MyTaskSet
    min_wait = 45
    max_wait = 50
<commit_msg>adding<commit_after># Copyright 2015-2015 Amazon.com, Inc. or its affiliates. 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. A copy of the License is located at
#
#     http://aws.amazon.com/apache2.0/
#
# or in the "license" file accompanying this file. This file 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.

import os
import string
import random
from locust import HttpLocust, TaskSet, task

class MyTaskSet(TaskSet):
    @task(1000)
    def index(self):
        response = self.client.get("/")

    @task(3)
    def discover(self):
        self.client.get("/discover/")

    @task(2)
    def projects(self):
        self.client.get("/projects/unitarian-church-of-staten-island/")

    @task(1)
    def blog(self):
        self.client.get("/blog/")
        self.client.get("/welcome-to-the-new-blocpower/")

    @task(1)
    def about(self):
        self.client.get("/about-us/")

    # This task will 15 times for every 1000 runs of the above task
    # @task(15)
    # def about(self):
    #     self.client.get("/blog")

    # This task will run once for every 1000 runs of the above task
    # @task(1)
    # def about(self):
    #     id = id_generator()
    #     self.client.post("/signup", {"email": "example@example.com", "name": "Test"})

class MyLocust(HttpLocust):
    host = os.getenv('TARGET_URL', "http://localhost")
    task_set = MyTaskSet
    min_wait = 45
    max_wait = 50
<|endoftext|>


<commit_before>"""
Creates a `frozen <http://www.pyinstaller.org/>`_ 32-bit server to use for
`inter-process communication <https://en.wikipedia.org/wiki/Inter-process_communication>`_.

This module creates a 32-bit executable. The executable starts a 32-bit server to
host a 32-bit library. A client module running within a 64-bit Python interpreter
can communicate with the 32-bit library by sending requests to the server. The server
calls the library to execute the request and then the server sends a response back
to the client.

This module must be run from a 32-bit Python interpreter with both
`PyInstaller <http://www.pyinstaller.org/>`_ (to create the executable) and
`Python .NET <https://pypi.python.org/pypi/pythonnet/>`_ (to be able to load a .NET
Framework library) installed.
"""
import os
import sys
import shutil
import subprocess

try:
    from msl import loadlib
except ImportError:
    sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), '..', '..')))
    from msl import loadlib

if loadlib.IS_PYTHON2:
    from urllib import urlopen
elif loadlib.IS_PYTHON3:
    from urllib.request import urlopen
else:
    raise NotImplementedError('Python major version is not 2 or 3')

def main(spec=None):
    """Creates a `frozen <PyInstaller_>`_ 32-bit Python server.

    Uses PyInstaller_ to create a `frozen <PyInstaller_>`_ 32-bit Python executable.
    This executable starts a server, :class:`~.server32.Server32`, which hosts a Python
    module that can load a 32-bit library.

    .. _PyInstaller: http://www.pyinstaller.org/

    Parameters
    ----------
    spec : :obj:`str`, optional
        If you want to freeze using a PyInstaller_ .spec file then you can specify the
        path to the .spec file. Default is :obj:`None`.
    """
    if loadlib.IS_PYTHON_64BIT:
        print('Must run {} using a 32-bit Python interpreter'.format(os.path.basename(__file__)))
        sys.exit(0)

    try:
        import PyInstaller
    except ImportError:
        print('PyInstaller not found, run:')
        print('$ pip install pyinstaller')
        sys.exit(0)

    try:
        import clr
    except ImportError:
        print('pythonnet not found, run:')
        print('$ pip install pythonnet')
        sys.exit(0)

    # start the freezing process

    here = os.path.abspath(os.path.dirname(__file__))
    cmd = ['pyinstaller',
           '--distpath', here,
           '--noconfirm',
           ]
    if spec is None:
        spec_file = '{}.spec'.format(loadlib.SERVER_FILENAME)
        yn = input('A {0} file exists. You may want to run "python freeze_server32.py {0}"\n'
                   'Do you want to continue and overwrite the spec file (y/[n]):'.format(spec_file))
        if yn.lower() not in ('y', 'yes'):
            print('Aborted.')
            return
        cmd.extend(['--name', loadlib.SERVER_FILENAME,
                    '--onefile',
                    '--clean',
                    '--hidden-import', 'clr',
                    ])
        cmd.extend(_get_standard_modules())
        cmd.append(os.path.join(here, 'start_server32.py'))
    else:
        cmd.append(spec)
    subprocess.call(cmd)

    # the --version-file option for pyinstaller does not currently work on Windows, this is a fix
    verpatch = os.path.join(here, 'verpatch.exe')
    if loadlib.IS_WINDOWS and os.path.isfile(verpatch):
        ver = [verpatch,
               os.path.join(here, loadlib.SERVER_FILENAME),
               '/va', '{0}.{1}.{2}'.format(*loadlib.version_info) + '.0',
               '/pv', '{0}.{1}.{2}.{4}'.format(*sys.version_info),
               '/s', 'description', 'Access a 32-bit library from 64-bit Python',
               '/s', 'product', 'Python 32-bit server',
               '/s', 'copyright', loadlib.__copyright__]
        subprocess.call(ver)

    # cleanup
    shutil.rmtree('./build/' + loadlib.SERVER_FILENAME)
    if not os.listdir('./build'):
        shutil.rmtree('./build')
    if loadlib.IS_WINDOWS:
        # pyinstaller is able to include Python.Runtime.dll and Python.Runtime.dll.config
        # automatically in the build, so we don't need to keep the .spec file
        os.remove(loadlib.SERVER_FILENAME + '.spec')

    # create the .NET Framework config file
    loadlib.LoadLibrary.check_dot_net_config(os.path.join(here, loadlib.SERVER_FILENAME))

def _get_standard_modules():
    """
    Returns a list of standard python modules to include and exclude in the
    frozen application.

    PyInstaller does not automatically bundle all of the standard Python modules
    into the frozen application. This
    method parses the 'docs.python.org' website for the list of standard Python
    modules that are available.

    The 'pyinstaller --exclude-module' option ensures that the module is
    excluded from the frozen application.

    The 'pyinstaller --hidden-import' option ensures that the module is included
    into the frozen application (only if the module is available for the operating
    system that is running this script).

    Returns
    -------
    :obj:`list` of :obj:`str`
        A list of modules to be included and excluded.
    """

    # the frozen application is never meant to create GUIs or to add
    # support for building and installing Python modules
    ignore_list = ['__main__', 'distutils', 'ensurepip', 'tkinter', 'turtle']

    url = 'https://docs.python.org/{0}.{1}/py-modindex.html'.format(*sys.version_info)
    source = urlopen(url).read().decode().split('#module-')
    modules = [source[idx].split('"><code')[0] for idx in range(1, len(source))]

    included_modules, excluded_modules = [], []
    for module in modules:
        include_module = True
        for mod in ignore_list:
            if mod in module:
                excluded_modules.extend(['--exclude-module', module])
                include_module = False
                break
        if include_module:
            included_modules.extend(['--hidden-import', module])
    return included_modules + excluded_modules

if __name__ == '__main__':
    spec = None
    if len(sys.argv) > 1:
        if sys.argv[1].endswith('.spec'):
            spec = sys.argv[1]
        else:
            raise IOError('Must pass in a PyInstaller .spec file')
    main(spec)
<commit_msg>properly check if spec_file exists<commit_after>"""
Creates a `frozen <http://www.pyinstaller.org/>`_ 32-bit server to use for
`inter-process communication <https://en.wikipedia.org/wiki/Inter-process_communication>`_.

This module creates a 32-bit executable. The executable starts a 32-bit server to
host a 32-bit library. A client module running within a 64-bit Python interpreter
can communicate with the 32-bit library by sending requests to the server. The server
calls the library to execute the request and then the server sends a response back
to the client.

This module must be run from a 32-bit Python interpreter with both
`PyInstaller <http://www.pyinstaller.org/>`_ (to create the executable) and
`Python .NET <https://pypi.python.org/pypi/pythonnet/>`_ (to be able to load a .NET
Framework library) installed.
"""
import os
import sys
import shutil
import subprocess

try:
    from msl import loadlib
except ImportError:
    sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), '..', '..')))
    from msl import loadlib

if loadlib.IS_PYTHON2:
    from urllib import urlopen
elif loadlib.IS_PYTHON3:
    from urllib.request import urlopen
else:
    raise NotImplementedError('Python major version is not 2 or 3')

def main(spec=None):
    """Creates a `frozen <PyInstaller_>`_ 32-bit Python server.

    Uses PyInstaller_ to create a `frozen <PyInstaller_>`_ 32-bit Python executable.
    This executable starts a server, :class:`~.server32.Server32`, which hosts a Python
    module that can load a 32-bit library.

    .. _PyInstaller: http://www.pyinstaller.org/

    Parameters
    ----------
    spec : :obj:`str`, optional
        If you want to freeze using a PyInstaller_ .spec file then you can specify the
        path to the .spec file. Default is :obj:`None`.
    """
    if loadlib.IS_PYTHON_64BIT:
        print('Must run {} using a 32-bit Python interpreter'.format(os.path.basename(__file__)))
        sys.exit(0)

    try:
        import PyInstaller
    except ImportError:
        print('PyInstaller not found, run:')
        print('$ pip install pyinstaller')
        sys.exit(0)

    try:
        import clr
    except ImportError:
        print('pythonnet not found, run:')
        print('$ pip install pythonnet')
        sys.exit(0)

    # start the freezing process

    here = os.path.abspath(os.path.dirname(__file__))
    cmd = [
        'pyinstaller',
        '--distpath', here,
        '--noconfirm',
    ]

    if spec is None:
        spec_file = '{}.spec'.format(loadlib.SERVER_FILENAME)
        if os.path.exists(spec_file):
            yn = input('A {0} file exists. You may want to run "python freeze_server32.py {0}"\n'
                       'Do you want to continue and overwrite the spec file (y/[n]):'.format(spec_file))
            if yn.lower() not in ('y', 'yes'):
                print('Aborted.')
                return
        cmd.extend([
            '--name', loadlib.SERVER_FILENAME,
            '--onefile',
            '--clean',
            '--hidden-import', 'clr',
        ])
        cmd.extend(_get_standard_modules())
        cmd.append(os.path.join(here, 'start_server32.py'))
    else:
        cmd.append(spec)
    subprocess.call(cmd)

    # the --version-file option for pyinstaller does not currently work on Windows, this is a fix
    verpatch = os.path.join(here, 'verpatch.exe')
    if loadlib.IS_WINDOWS and os.path.isfile(verpatch):
        ver = [verpatch,
               os.path.join(here, loadlib.SERVER_FILENAME),
               '/va', '{0}.{1}.{2}'.format(*loadlib.version_info) + '.0',
               '/pv', '{0}.{1}.{2}.{4}'.format(*sys.version_info),
               '/s', 'description', 'Access a 32-bit library from 64-bit Python',
               '/s', 'product', 'Python 32-bit server',
               '/s', 'copyright', loadlib.__copyright__]
        subprocess.call(ver)

    # cleanup
    shutil.rmtree('./build/' + loadlib.SERVER_FILENAME)
    if not os.listdir('./build'):
        shutil.rmtree('./build')
    if loadlib.IS_WINDOWS:
        # pyinstaller is able to include Python.Runtime.dll and Python.Runtime.dll.config
        # automatically in the build, so we don't need to keep the .spec file
        os.remove(loadlib.SERVER_FILENAME + '.spec')

    # create the .NET Framework config file
    loadlib.LoadLibrary.check_dot_net_config(os.path.join(here, loadlib.SERVER_FILENAME))

def _get_standard_modules():
    """
    Returns a list of standard python modules to include and exclude in the
    frozen application.

    PyInstaller does not automatically bundle all of the standard Python modules
    into the frozen application. This
    method parses the 'docs.python.org' website for the list of standard Python
    modules that are available.

    The 'pyinstaller --exclude-module' option ensures that the module is
    excluded from the frozen application.

    The 'pyinstaller --hidden-import' option ensures that the module is included
    into the frozen application (only if the module is available for the operating
    system that is running this script).

    Returns
    -------
    :obj:`list` of :obj:`str`
        A list of modules to be included and excluded.
    """

    # the frozen application is never meant to create GUIs or to add
    # support for building and installing Python modules
    ignore_list = ['__main__', 'distutils', 'ensurepip', 'tkinter', 'turtle']

    url = 'https://docs.python.org/{0}.{1}/py-modindex.html'.format(*sys.version_info)
    source = urlopen(url).read().decode().split('#module-')
    modules = [source[idx].split('"><code')[0] for idx in range(1, len(source))]

    included_modules, excluded_modules = [], []
    for module in modules:
        include_module = True
        for mod in ignore_list:
            if mod in module:
                excluded_modules.extend(['--exclude-module', module])
                include_module = False
                break
        if include_module:
            included_modules.extend(['--hidden-import', module])
    return included_modules + excluded_modules

if __name__ == '__main__':
    spec = None
    if len(sys.argv) > 1:
        if sys.argv[1].endswith('.spec'):
            spec = sys.argv[1]
        else:
            raise IOError('Must pass in a PyInstaller .spec file')
    main(spec)
<|endoftext|>