<commit_before>import numpy
import cv2
import os
import re
import logging
from collections import namedtuple
try:
    unicode = unicode
except NameError:
    # 'unicode' is undefined, must be Python 3
    str = str
    unicode = str
    bytes = bytes
    basestring = (str, bytes)
else:
    # 'unicode' exists, must be Python 2
    str = str
    unicode = unicode
    bytes = str
    basestring = basestring

HomeCamConfig = namedtuple('HomeCamConfig',
                           ['recording_cam_id',
                            'recording_fps',
                            'recording_file_limit',
                            'recording_time_limit',
                            'recording_enable',
                            'recording_dir',
                            'recording_file_base',
                            'detection_scale_factor',
                            'detection_min_neighbours',
                            'detection_size',
                            'detection_cascade_files'],
                           verbose=False)

# frame          - The current frame (image)
# cascade_status - A dict of detection status (boolean) for all cascades.
#                  The current cascade file is the dict key.
#                  True => The cascade has detected an object, False => No detection.
# rectangles     - A dict of (x, y, width, height) tuples forming a rectangle
#                  of each match in the current frame. The current cascade
#                  file is the dict key.
HomeCamDetectionData = namedtuple('HomeCamDetectionData',
                                  ['frame',
                                   'cascade_status',
                                   'rectangles'],
                                  verbose=False)

# Hard coded BGR tuples for rectangles
BGR_RED = (0, 0, 255)
BGR_GREEN = (0, 255, 0)
BGR_BLUE = (255, 0, 0)
BGR_YELLOW = (0, 255, 255)
BGR_CYAN = (255, 255, 0)
BGR_MAGENTA = (255, 0, 255)
COLORS = [BGR_RED, BGR_GREEN, BGR_BLUE, BGR_YELLOW, BGR_CYAN, BGR_MAGENTA]

class HomeCamException(Exception):

    pass

class HomeCam:

    def __init__(self, config):

        self._logger = logging.getLogger(__name__)

        if config.detection_cascade_files is None:
            raise HomeCamException("Missing cascade file(s)")

        self._min_neighbours = config.detection_min_neighbours
        self._scale_factor = config.detection_scale_factor
        self._size = config.detection_size
        self._cascades = []
        self._save_frame = False
        color_cnt = 0

        for cascade_file in config.detection_cascade_files:
            cur_cascade = cv2.CascadeClassifier(cascade_file)
            if cur_cascade is None:
                raise HomeCamException("Bad cascade file")
            # Associate a color with the cascade.
            # The color will be used when drawing rectangles in the recordings
            # of all detections made with the cascade.
            rgb_tuple = COLORS[color_cnt % len(COLORS)]
            color_cnt += 1
            self._cascades.append((cascade_file, cur_cascade, rgb_tuple))

        self._video_capture = cv2.VideoCapture(int(config.recording_cam_id))
        if self._video_capture is None:
            raise HomeCamException("Unable to open camera")

        # Get the resolution of the capture device and use the same value for
        # recording
        width = self._video_capture.get(3)
        height = self._video_capture.get(4)
        self._recording_resolution = (int(width), int(height))

        if config.recording_enable:
            self._logger.info("Video recording enabled")
            self._recording_fps = config.recording_fps
            self._recording_file_limit = config.recording_file_limit
            self._recording_frame_limit = config.recording_time_limit * config.recording_fps

            self._recording_dir = config.recording_dir
            self._recording_file_base = config.recording_file_base
            self._recording_ext = '.avi'
            self._scan_video_files()
            self._open_new_video_file()
        else:
            self._outfile = None

    def _scan_video_files(self):

        directory = self._recording_dir
        base = self._recording_file_base
        ext = self._recording_ext
        regex = re.compile(base + '(\d+)')

        self._logger.info("Video files dir: %s. File base: %s",
                          directory, base)

        lowest_idx = 0x7fffffff
        highest_idx = 0
        nbr_of_files = 0
        for anyfile in os.listdir(directory):
            (anyfile_base, anyfile_ext) = os.path.splitext(anyfile)
            if not anyfile_ext == ext:
                continue

            m = regex.match(anyfile_base)
            if m is None:
                continue

            idx = int(m.group(1))
            if idx < lowest_idx:
                lowest_idx = idx
            if idx > highest_idx:
                highest_idx = idx
            nbr_of_files += 1

        self._nbr_of_outfiles = nbr_of_files
        if nbr_of_files == 0:
            # There are no logfiles stored in the log file directory
            self._logger.info("Videofile dir empty.")
            self._cur_outfile_index = 0
            self._cur_outfile_lowest_index = 0
        else:
            self._cur_outfile_index = highest_idx + 1
            self._cur_outfile_lowest_index = lowest_idx

        self._logger.info("Cur indices: highest = %d, lowest = %d",
                          self._cur_outfile_index, self._cur_outfile_lowest_index)

    def _open_new_video_file(self):

        directory = self._recording_dir
        base = self._recording_file_base
        ext = self._recording_ext

        new_file_name = directory + '/' + base + str(self._cur_outfile_index) + ext
        self._logger.info("Opening new output file: %s", new_file_name)
        fourcc = cv2.VideoWriter_fourcc(*'mjpa')
        self._logger.info("recording resoluton: {}".format(self._recording_resolution))
        self._outfile = cv2.VideoWriter(new_file_name, fourcc,
                                        self._recording_fps,
                                        self._recording_resolution)
        self._nbr_of_outfiles += 1
        self._cur_nbr_of_recorded_frames = 0

    def _remove_old_video_file(self):

        directory = self._recording_dir
        base = self._recording_file_base
        ext = self._recording_ext

        oldest_filename = directory + '/' + base + str(self._cur_outfile_lowest_index) + ext
        self._logger.info("Removing old output file: %s", oldest_filename)
        os.remove(oldest_filename)
        # Update oldest and current index by rescanning all outfiles
        self._scan_video_files()

    def _do_save_frame(self, frame):

        if self._cur_nbr_of_recorded_frames > self._recording_frame_limit:
            self._logger.info("Switching output file")
            self._outfile.release()
            self._cur_outfile_index += 1
            self._open_new_video_file()

        self._outfile.write(frame)

        self._cur_nbr_of_recorded_frames += 1
        if self._nbr_of_outfiles > self._recording_file_limit:
            self._remove_old_video_file()

    # Read one frame from the cam and process it.
    # Returns a HomeCamDetectionData named tuple containing all detection
    # data.
    def read_and_process_frame(self):

        # Initialize the detection data
        cascade_status = {}
        rectangles = {}
        for (cascade_file, cascade, rgb_tuple) in self._cascades:
            cascade_status[cascade_file] = False
            rectangles[cascade_file] = None

        # Check error conditions
        if not self._video_capture.isOpened():
            return HomeCamDetectionData(frame=None,
                                        cascade_status=cascade_status,
                                        rectangles=rectangles)

        ret, frame = self._video_capture.read()
        if not ret:
            return HomeCamDetectionData(frame=None,
                                        cascade_status=cascade_status,
                                        rectangles=rectangles)

        frame_gs = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)

        # Check detections for each cascade
        for (cascade_file, cascade, rgb_tuple) in self._cascades:
            obj = cascade.detectMultiScale(frame_gs,
                                           scaleFactor=self._scale_factor,
                                           minNeighbors=self._min_neighbours,
                                           minSize=(self._size, self._size))
            if len(obj) == 0:
                continue

            cascade_status[cascade_file] = True
            rectangles[cascade_file] = obj

            for (x, y, w, h) in obj:
                cv2.rectangle(frame, (x, y), (x+w, y+h), rgb_tuple, 2)

        if self._save_frame and self._outfile is not None:
            self._do_save_frame(frame)

        return HomeCamDetectionData(frame=frame_gs,
                                    cascade_status=cascade_status,
                                    rectangles=rectangles)

    def enable_frame_saving(self):

        self._save_frame = True

    def disable_frame_saving(self):

        self._save_frame = False

    def close(self):

        if self._video_capture is not None:
            self._logger.info("Closing video capture device")
            self._video_capture.release()

        if self._outfile is not None:
            self._logger.info("Closing video output file")
            self._outfile.release()
<commit_msg>home_cam.py: add capture frame resize<commit_after>import numpy
import cv2
import imutils
import os
import re
import logging
from collections import namedtuple
try:
    unicode = unicode
except NameError:
    # 'unicode' is undefined, must be Python 3
    str = str
    unicode = str
    bytes = bytes
    basestring = (str, bytes)
else:
    # 'unicode' exists, must be Python 2
    str = str
    unicode = unicode
    bytes = str
    basestring = basestring

HomeCamConfig = namedtuple('HomeCamConfig',
                           ['recording_cam_id',
                            'recording_fps',
                            'recording_file_limit',
                            'recording_time_limit',
                            'recording_enable',
                            'recording_dir',
                            'recording_file_base',
                            'detection_scale_factor',
                            'detection_min_neighbours',
                            'detection_size',
                            'detection_cascade_files'],
                           verbose=False)

# frame          - The current frame (image)
# cascade_status - A dict of detection status (boolean) for all cascades.
#                  The current cascade file is the dict key.
#                  True => The cascade has detected an object, False => No detection.
# rectangles     - A dict of (x, y, width, height) tuples forming a rectangle
#                  of each match in the current frame. The current cascade
#                  file is the dict key.
HomeCamDetectionData = namedtuple('HomeCamDetectionData',
                                  ['frame',
                                   'cascade_status',
                                   'rectangles'],
                                  verbose=False)

# Hard coded BGR tuples for rectangles
BGR_RED = (0, 0, 255)
BGR_GREEN = (0, 255, 0)
BGR_BLUE = (255, 0, 0)
BGR_YELLOW = (0, 255, 255)
BGR_CYAN = (255, 255, 0)
BGR_MAGENTA = (255, 0, 255)
COLORS = [BGR_RED, BGR_GREEN, BGR_BLUE, BGR_YELLOW, BGR_CYAN, BGR_MAGENTA]

class HomeCamException(Exception):

    pass

class HomeCam:

    def __init__(self, config):

        self._logger = logging.getLogger(__name__)

        if config.detection_cascade_files is None:
            raise HomeCamException("Missing cascade file(s)")

        self._min_neighbours = config.detection_min_neighbours
        self._scale_factor = config.detection_scale_factor
        self._size = config.detection_size
        self._cascades = []
        self._save_frame = False
        color_cnt = 0

        for cascade_file in config.detection_cascade_files:
            cur_cascade = cv2.CascadeClassifier(cascade_file)
            if cur_cascade is None:
                raise HomeCamException("Bad cascade file")
            # Associate a color with the cascade.
            # The color will be used when drawing rectangles in the recordings
            # of all detections made with the cascade.
            rgb_tuple = COLORS[color_cnt % len(COLORS)]
            color_cnt += 1
            self._cascades.append((cascade_file, cur_cascade, rgb_tuple))

        self._video_capture = cv2.VideoCapture(int(config.recording_cam_id))
        if self._video_capture is None:
            raise HomeCamException("Unable to open camera")

        # Get the resolution of the capture device and use the same value for
        # recording
        width = self._video_capture.get(3)
        height = self._video_capture.get(4)
        # We are going to resize all captured frames to a width of max 400 pixels,
        # so we must make sure the recording resolution matches the rescaled
        # frames.
        rec_width = min(400, width)
        rec_height = rec_width * height / width
        self._recording_resolution = (int(rec_width), int(rec_height))

        if config.recording_enable:
            self._logger.info("Video recording enabled")
            self._recording_fps = config.recording_fps
            self._recording_file_limit = config.recording_file_limit
            self._recording_frame_limit = config.recording_time_limit * config.recording_fps

            self._recording_dir = config.recording_dir
            self._recording_file_base = config.recording_file_base
            self._recording_ext = '.avi'
            self._scan_video_files()
            self._open_new_video_file()
        else:
            self._outfile = None

    def _scan_video_files(self):

        directory = self._recording_dir
        base = self._recording_file_base
        ext = self._recording_ext
        regex = re.compile(base + '(\d+)')

        self._logger.info("Video files dir: %s. File base: %s",
                          directory, base)

        lowest_idx = 0x7fffffff
        highest_idx = 0
        nbr_of_files = 0
        for anyfile in os.listdir(directory):
            (anyfile_base, anyfile_ext) = os.path.splitext(anyfile)
            if not anyfile_ext == ext:
                continue

            m = regex.match(anyfile_base)
            if m is None:
                continue

            idx = int(m.group(1))
            if idx < lowest_idx:
                lowest_idx = idx
            if idx > highest_idx:
                highest_idx = idx
            nbr_of_files += 1

        self._nbr_of_outfiles = nbr_of_files
        if nbr_of_files == 0:
            # There are no logfiles stored in the log file directory
            self._logger.info("Videofile dir empty.")
            self._cur_outfile_index = 0
            self._cur_outfile_lowest_index = 0
        else:
            self._cur_outfile_index = highest_idx + 1
            self._cur_outfile_lowest_index = lowest_idx

        self._logger.info("Cur indices: highest = %d, lowest = %d",
                          self._cur_outfile_index, self._cur_outfile_lowest_index)

    def _open_new_video_file(self):

        directory = self._recording_dir
        base = self._recording_file_base
        ext = self._recording_ext

        new_file_name = directory + '/' + base + str(self._cur_outfile_index) + ext
        self._logger.info("Opening new output file: %s", new_file_name)
        fourcc = cv2.VideoWriter_fourcc(*'mjpa')
        self._logger.info("recording resoluton: {}".format(self._recording_resolution))
        self._outfile = cv2.VideoWriter(new_file_name, fourcc,
                                        self._recording_fps,
                                        self._recording_resolution)
        self._nbr_of_outfiles += 1
        self._cur_nbr_of_recorded_frames = 0

    def _remove_old_video_file(self):

        directory = self._recording_dir
        base = self._recording_file_base
        ext = self._recording_ext

        oldest_filename = directory + '/' + base + str(self._cur_outfile_lowest_index) + ext
        self._logger.info("Removing old output file: %s", oldest_filename)
        os.remove(oldest_filename)
        # Update oldest and current index by rescanning all outfiles
        self._scan_video_files()

    def _do_save_frame(self, frame):

        if self._cur_nbr_of_recorded_frames > self._recording_frame_limit:
            self._logger.info("Switching output file")
            self._outfile.release()
            self._cur_outfile_index += 1
            self._open_new_video_file()

        self._outfile.write(frame)

        self._cur_nbr_of_recorded_frames += 1
        if self._nbr_of_outfiles > self._recording_file_limit:
            self._remove_old_video_file()

    # Read one frame from the cam and process it.
    # Returns a HomeCamDetectionData named tuple containing all detection
    # data.
    def read_and_process_frame(self):

        # Initialize the detection data
        cascade_status = {}
        rectangles = {}
        for (cascade_file, cascade, rgb_tuple) in self._cascades:
            cascade_status[cascade_file] = False
            rectangles[cascade_file] = None

        # Check error conditions
        if not self._video_capture.isOpened():
            return HomeCamDetectionData(frame=None,
                                        cascade_status=cascade_status,
                                        rectangles=rectangles)

        ret, frame = self._video_capture.read()
        if not ret:
            return HomeCamDetectionData(frame=None,
                                        cascade_status=cascade_status,
                                        rectangles=rectangles)

        # Resize the frame.
        # Below code snippet is taken from:
        # http://www.pyimagesearch.com/2015/11/09/pedestrian-detection-opencv/
        # load the image and resize it to (1) reduce detection time
        # and (2) improve detection accuracy
        frame = imutils.resize(frame, width=min(400, frame.shape[1]))

        frame_gs = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)

        # Check detections for each cascade
        for (cascade_file, cascade, rgb_tuple) in self._cascades:
            obj = cascade.detectMultiScale(frame_gs,
                                           scaleFactor=self._scale_factor,
                                           minNeighbors=self._min_neighbours,
                                           minSize=(self._size, self._size))
            if len(obj) == 0:
                continue

            cascade_status[cascade_file] = True
            rectangles[cascade_file] = obj

            for (x, y, w, h) in obj:
                cv2.rectangle(frame, (x, y), (x+w, y+h), rgb_tuple, 2)

        if self._save_frame and self._outfile is not None:
            self._do_save_frame(frame)

        return HomeCamDetectionData(frame=frame_gs,
                                    cascade_status=cascade_status,
                                    rectangles=rectangles)

    def enable_frame_saving(self):

        self._save_frame = True

    def disable_frame_saving(self):

        self._save_frame = False

    def close(self):

        if self._video_capture is not None:
            self._logger.info("Closing video capture device")
            self._video_capture.release()

        if self._outfile is not None:
            self._logger.info("Closing video output file")
            self._outfile.release()
<|endoftext|>


<commit_before># Copyright (c) 2012, Dylan Hackers.
# See License.txt for details.

import fcntl, os, subprocess, time
from textwrap import dedent

TIMEOUT = 50

LID_TEMPLATE = dedent("""\
  library: %(taskID)s
  executable: %(taskID)s
  files: library
         code
  """)

LIBRARY_TEMPLATE = dedent("""\
  module: dylan-user

  define library %(taskID)s
    use common-dylan;
    use io;
  end library;

  define module %(taskID)s
    use common-dylan;
    use format-out;
    use standard-io;
    use streams;
  end module;
  """)

DYLAN_TEMPLATE = dedent("""\
  module: %(taskID)s

  %(code)s
  """)

def writeToFile(fileName, data):
  f = open(fileName, 'w')
  f.write(data)
  f.close()

def makeNonBlocking(f):
  fd = f.fileno()
  fl = fcntl.fcntl(fd, fcntl.F_GETFL)
  fcntl.fcntl(fd, fcntl.F_SETFL, fl | os.O_NONBLOCK)

def safeRead(f):
  try:
    return f.read()
  except:
    return ''

class DylanBuilder(object):
  def __init__(self, taskID, code):
    self.directoryName = str(taskID)
    self.baseName = str(taskID)
    self.lidFileName = os.path.join(self.directoryName, self.baseName + '.lid')
    self.libraryFileName = os.path.join(self.directoryName, 'library.dylan')
    self.dylanFileName = os.path.join(self.directoryName, 'code.dylan')
    self.params = {
      'taskID': taskID,
      'code': code
    }

  def prepare(self):
    os.mkdir(self.directoryName)
    writeToFile(self.lidFileName, LID_TEMPLATE % self.params)
    writeToFile(self.libraryFileName, LIBRARY_TEMPLATE % self.params)
    writeToFile(self.dylanFileName, DYLAN_TEMPLATE % self.params)

  def compile(self):
    (ret, stdout, stderr) = self.runWithTimeout(TIMEOUT, ["/opt/opendylan-2011.1/bin/dylan-compiler", "-build", self.lidFileName])
    writeToFile(os.path.join(self.directoryName, 'compile-out.txt'), stdout)
    writeToFile(os.path.join(self.directoryName, 'compile-err.txt'), stderr)
    return (ret, stdout, stderr)

  def run(self):
    (ret, stdout, stderr) = self.runWithTimeout(TIMEOUT, [os.path.join("_build", "bin", self.baseName)])
    writeToFile(os.path.join(self.directoryName, 'run-out.txt'), stdout)
    writeToFile(os.path.join(self.directoryName, 'run-err.txt'), stderr)
    return (ret, stdout, stderr)

  def runWithTimeout(self, timeout, args):
    beginTime = time.time()
    elapsedTime = 0

    stdout = ''
    stderr = ''

    p = subprocess.Popen(args, stdin=subprocess.PIPE, stdout=subprocess.PIPE, stderr=subprocess.PIPE, shell=False)

    makeNonBlocking(p.stdout)
    makeNonBlocking(p.stderr)

    while p.poll() is None and elapsedTime < timeout:
      time.sleep(0.1)
      elapsedTime = time.time() - beginTime
      stdout += safeRead(p.stdout)
      stderr += safeRead(p.stderr)

    if elapsedTime >= timeout:
      p.kill()

    print 'STDOUT: ', stdout
    print 'STDERR: ', stderr

    return (p.returncode, stdout, stderr)

<commit_msg>Remove print statements.<commit_after># Copyright (c) 2012, Dylan Hackers.
# See License.txt for details.

import fcntl, os, subprocess, time
from textwrap import dedent

TIMEOUT = 50

LID_TEMPLATE = dedent("""\
  library: %(taskID)s
  executable: %(taskID)s
  files: library
         code
  """)

LIBRARY_TEMPLATE = dedent("""\
  module: dylan-user

  define library %(taskID)s
    use common-dylan;
    use io;
  end library;

  define module %(taskID)s
    use common-dylan;
    use format-out;
    use standard-io;
    use streams;
  end module;
  """)

DYLAN_TEMPLATE = dedent("""\
  module: %(taskID)s

  %(code)s
  """)

def writeToFile(fileName, data):
  f = open(fileName, 'w')
  f.write(data)
  f.close()

def makeNonBlocking(f):
  fd = f.fileno()
  fl = fcntl.fcntl(fd, fcntl.F_GETFL)
  fcntl.fcntl(fd, fcntl.F_SETFL, fl | os.O_NONBLOCK)

def safeRead(f):
  try:
    return f.read()
  except:
    return ''

class DylanBuilder(object):
  def __init__(self, taskID, code):
    self.directoryName = str(taskID)
    self.baseName = str(taskID)
    self.lidFileName = os.path.join(self.directoryName, self.baseName + '.lid')
    self.libraryFileName = os.path.join(self.directoryName, 'library.dylan')
    self.dylanFileName = os.path.join(self.directoryName, 'code.dylan')
    self.params = {
      'taskID': taskID,
      'code': code
    }

  def prepare(self):
    os.mkdir(self.directoryName)
    writeToFile(self.lidFileName, LID_TEMPLATE % self.params)
    writeToFile(self.libraryFileName, LIBRARY_TEMPLATE % self.params)
    writeToFile(self.dylanFileName, DYLAN_TEMPLATE % self.params)

  def compile(self):
    (ret, stdout, stderr) = self.runWithTimeout(TIMEOUT, ["/opt/opendylan-2011.1/bin/dylan-compiler", "-build", self.lidFileName])
    writeToFile(os.path.join(self.directoryName, 'compile-out.txt'), stdout)
    writeToFile(os.path.join(self.directoryName, 'compile-err.txt'), stderr)
    return (ret, stdout, stderr)

  def run(self):
    (ret, stdout, stderr) = self.runWithTimeout(TIMEOUT, [os.path.join("_build", "bin", self.baseName)])
    writeToFile(os.path.join(self.directoryName, 'run-out.txt'), stdout)
    writeToFile(os.path.join(self.directoryName, 'run-err.txt'), stderr)
    return (ret, stdout, stderr)

  def runWithTimeout(self, timeout, args):
    beginTime = time.time()
    elapsedTime = 0

    stdout = ''
    stderr = ''

    p = subprocess.Popen(args, stdin=subprocess.PIPE, stdout=subprocess.PIPE, stderr=subprocess.PIPE, shell=False)

    makeNonBlocking(p.stdout)
    makeNonBlocking(p.stderr)

    while p.poll() is None and elapsedTime < timeout:
      time.sleep(0.1)
      elapsedTime = time.time() - beginTime
      stdout += safeRead(p.stdout)
      stderr += safeRead(p.stderr)

    if elapsedTime >= timeout:
      p.kill()

    return (p.returncode, stdout, stderr)

<|endoftext|>


<commit_before># Copyright 2015 Google Inc. All Rights Reserved.

# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at

#     http://www.apache.org/licenses/LICENSE-2.0

# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""User input module for OpenHTF.

Provides a plug which can be used to prompt the user for input. The prompt can
be displayed in the console, the OpenHTF web GUI, and custom frontends.
"""

import functools
import logging
import os
import platform
import select
import sys
import threading
from typing import Any, Callable, Dict, Optional, Text, Tuple, Union
import uuid

import attr
import openhtf
from openhtf import plugs
from openhtf.core import base_plugs
from openhtf.util import console_output

if platform.system() != 'Windows':
  import termios  # pylint: disable=g-import-not-at-top

_LOG = logging.getLogger(__name__)

PROMPT = '--> '

class PromptInputError(Exception):
  """Raised in the event that a prompt returns without setting the response."""

class MultiplePromptsError(Exception):
  """Raised if a prompt is invoked while there is an existing prompt."""

class PromptUnansweredError(Exception):
  """Raised when a prompt times out or otherwise comes back unanswered."""

@attr.s(slots=True, frozen=True)
class Prompt(object):
  id = attr.ib(type=Text)
  message = attr.ib(type=Text)
  text_input = attr.ib(type=bool)
  image_url = attr.ib(type=Optional[Text], default=None)

class ConsolePrompt(threading.Thread):
  """Thread that displays a prompt to the console and waits for a response.

  This should not be used for processes that run in the background.
  """

  def __init__(self,
               message: Text,
               callback: Callable[[Text], None],
               color: Text = ''):
    """Initializes a ConsolePrompt.

    Args:
      message: A string to be presented to the user.
      callback: A function to be called with the response string.
      color: An ANSI color code, or the empty string.
    """
    super(ConsolePrompt, self).__init__()
    self.daemon = True
    self._message = message
    self._callback = callback
    self._color = color
    self._stop_event = threading.Event()
    self._answered = False

  def stop(self) -> None:
    """Mark this ConsolePrompt as stopped."""
    self._stop_event.set()
    if not self._answered:
      console_output.cli_print(
          os.linesep, color=self._color, end='', logger=None)
      _LOG.debug('Stopping ConsolePrompt--prompt was answered from elsewhere.')

  def run(self) -> None:
    """Main logic for this thread to execute."""
    if platform.system() == 'Windows':
      # Windows doesn't support file-like objects for select(), so fall back
      # to raw_input().
      response = input(''.join((self._message, os.linesep, PROMPT)))
      self._answered = True
      self._callback(response)
      return

    # First, display the prompt to the console.
    console_output.cli_print(
        self._message, color=self._color, end=os.linesep, logger=None)
    console_output.cli_print(PROMPT, color=self._color, end='', logger=None)
    sys.stdout.flush()

    # Before reading, clear any lingering buffered terminal input.
    termios.tcflush(sys.stdin, termios.TCIFLUSH)

    line = ''
    while not self._stop_event.is_set():
      inputs, _, _ = select.select([sys.stdin], [], [], 0.001)
      if sys.stdin in inputs:
        new = os.read(sys.stdin.fileno(), 1024)
        if not new:
          # Hit EOF!
          # They hit ^D (to insert EOF). Tell them to hit ^C if they
          # want to actually quit.
          print('Hit ^C (Ctrl+c) to exit.')
          break
        line += new.decode('utf-8')
        if '\n' in line:
          response = line[:line.find('\n')]
          self._answered = True
          self._callback(response)
          return

class UserInput(base_plugs.FrontendAwareBasePlug):
  """Get user input from inside test phases.

  Attributes:
    last_response: None, or a pair of (prompt_id, response) indicating the last
      user response that was received by the plug.
  """

  def __init__(self):
    super(UserInput, self).__init__()
    self.last_response = None  # type: Optional[Tuple[Text, Text]]
    self._prompt = None  # type: Optional[Prompt]
    self._console_prompt = None  # type: Optional[ConsolePrompt]
    self._response = None  # type: Optional[Text]
    self._cond = threading.Condition(threading.RLock())

  def _asdict(self) -> Optional[Dict[Text, Any]]:
    """Return a dictionary representation of the current prompt."""
    with self._cond:
      if self._prompt is None:
        return None
      return {
          'id': self._prompt.id,
          'message': self._prompt.message,
          'text-input': self._prompt.text_input
      }

  def tearDown(self) -> None:
    self.remove_prompt()

  def remove_prompt(self) -> None:
    """Remove the prompt."""
    with self._cond:
      self._prompt = None
      if self._console_prompt:
        self._console_prompt.stop()
        self._console_prompt = None
      self.notify_update()

  def prompt(self,
             message: Text,
             text_input: bool = False,
             timeout_s: Union[int, float, None] = None,
             cli_color: Text = '',
             image_url: Optional[Text] = None) -> Text:
    """Display a prompt and wait for a response.

    Args:
      message: A string to be presented to the user.
      text_input: A boolean indicating whether the user must respond with text.
      timeout_s: Seconds to wait before raising a PromptUnansweredError.
      cli_color: An ANSI color code, or the empty string.
      image_url: Optional image URL to display or None.

    Returns:
      A string response, or the empty string if text_input was False.

    Raises:
      MultiplePromptsError: There was already an existing prompt.
      PromptUnansweredError: Timed out waiting for the user to respond.
    """
    self.start_prompt(message, text_input, cli_color, image_url)
    return self.wait_for_prompt(timeout_s)

  def start_prompt(self,
                   message: Text,
                   text_input: bool = False,
                   cli_color: Text = '',
                   image_url: Optional[Text] = None) -> Text:
    """Display a prompt.

    Args:
      message: A string to be presented to the user.
      text_input: A boolean indicating whether the user must respond with text.
      cli_color: An ANSI color code, or the empty string.
      image_url: Optional image URL to display or None.

    Raises:
      MultiplePromptsError: There was already an existing prompt.

    Returns:
      A string uniquely identifying the prompt.
    """
    with self._cond:
      if self._prompt:
        raise MultiplePromptsError(
            'Multiple concurrent prompts are not supported.')
      prompt_id = uuid.uuid4().hex
      _LOG.debug('Displaying prompt (%s): "%s"%s', prompt_id, message,
                 ', Expects text input.' if text_input else '')

      self._response = None
      self._prompt = Prompt(
          id=prompt_id,
          message=message,
          text_input=text_input,
          image_url=image_url)
      if sys.stdin.isatty():
        self._console_prompt = ConsolePrompt(
            message, functools.partial(self.respond, prompt_id), cli_color)
        self._console_prompt.start()

      self.notify_update()
      return prompt_id

  def wait_for_prompt(self, timeout_s: Union[int, float, None] = None) -> Text:
    """Wait for the user to respond to the current prompt.

    Args:
      timeout_s: Seconds to wait before raising a PromptUnansweredError.

    Returns:
      A string response, or the empty string if text_input was False.

    Raises:
      PromptUnansweredError: Timed out waiting for the user to respond.
    """
    with self._cond:
      if self._prompt:
        if timeout_s is None:
          self._cond.wait(3600 * 24 * 365)
        else:
          self._cond.wait(timeout_s)
      if self._response is None:
        raise PromptUnansweredError
      return self._response

  def respond(self, prompt_id: Text, response: Text) -> None:
    """Respond to the prompt with the given ID.

    If there is no active prompt or the given ID doesn't match the active
    prompt, do nothing.

    Args:
      prompt_id: A string uniquely identifying the prompt.
      response: A string response to the given prompt.
    """
    _LOG.debug('Responding to prompt (%s): "%s"', prompt_id, response)
    with self._cond:
      if not (self._prompt and self._prompt.id == prompt_id):
        return
      self._response = response
      self.last_response = (prompt_id, response)
      self.remove_prompt()
      self._cond.notifyAll()

def prompt_for_test_start(
    message: Text = 'Enter a DUT ID in order to start the test.',
    timeout_s: Union[int, float, None] = 60 * 60 * 24,
    validator: Callable[[Text], Text] = lambda sn: sn,
    cli_color: Text = '') -> openhtf.PhaseDescriptor:
  """Returns an OpenHTF phase for use as a prompt-based start trigger.

  Args:
    message: The message to display to the user.
    timeout_s: Seconds to wait before raising a PromptUnansweredError.
    validator: Function used to validate or modify the serial number.
    cli_color: An ANSI color code, or the empty string.
  """

  @openhtf.PhaseOptions(timeout_s=timeout_s)
  @plugs.plug(prompts=UserInput)
  def trigger_phase(test: openhtf.TestApi, prompts: UserInput) -> None:
    """Test start trigger that prompts the user for a DUT ID."""
    dut_id = prompts.prompt(
        message, text_input=True, timeout_s=timeout_s, cli_color=cli_color)
    test.test_record.dut_id = validator(dut_id)

  return trigger_phase
<commit_msg>Fixed bug with 'show image with user.prompt' One line was over written with a previous merge that stopped user.prompt with images from working.<commit_after># Copyright 2015 Google Inc. All Rights Reserved.

# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at

#     http://www.apache.org/licenses/LICENSE-2.0

# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""User input module for OpenHTF.

Provides a plug which can be used to prompt the user for input. The prompt can
be displayed in the console, the OpenHTF web GUI, and custom frontends.
"""

import functools
import logging
import os
import platform
import select
import sys
import threading
from typing import Any, Callable, Dict, Optional, Text, Tuple, Union
import uuid

import attr
import openhtf
from openhtf import plugs
from openhtf.core import base_plugs
from openhtf.util import console_output

if platform.system() != 'Windows':
  import termios  # pylint: disable=g-import-not-at-top

_LOG = logging.getLogger(__name__)

PROMPT = '--> '

class PromptInputError(Exception):
  """Raised in the event that a prompt returns without setting the response."""

class MultiplePromptsError(Exception):
  """Raised if a prompt is invoked while there is an existing prompt."""

class PromptUnansweredError(Exception):
  """Raised when a prompt times out or otherwise comes back unanswered."""

@attr.s(slots=True, frozen=True)
class Prompt(object):
  id = attr.ib(type=Text)
  message = attr.ib(type=Text)
  text_input = attr.ib(type=bool)
  image_url = attr.ib(type=Optional[Text], default=None)

class ConsolePrompt(threading.Thread):
  """Thread that displays a prompt to the console and waits for a response.

  This should not be used for processes that run in the background.
  """

  def __init__(self,
               message: Text,
               callback: Callable[[Text], None],
               color: Text = ''):
    """Initializes a ConsolePrompt.

    Args:
      message: A string to be presented to the user.
      callback: A function to be called with the response string.
      color: An ANSI color code, or the empty string.
    """
    super(ConsolePrompt, self).__init__()
    self.daemon = True
    self._message = message
    self._callback = callback
    self._color = color
    self._stop_event = threading.Event()
    self._answered = False

  def stop(self) -> None:
    """Mark this ConsolePrompt as stopped."""
    self._stop_event.set()
    if not self._answered:
      console_output.cli_print(
          os.linesep, color=self._color, end='', logger=None)
      _LOG.debug('Stopping ConsolePrompt--prompt was answered from elsewhere.')

  def run(self) -> None:
    """Main logic for this thread to execute."""
    if platform.system() == 'Windows':
      # Windows doesn't support file-like objects for select(), so fall back
      # to raw_input().
      response = input(''.join((self._message, os.linesep, PROMPT)))
      self._answered = True
      self._callback(response)
      return

    # First, display the prompt to the console.
    console_output.cli_print(
        self._message, color=self._color, end=os.linesep, logger=None)
    console_output.cli_print(PROMPT, color=self._color, end='', logger=None)
    sys.stdout.flush()

    # Before reading, clear any lingering buffered terminal input.
    termios.tcflush(sys.stdin, termios.TCIFLUSH)

    line = ''
    while not self._stop_event.is_set():
      inputs, _, _ = select.select([sys.stdin], [], [], 0.001)
      if sys.stdin in inputs:
        new = os.read(sys.stdin.fileno(), 1024)
        if not new:
          # Hit EOF!
          # They hit ^D (to insert EOF). Tell them to hit ^C if they
          # want to actually quit.
          print('Hit ^C (Ctrl+c) to exit.')
          break
        line += new.decode('utf-8')
        if '\n' in line:
          response = line[:line.find('\n')]
          self._answered = True
          self._callback(response)
          return

class UserInput(base_plugs.FrontendAwareBasePlug):
  """Get user input from inside test phases.

  Attributes:
    last_response: None, or a pair of (prompt_id, response) indicating the last
      user response that was received by the plug.
  """

  def __init__(self):
    super(UserInput, self).__init__()
    self.last_response = None  # type: Optional[Tuple[Text, Text]]
    self._prompt = None  # type: Optional[Prompt]
    self._console_prompt = None  # type: Optional[ConsolePrompt]
    self._response = None  # type: Optional[Text]
    self._cond = threading.Condition(threading.RLock())

  def _asdict(self) -> Optional[Dict[Text, Any]]:
    """Return a dictionary representation of the current prompt."""
    with self._cond:
      if self._prompt is None:
        return None
      return {
          'id': self._prompt.id,
          'message': self._prompt.message,
          'text-input': self._prompt.text_input,
          'image-url': self._prompt.image_url
      }

  def tearDown(self) -> None:
    self.remove_prompt()

  def remove_prompt(self) -> None:
    """Remove the prompt."""
    with self._cond:
      self._prompt = None
      if self._console_prompt:
        self._console_prompt.stop()
        self._console_prompt = None
      self.notify_update()

  def prompt(self,
             message: Text,
             text_input: bool = False,
             timeout_s: Union[int, float, None] = None,
             cli_color: Text = '',
             image_url: Optional[Text] = None) -> Text:
    """Display a prompt and wait for a response.

    Args:
      message: A string to be presented to the user.
      text_input: A boolean indicating whether the user must respond with text.
      timeout_s: Seconds to wait before raising a PromptUnansweredError.
      cli_color: An ANSI color code, or the empty string.
      image_url: Optional image URL to display or None.

    Returns:
      A string response, or the empty string if text_input was False.

    Raises:
      MultiplePromptsError: There was already an existing prompt.
      PromptUnansweredError: Timed out waiting for the user to respond.
    """
    self.start_prompt(message, text_input, cli_color, image_url)
    return self.wait_for_prompt(timeout_s)

  def start_prompt(self,
                   message: Text,
                   text_input: bool = False,
                   cli_color: Text = '',
                   image_url: Optional[Text] = None) -> Text:
    """Display a prompt.

    Args:
      message: A string to be presented to the user.
      text_input: A boolean indicating whether the user must respond with text.
      cli_color: An ANSI color code, or the empty string.
      image_url: Optional image URL to display or None.

    Raises:
      MultiplePromptsError: There was already an existing prompt.

    Returns:
      A string uniquely identifying the prompt.
    """
    with self._cond:
      if self._prompt:
        raise MultiplePromptsError(
            'Multiple concurrent prompts are not supported.')
      prompt_id = uuid.uuid4().hex
      _LOG.debug('Displaying prompt (%s): "%s"%s', prompt_id, message,
                 ', Expects text input.' if text_input else '')

      self._response = None
      self._prompt = Prompt(
          id=prompt_id,
          message=message,
          text_input=text_input,
          image_url=image_url)
      if sys.stdin.isatty():
        self._console_prompt = ConsolePrompt(
            message, functools.partial(self.respond, prompt_id), cli_color)
        self._console_prompt.start()

      self.notify_update()
      return prompt_id

  def wait_for_prompt(self, timeout_s: Union[int, float, None] = None) -> Text:
    """Wait for the user to respond to the current prompt.

    Args:
      timeout_s: Seconds to wait before raising a PromptUnansweredError.

    Returns:
      A string response, or the empty string if text_input was False.

    Raises:
      PromptUnansweredError: Timed out waiting for the user to respond.
    """
    with self._cond:
      if self._prompt:
        if timeout_s is None:
          self._cond.wait(3600 * 24 * 365)
        else:
          self._cond.wait(timeout_s)
      if self._response is None:
        raise PromptUnansweredError
      return self._response

  def respond(self, prompt_id: Text, response: Text) -> None:
    """Respond to the prompt with the given ID.

    If there is no active prompt or the given ID doesn't match the active
    prompt, do nothing.

    Args:
      prompt_id: A string uniquely identifying the prompt.
      response: A string response to the given prompt.
    """
    _LOG.debug('Responding to prompt (%s): "%s"', prompt_id, response)
    with self._cond:
      if not (self._prompt and self._prompt.id == prompt_id):
        return
      self._response = response
      self.last_response = (prompt_id, response)
      self.remove_prompt()
      self._cond.notifyAll()

def prompt_for_test_start(
    message: Text = 'Enter a DUT ID in order to start the test.',
    timeout_s: Union[int, float, None] = 60 * 60 * 24,
    validator: Callable[[Text], Text] = lambda sn: sn,
    cli_color: Text = '') -> openhtf.PhaseDescriptor:
  """Returns an OpenHTF phase for use as a prompt-based start trigger.

  Args:
    message: The message to display to the user.
    timeout_s: Seconds to wait before raising a PromptUnansweredError.
    validator: Function used to validate or modify the serial number.
    cli_color: An ANSI color code, or the empty string.
  """

  @openhtf.PhaseOptions(timeout_s=timeout_s)
  @plugs.plug(prompts=UserInput)
  def trigger_phase(test: openhtf.TestApi, prompts: UserInput) -> None:
    """Test start trigger that prompts the user for a DUT ID."""
    dut_id = prompts.prompt(
        message, text_input=True, timeout_s=timeout_s, cli_color=cli_color)
    test.test_record.dut_id = validator(dut_id)

  return trigger_phase
<|endoftext|>


<commit_before>"""
This module contains the HTTP fetcher interface and several implementations.
"""
import urllib2
import time
import cStringIO

class OpenIDHTTPFetcher(object):
    """
    This class is the interface for HTTP fetchers the OpenID consumer
    library uses.  This interface is only important if you need to
    write a new fetcher for some reason.
    """

    def get(self, url):
        """
        This performs an HTTP get, following redirects along the way.

        @return: This returns a three-tuple on success.  The first value
            is the http return code. The second value is the final url
            that was fetched, after following any redirects.  The third
            value is the data that was retrieved from the site.  If the
            fetch didn't succeed, return C{None}.

        @rtype: C{None} or (C{int}, C{str}, C{str})

        @raise Exception: There are also failure conditions where this
            may raise an exception, rather than returning None.  When
            this occurs, the OpenID consumer library doesn't attempt
            to handle the exception at all, leaving it for the user of
            the library to handle.
        """
        raise NotImplementedError

    def post(self, url, body):
        """
        This performs an HTTP post.  If it makes sense, it will follow
        redirects along the way.

        @return: This returns a three-tuple on success.  The first value
            is the http return code. The second value is the final url
            that was fetched, after following any redirects.  The third
            value is the data that was retrieved from the site.  If the
            fetch didn't succeed, return C{None}.

        @rtype: C{None} or (C{int}, C{str}, C{str})

        @raise Exception: There are also failure conditions where this
            may raise an exception, rather than returning None.  When
            this occurs, the OpenID consumer library doesn't attempt
            to handle the exception at all, leaving it for the user of
            the library to handle.
        """
        raise NotImplementedError

# try to import pycurl, which will let us use ParanoidHTTPClient
try:
    import pycurl
except ImportError:
    pycurl = None

def _getHTTPFetcher(raise_exceptions=False):
    if pycurl is None:
        fetcher = UrllibFetcher()
    else:
        fetcher = ParanoidHTTPFetcher()

    if not raise_exceptions:
        fetcher = ExceptionCatchingFetcher(fetcher)

    return fetcher

def allowedURL(url):
    return url.startswith('http://') or url.startswith('https://')

class ExceptionCatchingFetcher(OpenIDHTTPFetcher):
    def __init__(self, fetcher):
        self.fetcher = fetcher

    def _fetch(self, func, *args):
        try:
            return func(*args)
        except (SystemExit, KeyboardInterrupt):
            raise
        except:
            return None

    def post(self, url, body):
        return self._fetch(self.fetcher.post, url, body)

    def get(self, url):
        return self._fetch(self.fetcher.get, url)

class UrllibFetcher(OpenIDHTTPFetcher):
    """
    A very primitive C{L{OpenIDHTTPFetcher}} that uses urllib2.  Not
    recommended for use if anything else is available.
    """
    def _fetch(self, req):
        try:
            f = urllib2.urlopen(req)
            try:
                data = f.read()
            finally:
                f.close()
        except urllib2.HTTPError, why:
            try:
                data = why.read()
                return (why.code, why.geturl(), data)
            finally:
                why.close()
        else:
            if hasattr(f, 'code'):
                code = f.code
            else:
                code = 200

            return (code, f.geturl(), data)

    def get(self, url):
        if not allowedURL(url):
            raise ValueError('Bad URL scheme: %r' % (url,))

        return self._fetch(url)

    def post(self, url, body):
        if not allowedURL(url):
            raise ValueError('Bad URL scheme: %r' % (url,))

        req = urllib2.Request(url, body)
        return self._fetch(req)

class OpenIDHTTPError(Exception):
    pass

class ParanoidHTTPFetcher(OpenIDHTTPFetcher):
    """
    A paranoid C{L{OpenIDHTTPFetcher}} that uses pycurl for fetching.
    See U{http://pycurl.sourceforge.net/}.
    """
    ALLOWED_TIME = 20 # seconds

    def __init__(self):
        OpenIDHTTPFetcher.__init__(self)
        if pycurl is None:
            raise RuntimeError('Cannot find pycurl library')

    def _findRedirect(self, headers):
        headers.seek(0)
        for line in headers:
            if line.startswith('Location: '):
                return line[9:].strip()
        return None

    def _checkURL(self, url):
        # XXX: make sure url is well-formed and routeable
        return allowedURL(url)

    def get(self, url):
        c = pycurl.Curl()
        try:
            c.setopt(pycurl.NOSIGNAL, 1)

            stop = int(time.time()) + self.ALLOWED_TIME
            off = self.ALLOWED_TIME
            while off > 0:
                if not self._checkURL(url):
                    raise OpenIDHTTPError(
                        "Fetching URL not allowed: %r" % (url,))

                data = cStringIO.StringIO()
                headers = cStringIO.StringIO()
                c.setopt(pycurl.WRITEFUNCTION, data.write)
                c.setopt(pycurl.HEADERFUNCTION, headers.write)
                c.setopt(pycurl.TIMEOUT, off)
                c.setopt(pycurl.URL, url)

                c.perform()

                code = c.getinfo(pycurl.RESPONSE_CODE)
                if code in [301, 302, 303, 307]:
                    url = self._findRedirect(headers)
                else:
                    return code, url, data.getvalue()

                off = stop - int(time.time())

            raise OpenIDHTTPError("Timed out fetching: %r" % (url,))
        finally:
            c.close()

    def post(self, url, body):
        if not self._checkURL(url):
            raise OpenIDHTTPError("Fetching URL not allowed: %r" % (url,))

        c = pycurl.Curl()
        try:
            c.setopt(pycurl.NOSIGNAL, 1)
            c.setopt(pycurl.POST, 1)
            c.setopt(pycurl.POSTFIELDS, body)
            c.setopt(pycurl.TIMEOUT, self.ALLOWED_TIME)
            c.setopt(pycurl.URL, url)

            data = cStringIO.StringIO()
            c.setopt(pycurl.WRITEFUNCTION, data.write)

            c.perform()

            code = c.getinfo(pycurl.RESPONSE_CODE)
            return code, url, data.getvalue()
        finally:
            c.close()
<commit_msg>[project @ Change the name of the allowedURL function]<commit_after>"""
This module contains the HTTP fetcher interface and several implementations.
"""
import urllib2
import time
import cStringIO

class OpenIDHTTPFetcher(object):
    """
    This class is the interface for HTTP fetchers the OpenID consumer
    library uses.  This interface is only important if you need to
    write a new fetcher for some reason.
    """

    def get(self, url):
        """
        This performs an HTTP get, following redirects along the way.

        @return: This returns a three-tuple on success.  The first value
            is the http return code. The second value is the final url
            that was fetched, after following any redirects.  The third
            value is the data that was retrieved from the site.  If the
            fetch didn't succeed, return C{None}.

        @rtype: C{None} or (C{int}, C{str}, C{str})

        @raise Exception: There are also failure conditions where this
            may raise an exception, rather than returning None.  When
            this occurs, the OpenID consumer library doesn't attempt
            to handle the exception at all, leaving it for the user of
            the library to handle.
        """
        raise NotImplementedError

    def post(self, url, body):
        """
        This performs an HTTP post.  If it makes sense, it will follow
        redirects along the way.

        @return: This returns a three-tuple on success.  The first value
            is the http return code. The second value is the final url
            that was fetched, after following any redirects.  The third
            value is the data that was retrieved from the site.  If the
            fetch didn't succeed, return C{None}.

        @rtype: C{None} or (C{int}, C{str}, C{str})

        @raise Exception: There are also failure conditions where this
            may raise an exception, rather than returning None.  When
            this occurs, the OpenID consumer library doesn't attempt
            to handle the exception at all, leaving it for the user of
            the library to handle.
        """
        raise NotImplementedError

# try to import pycurl, which will let us use ParanoidHTTPClient
try:
    import pycurl
except ImportError:
    pycurl = None

def _getHTTPFetcher(raise_exceptions=False):
    if pycurl is None:
        fetcher = UrllibFetcher()
    else:
        fetcher = ParanoidHTTPFetcher()

    if not raise_exceptions:
        fetcher = ExceptionCatchingFetcher(fetcher)

    return fetcher

def _allowedURL(url):
    return url.startswith('http://') or url.startswith('https://')

class ExceptionCatchingFetcher(OpenIDHTTPFetcher):
    def __init__(self, fetcher):
        self.fetcher = fetcher

    def _fetch(self, func, *args):
        try:
            return func(*args)
        except (SystemExit, KeyboardInterrupt):
            raise
        except:
            return None

    def post(self, url, body):
        return self._fetch(self.fetcher.post, url, body)

    def get(self, url):
        return self._fetch(self.fetcher.get, url)

class UrllibFetcher(OpenIDHTTPFetcher):
    """
    A very primitive C{L{OpenIDHTTPFetcher}} that uses urllib2.  Not
    recommended for use if anything else is available.
    """
    def _fetch(self, req):
        try:
            f = urllib2.urlopen(req)
            try:
                data = f.read()
            finally:
                f.close()
        except urllib2.HTTPError, why:
            try:
                data = why.read()
                return (why.code, why.geturl(), data)
            finally:
                why.close()
        else:
            if hasattr(f, 'code'):
                code = f.code
            else:
                code = 200

            return (code, f.geturl(), data)

    def get(self, url):
        if not _allowedURL(url):
            raise ValueError('Bad URL scheme: %r' % (url,))

        return self._fetch(url)

    def post(self, url, body):
        if not _allowedURL(url):
            raise ValueError('Bad URL scheme: %r' % (url,))

        req = urllib2.Request(url, body)
        return self._fetch(req)

class OpenIDHTTPError(Exception):
    pass

class ParanoidHTTPFetcher(OpenIDHTTPFetcher):
    """
    A paranoid C{L{OpenIDHTTPFetcher}} that uses pycurl for fetching.
    See U{http://pycurl.sourceforge.net/}.
    """
    ALLOWED_TIME = 20 # seconds

    def __init__(self):
        OpenIDHTTPFetcher.__init__(self)
        if pycurl is None:
            raise RuntimeError('Cannot find pycurl library')

    def _findRedirect(self, headers):
        headers.seek(0)
        for line in headers:
            if line.startswith('Location: '):
                return line[9:].strip()
        return None

    def _checkURL(self, url):
        # XXX: make sure url is well-formed and routeable
        return _allowedURL(url)

    def get(self, url):
        c = pycurl.Curl()
        try:
            c.setopt(pycurl.NOSIGNAL, 1)

            stop = int(time.time()) + self.ALLOWED_TIME
            off = self.ALLOWED_TIME
            while off > 0:
                if not self._checkURL(url):
                    raise OpenIDHTTPError(
                        "Fetching URL not allowed: %r" % (url,))

                data = cStringIO.StringIO()
                headers = cStringIO.StringIO()
                c.setopt(pycurl.WRITEFUNCTION, data.write)
                c.setopt(pycurl.HEADERFUNCTION, headers.write)
                c.setopt(pycurl.TIMEOUT, off)
                c.setopt(pycurl.URL, url)

                c.perform()

                code = c.getinfo(pycurl.RESPONSE_CODE)
                if code in [301, 302, 303, 307]:
                    url = self._findRedirect(headers)
                else:
                    return code, url, data.getvalue()

                off = stop - int(time.time())

            raise OpenIDHTTPError("Timed out fetching: %r" % (url,))
        finally:
            c.close()

    def post(self, url, body):
        if not self._checkURL(url):
            raise OpenIDHTTPError("Fetching URL not allowed: %r" % (url,))

        c = pycurl.Curl()
        try:
            c.setopt(pycurl.NOSIGNAL, 1)
            c.setopt(pycurl.POST, 1)
            c.setopt(pycurl.POSTFIELDS, body)
            c.setopt(pycurl.TIMEOUT, self.ALLOWED_TIME)
            c.setopt(pycurl.URL, url)

            data = cStringIO.StringIO()
            c.setopt(pycurl.WRITEFUNCTION, data.write)

            c.perform()

            code = c.getinfo(pycurl.RESPONSE_CODE)
            return code, url, data.getvalue()
        finally:
            c.close()
<|endoftext|>


<commit_before>#  -*- coding: utf-8 -*-
#  vim: tabstop=4 shiftwidth=4 softtabstop=4

#  Copyright (c) 2017, GEM Foundation

#  OpenQuake is free software: you can redistribute it and/or modify it
#  under the terms of the GNU Affero General Public License as published
#  by the Free Software Foundation, either version 3 of the License, or
#  (at your option) any later version.

#  OpenQuake is distributed in the hope that it will be useful,
#  but WITHOUT ANY WARRANTY; without even the implied warranty of
#  MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
#  GNU Affero General Public License for more details.

#  You should have received a copy of the GNU Affero General Public License
#  along with OpenQuake.  If not, see <http://www.gnu.org/licenses/>.
import re
import zmq
import logging

context = zmq.Context()

# from integer socket_type to string
SOCKTYPE = {zmq.REQ: 'REQ', zmq.REP: 'REP',
            zmq.PUSH: 'PUSH', zmq.PULL: 'PULL',
            zmq.ROUTER: 'ROUTER', zmq.DEALER: 'DEALER'}

def bind(end_point, socket_type):
    """
    Bind to a zmq URL; raise a proper error if the URL is invalid; return
    a zmq socket.
    """
    sock = context.socket(socket_type)
    try:
        sock.bind(end_point)
    except zmq.error.ZMQError as exc:
        sock.close()
        raise exc.__class__('%s: %s' % (exc, end_point))
    return sock

def connect(end_point, socket_type):
    """
    Connect to a zmq URL; raise a proper error if the URL is invalid; return
    a zmq socket.
    """
    sock = context.socket(socket_type)
    try:
        sock.connect(end_point)
    except zmq.error.ZMQError as exc:
        sock.close()
        raise exc.__class__('%s: %s' % (exc, end_point))
    return sock

class Socket(object):
    """
    A Socket class to be used with code like the following::

     # server
     with Socket('tcp://127.0.0.1:9000', zmq.REP, 'bind') as sock:
         for tup in sock:
             sock.send(tup)

     # client
     with Socket('tcp://127.0.0.1:9000', zmq.REQ, 'connect') as sock:
        sock.send(tup)

    It also support zmq.PULL/zmq.PUSH sockets, which are asynchronous.

    :param end_point: zmq end point string
    :param socket_type: zmq socket type (integer)
    :param mode: default 'bind', accepts also 'connect'
    :param timeout: default 5000 ms, used when polling the underlying socket
    """
    def __init__(self, end_point, socket_type, mode, timeout=5000):
        assert 'localhost' not in end_point, 'Use 127.0.0.1 instead'
        assert socket_type in (zmq.REP, zmq.REQ, zmq.PULL, zmq.PUSH)
        assert mode in ('bind', 'connect'), mode
        self.end_point = end_point
        self.socket_type = socket_type
        self.mode = mode
        self.timeout = timeout
        self.running = False

    def __enter__(self):
        """Instantiate the underlying zmq socket"""
        # first check if the end_point ends in :<min_port>-<max_port>
        port_range = re.search(r':(\d+)-(\d+)$', self.end_point)
        if port_range:
            assert self.mode == 'bind', self.mode
            p1, p2 = map(int, port_range.groups())
            end_point = self.end_point.rsplit(':', 1)[0]  # strip port range
            self.zsocket = context.socket(self.socket_type)
            port = self.zsocket.bind_to_random_port(end_point, p1, p2)
            # NB: will raise a ZMQBindError if no port is available
            self.port = port
            self.backurl = '%s:%d' % (end_point, port)
        elif self.mode == 'bind':
            self.zsocket = bind(self.end_point, self.socket_type)
            self.backurl = self.end_point
        else:  # connect
            self.zsocket = connect(self.end_point, self.socket_type)
        port = re.search(r':(\d+)$', self.end_point)
        if port:
            self.port = int(port.group(1))
        self.zsocket.__enter__()
        self.num_sent = 0
        return self

    def __exit__(self, *args):
        self.zsocket.__exit__(*args)
        del self.zsocket

    def __iter__(self):
        """
        Iterate on the socket and yield the received arguments. Exits if

        1. the flag .running is set to False
        2. the message 'stop' is sent
        3. SIGINT is sent
        4. SIGTERM is sent
        """
        # works with zmq.REP and zmq.PULL sockets
        self.running = True
        while self.running:
            try:
                if self.zsocket.poll(self.timeout):
                    args = self.zsocket.recv_pyobj()
                else:
                    logging.warn('Timeout in %s', self)
                    continue
            except (KeyboardInterrupt, zmq.ZMQError):
                # sending SIGTERM raises ZMQError
                break
            if args == 'stop':
                if self.socket_type == zmq.REP:
                    self.send((None, None, None))
                break
            else:
                yield args

    def send(self, obj):
        """
        Send an object to the remote server; block and return the reply
        if the socket type is REQ.

        :param obj:
            the Python object to send
        """
        self.zsocket.send_pyobj(obj)
        self.num_sent += 1
        if self.socket_type == zmq.REQ:
            return self.zsocket.recv_pyobj()

    def __repr__(self):
        end_point = getattr(self, 'backurl', self.end_point)
        return '<%s %s %s %s>' % (self.__class__.__name__, end_point,
                                  SOCKTYPE[self.socket_type], self.mode)
<commit_msg>Refined warning [skip hazardlib][demos]<commit_after>#  -*- coding: utf-8 -*-
#  vim: tabstop=4 shiftwidth=4 softtabstop=4

#  Copyright (c) 2017, GEM Foundation

#  OpenQuake is free software: you can redistribute it and/or modify it
#  under the terms of the GNU Affero General Public License as published
#  by the Free Software Foundation, either version 3 of the License, or
#  (at your option) any later version.

#  OpenQuake is distributed in the hope that it will be useful,
#  but WITHOUT ANY WARRANTY; without even the implied warranty of
#  MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
#  GNU Affero General Public License for more details.

#  You should have received a copy of the GNU Affero General Public License
#  along with OpenQuake.  If not, see <http://www.gnu.org/licenses/>.
import re
import zmq
import logging

context = zmq.Context()

# from integer socket_type to string
SOCKTYPE = {zmq.REQ: 'REQ', zmq.REP: 'REP',
            zmq.PUSH: 'PUSH', zmq.PULL: 'PULL',
            zmq.ROUTER: 'ROUTER', zmq.DEALER: 'DEALER'}

def bind(end_point, socket_type):
    """
    Bind to a zmq URL; raise a proper error if the URL is invalid; return
    a zmq socket.
    """
    sock = context.socket(socket_type)
    try:
        sock.bind(end_point)
    except zmq.error.ZMQError as exc:
        sock.close()
        raise exc.__class__('%s: %s' % (exc, end_point))
    return sock

def connect(end_point, socket_type):
    """
    Connect to a zmq URL; raise a proper error if the URL is invalid; return
    a zmq socket.
    """
    sock = context.socket(socket_type)
    try:
        sock.connect(end_point)
    except zmq.error.ZMQError as exc:
        sock.close()
        raise exc.__class__('%s: %s' % (exc, end_point))
    return sock

class Socket(object):
    """
    A Socket class to be used with code like the following::

     # server
     with Socket('tcp://127.0.0.1:9000', zmq.REP, 'bind') as sock:
         for tup in sock:
             sock.send(tup)

     # client
     with Socket('tcp://127.0.0.1:9000', zmq.REQ, 'connect') as sock:
        sock.send(tup)

    It also support zmq.PULL/zmq.PUSH sockets, which are asynchronous.

    :param end_point: zmq end point string
    :param socket_type: zmq socket type (integer)
    :param mode: default 'bind', accepts also 'connect'
    :param timeout: default 5000 ms, used when polling the underlying socket
    """
    def __init__(self, end_point, socket_type, mode, timeout=5000):
        assert 'localhost' not in end_point, 'Use 127.0.0.1 instead'
        assert socket_type in (zmq.REP, zmq.REQ, zmq.PULL, zmq.PUSH)
        assert mode in ('bind', 'connect'), mode
        self.end_point = end_point
        self.socket_type = socket_type
        self.mode = mode
        self.timeout = timeout
        self.running = False

    def __enter__(self):
        """Instantiate the underlying zmq socket"""
        # first check if the end_point ends in :<min_port>-<max_port>
        port_range = re.search(r':(\d+)-(\d+)$', self.end_point)
        if port_range:
            assert self.mode == 'bind', self.mode
            p1, p2 = map(int, port_range.groups())
            end_point = self.end_point.rsplit(':', 1)[0]  # strip port range
            self.zsocket = context.socket(self.socket_type)
            port = self.zsocket.bind_to_random_port(end_point, p1, p2)
            # NB: will raise a ZMQBindError if no port is available
            self.port = port
            self.backurl = '%s:%d' % (end_point, port)
        elif self.mode == 'bind':
            self.zsocket = bind(self.end_point, self.socket_type)
            self.backurl = self.end_point
        else:  # connect
            self.zsocket = connect(self.end_point, self.socket_type)
        port = re.search(r':(\d+)$', self.end_point)
        if port:
            self.port = int(port.group(1))
        self.zsocket.__enter__()
        self.num_sent = 0
        return self

    def __exit__(self, *args):
        self.zsocket.__exit__(*args)
        del self.zsocket

    def __iter__(self):
        """
        Iterate on the socket and yield the received arguments. Exits if

        1. the flag .running is set to False
        2. the message 'stop' is sent
        3. SIGINT is sent
        4. SIGTERM is sent
        """
        # works with zmq.REP and zmq.PULL sockets
        self.running = True
        while self.running:
            try:
                if self.zsocket.poll(self.timeout):
                    args = self.zsocket.recv_pyobj()
                else:
                    # wait a bit more; print a warning for PULL sockets
                    if self.socket_type == 'PULL':
                        logging.warn('Timeout in %s', self)
                    continue
            except (KeyboardInterrupt, zmq.ZMQError):
                # sending SIGTERM raises ZMQError
                break
            if args == 'stop':
                if self.socket_type == zmq.REP:
                    self.send((None, None, None))
                break
            else:
                yield args

    def send(self, obj):
        """
        Send an object to the remote server; block and return the reply
        if the socket type is REQ.

        :param obj:
            the Python object to send
        """
        self.zsocket.send_pyobj(obj)
        self.num_sent += 1
        if self.socket_type == zmq.REQ:
            return self.zsocket.recv_pyobj()

    def __repr__(self):
        end_point = getattr(self, 'backurl', self.end_point)
        return '<%s %s %s %s>' % (self.__class__.__name__, end_point,
                                  SOCKTYPE[self.socket_type], self.mode)
<|endoftext|>


<commit_before>#  -*- coding: utf-8 -*-
#  vim: tabstop=4 shiftwidth=4 softtabstop=4

#  Copyright (c) 2016, GEM Foundation

#  OpenQuake is free software: you can redistribute it and/or modify it
#  under the terms of the GNU Affero General Public License as published
#  by the Free Software Foundation, either version 3 of the License, or
#  (at your option) any later version.

#  OpenQuake is distributed in the hope that it will be useful,
#  but WITHOUT ANY WARRANTY; without even the implied warranty of
#  MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
#  GNU Affero General Public License for more details.

#  You should have received a copy of the GNU Affero General Public License
#  along with OpenQuake.  If not, see <http://www.gnu.org/licenses/>.
import sys
import subprocess
from time import sleep
from openquake.risklib import valid
from openquake.baselib import sap
from openquake.engine import config
from openquake.server.dbserver import get_status

def rundjango(subcmd, hostport=None):
    args = [sys.executable, '-m', 'openquake.server.manage', subcmd]
    if hostport:
        args.append(hostport)
    subprocess.call(args)

@sap.Script
def webui(cmd, hostport='127.0.0.1:8800'):
    """
    start the webui server in foreground or perform other operation on the
    django application
    """
    dbstatus = get_status()
    if dbstatus == 'not-running':
        if valid.boolean(config.get('dbserver', 'multi_user')):
            sys.exit('Please start the DbServer: '
                     'see the documentation for details')
        subprocess.Popen([sys.executable, '-m', 'openquake.server.dbserver'])
        counter = 0
        while dbstatus == 'not-running':
            if counter >= 10:
                sys.exit('The DbServer cannot be started. '
                         'Please check the configuration')
            sleep(1)
            dbstatus = get_status()
            counter += 1

    if cmd == 'start':
        rundjango('runserver', hostport)
    elif cmd == 'syncdb':
        rundjango('syncdb')

webui.arg('cmd', 'webui command', choices='start syncdb'.split())
webui.arg('hostport', 'a string of the form <hostname:port>')
<commit_msg>Reduce the timeout to 5 secs<commit_after>#  -*- coding: utf-8 -*-
#  vim: tabstop=4 shiftwidth=4 softtabstop=4

#  Copyright (c) 2016, GEM Foundation

#  OpenQuake is free software: you can redistribute it and/or modify it
#  under the terms of the GNU Affero General Public License as published
#  by the Free Software Foundation, either version 3 of the License, or
#  (at your option) any later version.

#  OpenQuake is distributed in the hope that it will be useful,
#  but WITHOUT ANY WARRANTY; without even the implied warranty of
#  MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
#  GNU Affero General Public License for more details.

#  You should have received a copy of the GNU Affero General Public License
#  along with OpenQuake.  If not, see <http://www.gnu.org/licenses/>.
import sys
import subprocess
from time import sleep
from openquake.risklib import valid
from openquake.baselib import sap
from openquake.engine import config
from openquake.server.dbserver import get_status

def rundjango(subcmd, hostport=None):
    args = [sys.executable, '-m', 'openquake.server.manage', subcmd]
    if hostport:
        args.append(hostport)
    subprocess.call(args)

@sap.Script
def webui(cmd, hostport='127.0.0.1:8800'):
    """
    start the webui server in foreground or perform other operation on the
    django application
    """
    dbstatus = get_status()
    if dbstatus == 'not-running':
        if valid.boolean(config.get('dbserver', 'multi_user')):
            sys.exit('Please start the DbServer: '
                     'see the documentation for details')
        subprocess.Popen([sys.executable, '-m', 'openquake.server.dbserver'])
        counter = 0
        while dbstatus == 'not-running':
            if counter >= 5:
                sys.exit('The DbServer cannot be started. '
                         'Please check the configuration')
            sleep(1)
            dbstatus = get_status()
            counter += 1

    if cmd == 'start':
        rundjango('runserver', hostport)
    elif cmd == 'syncdb':
        rundjango('syncdb')

webui.arg('cmd', 'webui command', choices='start syncdb'.split())
webui.arg('hostport', 'a string of the form <hostname:port>')
<|endoftext|>


<commit_before># -*- coding: utf-8 -*-
# vim: tabstop=4 shiftwidth=4 softtabstop=4
#
# Copyright (C) 2014-2017 GEM Foundation
#
# OpenQuake is free software: you can redistribute it and/or modify it
# under the terms of the GNU Affero General Public License as published
# by the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
# OpenQuake is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU Affero General Public License for more details.
#
# You should have received a copy of the GNU Affero General Public License
# along with OpenQuake. If not, see <http://www.gnu.org/licenses/>.
"""\
It is possible to save a Node object into a NRML file by using the
function ``write(nodes, output)`` where output is a file
object. If you want to make sure that the generated file is valid
according to the NRML schema just open it in 'w+' mode: immediately
after writing it will be read and validated. It is also possible to
convert a NRML file into a Node object with the routine
``read(node, input)`` where input is the path name of the
NRML file or a file object opened for reading. The file will be
validated as soon as opened.

For instance an exposure file like the following::

  <?xml version='1.0' encoding='utf-8'?>
  <nrml xmlns="http://openquake.org/xmlns/nrml/0.4"
        xmlns:gml="http://www.opengis.net/gml">
    <exposureModel
        id="my_exposure_model_for_population"
        category="population"
        taxonomySource="fake population datasource">

      <description>
        Sample population
      </description>

      <assets>
        <asset id="asset_01" number="7" taxonomy="IT-PV">
            <location lon="9.15000" lat="45.16667" />
        </asset>

        <asset id="asset_02" number="7" taxonomy="IT-CE">
            <location lon="9.15333" lat="45.12200" />
        </asset>
      </assets>
    </exposureModel>
  </nrml>

can be converted as follows:

>> nrml = read(<path_to_the_exposure_file.xml>)

Then subnodes and attributes can be conveniently accessed:

>> nrml.exposureModel.assets[0]['taxonomy']
'IT-PV'
>> nrml.exposureModel.assets[0]['id']
'asset_01'
>> nrml.exposureModel.assets[0].location['lon']
'9.15000'
>> nrml.exposureModel.assets[0].location['lat']
'45.16667'

The Node class provides no facility to cast strings into Python types;
this is a job for the Node class which can be subclassed and
supplemented by a dictionary of validators.
"""
from __future__ import print_function
import io
import re
import sys
import copy
import decimal
import logging
import operator
import collections

import numpy

from openquake.baselib.general import CallableDict, groupby
from openquake.baselib.node import (
    node_to_xml, context, Node, striptag, ValidatingXmlParser, floatformat)
from openquake.hazardlib import valid, sourceconverter

F64 = numpy.float64
NAMESPACE = 'http://openquake.org/xmlns/nrml/0.4'
NRML05 = 'http://openquake.org/xmlns/nrml/0.5'
GML_NAMESPACE = 'http://www.opengis.net/gml'
SERIALIZE_NS_MAP = {None: NAMESPACE, 'gml': GML_NAMESPACE}
PARSE_NS_MAP = {'nrml': NAMESPACE, 'gml': GML_NAMESPACE}

class DuplicatedID(Exception):
    """Raised when two sources with the same ID are found in a source model"""

def get_tag_version(nrml_node):
    """
    Extract from a node of kind NRML the tag and the version. For instance
    from '{http://openquake.org/xmlns/nrml/0.4}fragilityModel' one gets
    the pair ('fragilityModel', 'nrml/0.4').
    """
    version, tag = re.search(r'(nrml/[\d\.]+)\}(\w+)', nrml_node.tag).groups()
    return tag, version

def parse(fname, *args):
    """
    Parse a NRML file and return an associated Python object. It works by
    calling nrml.read() and node_to_obj() in sequence.
    """
    [node] = read(fname)
    return node_to_obj(node, fname, *args)

node_to_obj = CallableDict(keyfunc=get_tag_version, keymissing=lambda n, f: n)
# dictionary of functions with at least two arguments, node and fname

@node_to_obj.add(('sourceModel', 'nrml/0.4'))
def get_source_model_04(node, fname, converter):
    sources = []
    source_ids = set()
    converter.fname = fname
    for no, src_node in enumerate(node, 1):
        src = converter.convert_node(src_node)
        if src.source_id in source_ids:
            raise DuplicatedID(
                'The source ID %s is duplicated!' % src.source_id)
        sources.append(src)
        source_ids.add(src.source_id)
        if no % 10000 == 0:  # log every 10,000 sources parsed
            logging.info('Instantiated %d sources from %s', no, fname)
    if no % 10000 != 0:
        logging.info('Instantiated %d sources from %s', no, fname)
    groups = groupby(
        sources, operator.attrgetter('tectonic_region_type'))
    return sorted(sourceconverter.SourceGroup(trt, srcs)
                  for trt, srcs in groups.items())

@node_to_obj.add(('sourceModel', 'nrml/0.5'))
def get_source_model_05(node, fname, converter):
    converter.fname = fname
    groups = []  # expect a sequence of sourceGroup nodes
    for src_group in node:
        with context(fname, src_group):
            if 'sourceGroup' not in src_group.tag:
                raise ValueError('expected sourceGroup')
        groups.append(converter.convert_node(src_group))
    return sorted(groups)

validators = {
    'strike': valid.strike_range,
    'dip': valid.dip_range,
    'rake': valid.rake_range,
    'magnitude': valid.positivefloat,
    'lon': valid.longitude,
    'lat': valid.latitude,
    'depth': valid.float_,
    'upperSeismoDepth': valid.float_,
    'lowerSeismoDepth': valid.float_,
    'posList': valid.posList,
    'pos': valid.lon_lat,
    'aValue': float,
    'bValue': valid.positivefloat,
    'magScaleRel': valid.mag_scale_rel,
    'tectonicRegion': str,
    'ruptAspectRatio': valid.positivefloat,
    'maxMag': valid.positivefloat,
    'minMag': valid.positivefloat,
    'binWidth': valid.positivefloat,
    'probability': valid.probability,
    'occurRates': valid.probabilities,
    'probs_occur': valid.pmf,
    'weight': valid.probability,
    'uncertaintyWeight': decimal.Decimal,
    'alongStrike': valid.probability,
    'downDip': valid.probability,
    'totalMomentRate': valid.positivefloat,
    'characteristicRate': valid.positivefloat,
    'characteristicMag': valid.positivefloat,
    'magnitudes': valid.positivefloats,
    'id': valid.simple_id,
    'rupture.id': valid.utf8,  # event tag
    'discretization': valid.compose(valid.positivefloat, valid.nonzero),
    'IML': valid.positivefloats,  # used in NRML 0.4
    'imt': valid.intensity_measure_type,
    'imls': valid.positivefloats,
    'poes': valid.positivefloats,
    'description': valid.utf8_not_empty,
    'noDamageLimit': valid.NoneOr(valid.positivefloat),
    'investigationTime': valid.positivefloat,
    'poEs': valid.probabilities,
    'gsimTreePath': lambda v: v.split('_'),
    'sourceModelTreePath': lambda v: v.split('_'),
    'poE': valid.probability,
    'IMLs': valid.positivefloats,
    'pos': valid.lon_lat,
    'IMT': str,
    'saPeriod': valid.positivefloat,
    'saDamping': valid.positivefloat,
    'quantileValue': valid.positivefloat,
    'investigationTime': valid.positivefloat,
    'poE': valid.probability,
    'periods': valid.positivefloats,
    'pos': valid.lon_lat,
    'IMLs': valid.positivefloats,
    'lon': valid.longitude,
    'lat': valid.latitude,
    'magBinEdges': valid.integers,
    'distBinEdges': valid.integers,
    'epsBinEdges': valid.integers,
    'lonBinEdges': valid.longitudes,
    'latBinEdges': valid.latitudes,
    'type': valid.namelist,
    'dims': valid.positiveints,
    'poE': valid.probability,
    'iml': valid.positivefloat,
    'index': valid.positiveints,
    'value': valid.positivefloat,
    'assetLifeExpectancy': valid.positivefloat,
    'interestRate': valid.positivefloat,
    'statistics': valid.Choice('mean', 'quantile'),
    'pos': valid.lon_lat,
    'gmv': valid.positivefloat,
    'spacing': valid.positivefloat,
    'srcs_weights': valid.positivefloats,
}

class SourceModelParser(object):
    """
    A source model parser featuring a cache.

    :param converter:
        :class:`openquake.commonlib.source.SourceConverter` instance
    """
    def __init__(self, converter):
        self.converter = converter
        self.groups = {}  # cache fname -> groups
        self.fname_hits = collections.Counter()  # fname -> number of calls

    def parse_src_groups(self, fname, apply_uncertainties=None):
        """
        :param fname:
            the full pathname of the source model file
        :param apply_uncertainties:
            a function modifying the sources (or None)
        """
        try:
            groups = self.groups[fname]
        except KeyError:
            groups = self.groups[fname] = self.parse_groups(fname)
        # NB: deepcopy is *essential* here
        groups = [copy.deepcopy(g) for g in groups]
        for group in groups:
            nrup = 0
            for src in group:
                if apply_uncertainties:
                    apply_uncertainties(src)
                    src.num_ruptures = src.count_ruptures()
                    nrup += src.num_ruptures
            # NB: if the user sets a wrong discretization parameter
            # the call to `.count_ruptures()` can be ultra-slow
            logging.debug("%s, %s: parsed %d source(s) with %d ruptures",
                          fname, group.trt, len(group), nrup)
        self.fname_hits[fname] += 1
        return groups

    def parse_groups(self, fname):
        """
        Parse all the groups and return them ordered by number of sources.
        It does not count the ruptures, so it is relatively fast.

        :param fname:
            the full pathname of the source model file
        """
        return parse(fname, self.converter)

def read(source, chatty=True, stop=None):
    """
    Convert a NRML file into a validated Node object. Keeps
    the entire tree in memory.

    :param source:
        a file name or file object open for reading
    """
    vparser = ValidatingXmlParser(validators, stop)
    nrml = vparser.parse_file(source)
    if striptag(nrml.tag) != 'nrml':
        raise ValueError('%s: expected a node of kind nrml, got %s' %
                         (source, nrml.tag))
    # extract the XML namespace URL ('http://openquake.org/xmlns/nrml/0.5')
    xmlns = nrml.tag.split('}')[0][1:]
    if xmlns != NRML05 and chatty:
        # for the moment NRML04 is still supported, so we hide the warning
        logging.debug('%s is at an outdated version: %s', source, xmlns)
    nrml['xmlns'] = xmlns
    nrml['xmlns:gml'] = GML_NAMESPACE
    return nrml

def write(nodes, output=sys.stdout, fmt='%.7E', gml=True, xmlns=None):
    """
    Convert nodes into a NRML file. output must be a file
    object open in write mode. If you want to perform a
    consistency check, open it in read-write mode, then it will
    be read after creation and validated.

    :params nodes: an iterable over Node objects
    :params output: a file-like object in write or read-write mode
    :param fmt: format used for writing the floats (default '%.7E')
    :param gml: add the http://www.opengis.net/gml namespace
    :param xmlns: NRML namespace like http://openquake.org/xmlns/nrml/0.4
    """
    root = Node('nrml', nodes=nodes)
    namespaces = {xmlns or NRML05: ''}
    if gml:
        namespaces[GML_NAMESPACE] = 'gml:'
    with floatformat(fmt):
        node_to_xml(root, output, namespaces)
    if hasattr(output, 'mode') and '+' in output.mode:  # read-write mode
        output.seek(0)
        read(output)  # validate the written file

def convert(node):
    """
    Convert a node into a string in NRML format
    """
    with io.BytesIO() as f:
        write([node], f)
        return f.getvalue().decode('utf-8')

if __name__ == '__main__':
    import sys
    for fname in sys.argv[1:]:
        print('****** %s ******' % fname)
        print(read(fname).to_str())
        print()
<commit_msg>Fixed occurRates<commit_after># -*- coding: utf-8 -*-
# vim: tabstop=4 shiftwidth=4 softtabstop=4
#
# Copyright (C) 2014-2017 GEM Foundation
#
# OpenQuake is free software: you can redistribute it and/or modify it
# under the terms of the GNU Affero General Public License as published
# by the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
# OpenQuake is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU Affero General Public License for more details.
#
# You should have received a copy of the GNU Affero General Public License
# along with OpenQuake. If not, see <http://www.gnu.org/licenses/>.
"""\
It is possible to save a Node object into a NRML file by using the
function ``write(nodes, output)`` where output is a file
object. If you want to make sure that the generated file is valid
according to the NRML schema just open it in 'w+' mode: immediately
after writing it will be read and validated. It is also possible to
convert a NRML file into a Node object with the routine
``read(node, input)`` where input is the path name of the
NRML file or a file object opened for reading. The file will be
validated as soon as opened.

For instance an exposure file like the following::

  <?xml version='1.0' encoding='utf-8'?>
  <nrml xmlns="http://openquake.org/xmlns/nrml/0.4"
        xmlns:gml="http://www.opengis.net/gml">
    <exposureModel
        id="my_exposure_model_for_population"
        category="population"
        taxonomySource="fake population datasource">

      <description>
        Sample population
      </description>

      <assets>
        <asset id="asset_01" number="7" taxonomy="IT-PV">
            <location lon="9.15000" lat="45.16667" />
        </asset>

        <asset id="asset_02" number="7" taxonomy="IT-CE">
            <location lon="9.15333" lat="45.12200" />
        </asset>
      </assets>
    </exposureModel>
  </nrml>

can be converted as follows:

>> nrml = read(<path_to_the_exposure_file.xml>)

Then subnodes and attributes can be conveniently accessed:

>> nrml.exposureModel.assets[0]['taxonomy']
'IT-PV'
>> nrml.exposureModel.assets[0]['id']
'asset_01'
>> nrml.exposureModel.assets[0].location['lon']
'9.15000'
>> nrml.exposureModel.assets[0].location['lat']
'45.16667'

The Node class provides no facility to cast strings into Python types;
this is a job for the Node class which can be subclassed and
supplemented by a dictionary of validators.
"""
from __future__ import print_function
import io
import re
import sys
import copy
import decimal
import logging
import operator
import collections

import numpy

from openquake.baselib.general import CallableDict, groupby
from openquake.baselib.node import (
    node_to_xml, context, Node, striptag, ValidatingXmlParser, floatformat)
from openquake.hazardlib import valid, sourceconverter

F64 = numpy.float64
NAMESPACE = 'http://openquake.org/xmlns/nrml/0.4'
NRML05 = 'http://openquake.org/xmlns/nrml/0.5'
GML_NAMESPACE = 'http://www.opengis.net/gml'
SERIALIZE_NS_MAP = {None: NAMESPACE, 'gml': GML_NAMESPACE}
PARSE_NS_MAP = {'nrml': NAMESPACE, 'gml': GML_NAMESPACE}

class DuplicatedID(Exception):
    """Raised when two sources with the same ID are found in a source model"""

def get_tag_version(nrml_node):
    """
    Extract from a node of kind NRML the tag and the version. For instance
    from '{http://openquake.org/xmlns/nrml/0.4}fragilityModel' one gets
    the pair ('fragilityModel', 'nrml/0.4').
    """
    version, tag = re.search(r'(nrml/[\d\.]+)\}(\w+)', nrml_node.tag).groups()
    return tag, version

def parse(fname, *args):
    """
    Parse a NRML file and return an associated Python object. It works by
    calling nrml.read() and node_to_obj() in sequence.
    """
    [node] = read(fname)
    return node_to_obj(node, fname, *args)

node_to_obj = CallableDict(keyfunc=get_tag_version, keymissing=lambda n, f: n)
# dictionary of functions with at least two arguments, node and fname

@node_to_obj.add(('sourceModel', 'nrml/0.4'))
def get_source_model_04(node, fname, converter):
    sources = []
    source_ids = set()
    converter.fname = fname
    for no, src_node in enumerate(node, 1):
        src = converter.convert_node(src_node)
        if src.source_id in source_ids:
            raise DuplicatedID(
                'The source ID %s is duplicated!' % src.source_id)
        sources.append(src)
        source_ids.add(src.source_id)
        if no % 10000 == 0:  # log every 10,000 sources parsed
            logging.info('Instantiated %d sources from %s', no, fname)
    if no % 10000 != 0:
        logging.info('Instantiated %d sources from %s', no, fname)
    groups = groupby(
        sources, operator.attrgetter('tectonic_region_type'))
    return sorted(sourceconverter.SourceGroup(trt, srcs)
                  for trt, srcs in groups.items())

@node_to_obj.add(('sourceModel', 'nrml/0.5'))
def get_source_model_05(node, fname, converter):
    converter.fname = fname
    groups = []  # expect a sequence of sourceGroup nodes
    for src_group in node:
        with context(fname, src_group):
            if 'sourceGroup' not in src_group.tag:
                raise ValueError('expected sourceGroup')
        groups.append(converter.convert_node(src_group))
    return sorted(groups)

validators = {
    'strike': valid.strike_range,
    'dip': valid.dip_range,
    'rake': valid.rake_range,
    'magnitude': valid.positivefloat,
    'lon': valid.longitude,
    'lat': valid.latitude,
    'depth': valid.float_,
    'upperSeismoDepth': valid.float_,
    'lowerSeismoDepth': valid.float_,
    'posList': valid.posList,
    'pos': valid.lon_lat,
    'aValue': float,
    'bValue': valid.positivefloat,
    'magScaleRel': valid.mag_scale_rel,
    'tectonicRegion': str,
    'ruptAspectRatio': valid.positivefloat,
    'maxMag': valid.positivefloat,
    'minMag': valid.positivefloat,
    'binWidth': valid.positivefloat,
    'probability': valid.probability,
    'occurRates': valid.positivefloats,  # they can be > 1
    'probs_occur': valid.pmf,
    'weight': valid.probability,
    'uncertaintyWeight': decimal.Decimal,
    'alongStrike': valid.probability,
    'downDip': valid.probability,
    'totalMomentRate': valid.positivefloat,
    'characteristicRate': valid.positivefloat,
    'characteristicMag': valid.positivefloat,
    'magnitudes': valid.positivefloats,
    'id': valid.simple_id,
    'rupture.id': valid.utf8,  # event tag
    'discretization': valid.compose(valid.positivefloat, valid.nonzero),
    'IML': valid.positivefloats,  # used in NRML 0.4
    'imt': valid.intensity_measure_type,
    'imls': valid.positivefloats,
    'poes': valid.positivefloats,
    'description': valid.utf8_not_empty,
    'noDamageLimit': valid.NoneOr(valid.positivefloat),
    'investigationTime': valid.positivefloat,
    'poEs': valid.probabilities,
    'gsimTreePath': lambda v: v.split('_'),
    'sourceModelTreePath': lambda v: v.split('_'),
    'poE': valid.probability,
    'IMLs': valid.positivefloats,
    'pos': valid.lon_lat,
    'IMT': str,
    'saPeriod': valid.positivefloat,
    'saDamping': valid.positivefloat,
    'quantileValue': valid.positivefloat,
    'investigationTime': valid.positivefloat,
    'poE': valid.probability,
    'periods': valid.positivefloats,
    'pos': valid.lon_lat,
    'IMLs': valid.positivefloats,
    'lon': valid.longitude,
    'lat': valid.latitude,
    'magBinEdges': valid.integers,
    'distBinEdges': valid.integers,
    'epsBinEdges': valid.integers,
    'lonBinEdges': valid.longitudes,
    'latBinEdges': valid.latitudes,
    'type': valid.namelist,
    'dims': valid.positiveints,
    'poE': valid.probability,
    'iml': valid.positivefloat,
    'index': valid.positiveints,
    'value': valid.positivefloat,
    'assetLifeExpectancy': valid.positivefloat,
    'interestRate': valid.positivefloat,
    'statistics': valid.Choice('mean', 'quantile'),
    'pos': valid.lon_lat,
    'gmv': valid.positivefloat,
    'spacing': valid.positivefloat,
    'srcs_weights': valid.positivefloats,
}

class SourceModelParser(object):
    """
    A source model parser featuring a cache.

    :param converter:
        :class:`openquake.commonlib.source.SourceConverter` instance
    """
    def __init__(self, converter):
        self.converter = converter
        self.groups = {}  # cache fname -> groups
        self.fname_hits = collections.Counter()  # fname -> number of calls

    def parse_src_groups(self, fname, apply_uncertainties=None):
        """
        :param fname:
            the full pathname of the source model file
        :param apply_uncertainties:
            a function modifying the sources (or None)
        """
        try:
            groups = self.groups[fname]
        except KeyError:
            groups = self.groups[fname] = self.parse_groups(fname)
        # NB: deepcopy is *essential* here
        groups = [copy.deepcopy(g) for g in groups]
        for group in groups:
            nrup = 0
            for src in group:
                if apply_uncertainties:
                    apply_uncertainties(src)
                    src.num_ruptures = src.count_ruptures()
                    nrup += src.num_ruptures
            # NB: if the user sets a wrong discretization parameter
            # the call to `.count_ruptures()` can be ultra-slow
            logging.debug("%s, %s: parsed %d source(s) with %d ruptures",
                          fname, group.trt, len(group), nrup)
        self.fname_hits[fname] += 1
        return groups

    def parse_groups(self, fname):
        """
        Parse all the groups and return them ordered by number of sources.
        It does not count the ruptures, so it is relatively fast.

        :param fname:
            the full pathname of the source model file
        """
        return parse(fname, self.converter)

def read(source, chatty=True, stop=None):
    """
    Convert a NRML file into a validated Node object. Keeps
    the entire tree in memory.

    :param source:
        a file name or file object open for reading
    """
    vparser = ValidatingXmlParser(validators, stop)
    nrml = vparser.parse_file(source)
    if striptag(nrml.tag) != 'nrml':
        raise ValueError('%s: expected a node of kind nrml, got %s' %
                         (source, nrml.tag))
    # extract the XML namespace URL ('http://openquake.org/xmlns/nrml/0.5')
    xmlns = nrml.tag.split('}')[0][1:]
    if xmlns != NRML05 and chatty:
        # for the moment NRML04 is still supported, so we hide the warning
        logging.debug('%s is at an outdated version: %s', source, xmlns)
    nrml['xmlns'] = xmlns
    nrml['xmlns:gml'] = GML_NAMESPACE
    return nrml

def write(nodes, output=sys.stdout, fmt='%.7E', gml=True, xmlns=None):
    """
    Convert nodes into a NRML file. output must be a file
    object open in write mode. If you want to perform a
    consistency check, open it in read-write mode, then it will
    be read after creation and validated.

    :params nodes: an iterable over Node objects
    :params output: a file-like object in write or read-write mode
    :param fmt: format used for writing the floats (default '%.7E')
    :param gml: add the http://www.opengis.net/gml namespace
    :param xmlns: NRML namespace like http://openquake.org/xmlns/nrml/0.4
    """
    root = Node('nrml', nodes=nodes)
    namespaces = {xmlns or NRML05: ''}
    if gml:
        namespaces[GML_NAMESPACE] = 'gml:'
    with floatformat(fmt):
        node_to_xml(root, output, namespaces)
    if hasattr(output, 'mode') and '+' in output.mode:  # read-write mode
        output.seek(0)
        read(output)  # validate the written file

def convert(node):
    """
    Convert a node into a string in NRML format
    """
    with io.BytesIO() as f:
        write([node], f)
        return f.getvalue().decode('utf-8')

if __name__ == '__main__':
    import sys
    for fname in sys.argv[1:]:
        print('****** %s ******' % fname)
        print(read(fname).to_str())
        print()
<|endoftext|>