<commit_before>"""
Create the tfrecord files for a dataset.

A lot of this code comes from the tensorflow inception example, so here is their license:

# Copyright 2016 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.
# ==============================================================================

"""
from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

from datetime import datetime
import os
from Queue import Queue
import random
import sys
import threading

import numpy as np
import tensorflow as tf

def _int64_feature(value):
    """Wrapper for inserting int64 features into Example proto."""
    if not isinstance(value, list):
        value = [value]
    return tf.train.Feature(int64_list=tf.train.Int64List(value=value))

def _float_feature(value):
    """Wrapper for inserting float features into Example proto."""
    if not isinstance(value, list):
        value = [value]
    return tf.train.Feature(float_list=tf.train.FloatList(value=value))

def _bytes_feature(value):
    """Wrapper for inserting bytes features into Example proto."""
    if not isinstance(value, list):
        value = [value]
    return tf.train.Feature(bytes_list=tf.train.BytesList(value=value))

def _convert_to_example(image_example, image_buffer, height, width):
    """Build an Example proto for an example.
    Args:
      image_example: dict, an image example
      image_buffer: string, JPEG encoding of RGB image
      height: integer, image height in pixels
      width: integer, image width in pixels
    Returns:
      Example proto
    """

    # Required
    filename = str(image_example['filename'])
    image_id = str(image_example['id'])

    # Not required
    fileurl = str(image_example.get('fileurl', ''))

    # Class label for the whole image
    image_class = image_example.get('class', {})
    class_label = image_class.get('label', 0)
    class_text = str(image_class.get('text', ''))

    # Objects
    image_objects = image_example.get('object', {})
    object_count = image_objects.get('count', 0)

    # Bounding Boxes
    image_bboxes = image_objects.get('bbox', {})
    xmin = image_bboxes.get('xmin', [])
    xmax = image_bboxes.get('xmax', [])
    ymin = image_bboxes.get('ymin', [])
    ymax = image_bboxes.get('ymax', [])
    bbox_labels = image_bboxes.get('label', [])
    bbox_scores = image_bboxes.get('score', [])
    bbox_text = map(str, image_bboxes.get('text', []))

    # Parts
    image_parts = image_objects.get('parts', {})
    parts_x = image_parts.get('x', [])
    parts_y = image_parts.get('y', [])
    parts_v = image_parts.get('v', [])
    parts_s = image_parts.get('score', [])

    # Areas
    object_areas = image_objects.get('area', [])

    # Ids
    object_ids = map(str, image_objects.get('id', []))

    # Any extra data (e.g. stringified json)
    extra_info = str(image_class.get('extra', ''))

    colorspace = 'RGB'
    channels = 3
    image_format = 'JPEG'

    example = tf.train.Example(features=tf.train.Features(feature={
        'image/height': _int64_feature(height),
        'image/width': _int64_feature(width),
        'image/colorspace': _bytes_feature(colorspace),
        'image/channels': _int64_feature(channels),
        'image/format': _bytes_feature(image_format),
        'image/filename': _bytes_feature(filename),
        'image/fileurl': _bytes_feature(fileurl),
        'image/id': _bytes_feature(image_id),
        'image/encoded': _bytes_feature(image_buffer),
        'image/extra': _bytes_feature(extra_info),
        'image/class/label': _int64_feature(class_label),
        'image/class/text': _bytes_feature(class_text),
        'image/object/bbox/xmin': _float_feature(xmin),
        'image/object/bbox/xmax': _float_feature(xmax),
        'image/object/bbox/ymin': _float_feature(ymin),
        'image/object/bbox/ymax': _float_feature(ymax),
        'image/object/bbox/label': _int64_feature(bbox_labels),
        'image/object/bbox/text': _bytes_feature(bbox_text),
        'image/object/bbox/score' : _float_feature(bbox_scores),
        'image/object/parts/x' : _float_feature(parts_x),
        'image/object/parts/y' : _float_feature(parts_y),
        'image/object/parts/v' : _int64_feature(parts_v),
        'image/object/parts/score' : _float_feature(parts_s),
        'image/object/count' : _int64_feature(object_count),
        'image/object/area' : _float_feature(object_areas),
        'image/object/id' : _int64_feature(object_ids)
    }))
    return example

class ImageCoder(object):
    """Helper class that provides TensorFlow image coding utilities."""

    def __init__(self):
        # Create a single Session to run all image coding calls.
        self._sess = tf.Session()

        # Initializes function that converts PNG to JPEG data.
        self._png_data = tf.placeholder(dtype=tf.string)
        image = tf.image.decode_png(self._png_data, channels=3)
        self._png_to_jpeg = tf.image.encode_jpeg(image, format='rgb', quality=100)

        # Initializes function that decodes RGB JPEG data.
        self._decode_jpeg_data = tf.placeholder(dtype=tf.string)
        self._decode_jpeg = tf.image.decode_jpeg(self._decode_jpeg_data, channels=3)

    def png_to_jpeg(self, image_data):
        # Convert the image data from png to jpg
        return self._sess.run(self._png_to_jpeg,
                              feed_dict={self._png_data: image_data})

    def decode_jpeg(self, image_data):
        # Decode the image data as a jpeg image
        image = self._sess.run(self._decode_jpeg,
                               feed_dict={self._decode_jpeg_data: image_data})
        assert len(image.shape) == 3, "JPEG needs to have height x width x channels"
        assert image.shape[2] == 3, "JPEG needs to have 3 channels"
        return image

def _is_png(filename):
    """Determine if a file contains a PNG format image.
    Args:
      filename: string, path of the image file.
    Returns:
      boolean indicating if the image is a PNG.
    """
    _, file_extension = os.path.splitext(filename)
    return file_extension == '.png'

def _process_image(filename, coder):
    """Process a single image file.
    Args:
      filename: string, path to an image file e.g., '/path/to/example.JPG'.
      coder: instance of ImageCoder to provide TensorFlow image coding utils.
    Returns:
      image_buffer: string, JPEG encoding of RGB image.
      height: integer, image height in pixels.
      width: integer, image width in pixels.
    """
    # Read the image file.
    image_data = tf.gfile.FastGFile(filename, 'r').read()

    # Clean the dirty data.
    if _is_png(filename):
        # 1 image is a PNG.
        #print('Converting PNG to JPEG for %s' % filename)
        image_data = coder.png_to_jpeg(image_data)

    # Decode the RGB JPEG.
    image = coder.decode_jpeg(image_data)

    # Check that image converted to RGB
    assert len(image.shape) == 3
    height = image.shape[0]
    width = image.shape[1]
    assert image.shape[2] == 3

    return image_data, height, width

def _process_image_files_batch(coder, thread_index, ranges, name, output_directory,
                               dataset, num_shards, error_queue):
    """Processes and saves list of images as TFRecord in 1 thread.
    Args:
      coder: instance of ImageCoder to provide TensorFlow image coding utils.
      thread_index: integer, unique batch to run index is within [0, len(ranges)).
      ranges: list of pairs of integers specifying ranges of each batches to
        analyze in parallel.
      name: string, unique identifier specifying the data set (e.g. `train` or `test`)
      output_directory: string, file path to store the tfrecord files.
      dataset: list, a list of image example dicts
      num_shards: integer number of shards for this data set.
      error_queue: Queue, a queue to place image examples that failed.
    """
    # Each thread produces N shards where N = int(num_shards / num_threads).
    # For instance, if num_shards = 128, and the num_threads = 2, then the first
    # thread would produce shards [0, 64).
    num_threads = len(ranges)
    assert not num_shards % num_threads
    num_shards_per_batch = int(num_shards / num_threads)

    shard_ranges = np.linspace(ranges[thread_index][0],
                               ranges[thread_index][1],
                               num_shards_per_batch + 1).astype(int)
    num_files_in_thread = ranges[thread_index][1] - ranges[thread_index][0]

    counter = 0
    error_counter = 0
    for s in xrange(num_shards_per_batch):
        # Generate a sharded version of the file name, e.g. 'train-00002-of-00010'
        shard = thread_index * num_shards_per_batch + s
        output_filename = '%s-%.5d-of-%.5d' % (name, shard, num_shards)
        output_file = os.path.join(output_directory, output_filename)
        writer = tf.python_io.TFRecordWriter(output_file)

        shard_counter = 0
        files_in_shard = np.arange(shard_ranges[s], shard_ranges[s + 1], dtype=int)
        for i in files_in_shard:

            image_example = dataset[i]

            filename = str(image_example['filename'])

            try:
                image_buffer, height, width = _process_image(filename, coder)

                example = _convert_to_example(image_example, image_buffer, height, width)
                writer.write(example.SerializeToString())
                shard_counter += 1
                counter += 1
            except Exception as e:
                error_counter += 1
                error_msg = repr(e)
                image_example['error_msg'] = error_msg
                error_queue.put(image_example)

            if not counter % 1000:
                print('%s [thread %d]: Processed %d of %d images in thread batch, with %d errors.' %
                      (datetime.now(), thread_index, counter, num_files_in_thread, error_counter))
                sys.stdout.flush()

        print('%s [thread %d]: Wrote %d images to %s, with %d errors.' %
              (datetime.now(), thread_index, shard_counter, output_file, error_counter))
        sys.stdout.flush()
        shard_counter = 0

    print('%s [thread %d]: Wrote %d images to %d shards, with %d errors.' %
          (datetime.now(), thread_index, counter, num_files_in_thread, error_counter))
    sys.stdout.flush()

def create(dataset, dataset_name, output_directory, num_shards, num_threads, shuffle=True):
    """Create the tfrecord files to be used to train or test a model.

    Args:
      dataset : [{
        "filename" : <REQUIRED: path to the image file>,
        "id" : <REQUIRED: id of the image>,
        "class" : {
          "label" : <[0, num_classes)>,
          "text" : <text description of class>
        },
        "object" : {
          "bbox" : {
            "xmin" : [],
            "xmax" : [],
            "ymin" : [],
            "ymax" : [],
            "label" : []
          }
        }
      }]

      dataset_name: a name for the dataset

      output_directory: path to a directory to write the tfrecord files

      num_shards: the number of tfrecord files to create

      num_threads: the number of threads to use

      shuffle : bool, should the image examples be shuffled or not prior to creating the tfrecords.

    Returns:
      list : a list of image examples that failed to process.
    """

    # Images in the tfrecords set must be shuffled properly
    if shuffle:
        random.shuffle(dataset)

    # Break all images into batches with a [ranges[i][0], ranges[i][1]].
    spacing = np.linspace(0, len(dataset), num_threads + 1).astype(np.int)
    ranges = []
    threads = []
    for i in xrange(len(spacing) - 1):
        ranges.append([spacing[i], spacing[i+1]])

    # Launch a thread for each batch.
    print('Launching %d threads for spacings: %s' % (num_threads, ranges))
    sys.stdout.flush()

    # Create a mechanism for monitoring when all threads are finished.
    coord = tf.train.Coordinator()

    # Create a generic TensorFlow-based utility for converting all image codings.
    coder = ImageCoder()

    # A Queue to hold the image examples that fail to process.
    error_queue = Queue()

    threads = []
    for thread_index in xrange(len(ranges)):
        args = (coder, thread_index, ranges, dataset_name, output_directory, dataset,
                num_shards, error_queue)
        t = threading.Thread(target=_process_image_files_batch, args=args)
        t.start()
        threads.append(t)

    # Wait for all the threads to terminate.
    coord.join(threads)
    print('%s: Finished writing all %d images in data set.' %
          (datetime.now(), len(dataset)))

    # Collect the errors
    errors = []
    while not error_queue.empty():
        errors.append(error_queue.get())
    print ('%d examples failed.' % (len(errors),))

    return errors
<commit_msg>adding check for encoded data<commit_after>"""
Create the tfrecord files for a dataset.

A lot of this code comes from the tensorflow inception example, so here is their license:

# Copyright 2016 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.
# ==============================================================================

"""
from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

from datetime import datetime
import os
from Queue import Queue
import random
import sys
import threading

import numpy as np
import tensorflow as tf

def _int64_feature(value):
    """Wrapper for inserting int64 features into Example proto."""
    if not isinstance(value, list):
        value = [value]
    return tf.train.Feature(int64_list=tf.train.Int64List(value=value))

def _float_feature(value):
    """Wrapper for inserting float features into Example proto."""
    if not isinstance(value, list):
        value = [value]
    return tf.train.Feature(float_list=tf.train.FloatList(value=value))

def _bytes_feature(value):
    """Wrapper for inserting bytes features into Example proto."""
    if not isinstance(value, list):
        value = [value]
    return tf.train.Feature(bytes_list=tf.train.BytesList(value=value))

def _convert_to_example(image_example, image_buffer, height, width, colorspace='RGB',
                        channels=3, image_format='JPEG'):
    """Build an Example proto for an example.
    Args:
      image_example: dict, an image example
      image_buffer: string, JPEG encoding of RGB image
      height: integer, image height in pixels
      width: integer, image width in pixels
    Returns:
      Example proto
    """

    # Required
    filename = str(image_example['filename'])
    image_id = str(image_example['id'])

    # Class label for the whole image
    image_class = image_example.get('class', {})
    class_label = image_class.get('label', 0)
    class_text = str(image_class.get('text', ''))
    class_conf = image_class.get('conf', 1.)

    # Objects
    image_objects = image_example.get('object', {})
    object_count = image_objects.get('count', 0)

    # Bounding Boxes
    image_bboxes = image_objects.get('bbox', {})
    xmin = image_bboxes.get('xmin', [])
    xmax = image_bboxes.get('xmax', [])
    ymin = image_bboxes.get('ymin', [])
    ymax = image_bboxes.get('ymax', [])
    bbox_scores = image_bboxes.get('score', [])
    bbox_labels = image_bboxes.get('label', [])
    bbox_text = map(str, image_bboxes.get('text', []))
    bbox_label_confs = image_bboxes.get('conf', [])

    # Parts
    image_parts = image_objects.get('parts', {})
    parts_x = image_parts.get('x', [])
    parts_y = image_parts.get('y', [])
    parts_v = image_parts.get('v', [])
    parts_s = image_parts.get('score', [])

    # Areas
    object_areas = image_objects.get('area', [])

    # Ids
    object_ids = map(str, image_objects.get('id', []))

    # Any extra data (e.g. stringified json)
    extra_info = str(image_class.get('extra', ''))

    example = tf.train.Example(features=tf.train.Features(feature={
        'image/height': _int64_feature(height),
        'image/width': _int64_feature(width),
        'image/colorspace': _bytes_feature(colorspace),
        'image/channels': _int64_feature(channels),
        'image/format': _bytes_feature(image_format),
        'image/filename': _bytes_feature(filename),
        'image/id': _bytes_feature(image_id),
        'image/encoded': _bytes_feature(image_buffer),
        'image/extra': _bytes_feature(extra_info),
        'image/class/label': _int64_feature(class_label),
        'image/class/text': _bytes_feature(class_text),
        'image/class/conf': _float_feature(class_conf),
        'image/object/bbox/xmin': _float_feature(xmin),
        'image/object/bbox/xmax': _float_feature(xmax),
        'image/object/bbox/ymin': _float_feature(ymin),
        'image/object/bbox/ymax': _float_feature(ymax),
        'image/object/bbox/label': _int64_feature(bbox_labels),
        'image/object/bbox/text': _bytes_feature(bbox_text),
        'image/object/bbox/conf': _float_feature(bbox_label_confs),
        'image/object/bbox/score' : _float_feature(bbox_scores),
        'image/object/parts/x' : _float_feature(parts_x),
        'image/object/parts/y' : _float_feature(parts_y),
        'image/object/parts/v' : _int64_feature(parts_v),
        'image/object/parts/score' : _float_feature(parts_s),
        'image/object/count' : _int64_feature(object_count),
        'image/object/area' : _float_feature(object_areas),
        'image/object/id' : _int64_feature(object_ids)
    }))
    return example

class ImageCoder(object):
    """Helper class that provides TensorFlow image coding utilities."""

    def __init__(self):
        # Create a single Session to run all image coding calls.
        self._sess = tf.Session()

        # Initializes function that converts PNG to JPEG data.
        self._png_data = tf.placeholder(dtype=tf.string)
        image = tf.image.decode_png(self._png_data, channels=3)
        self._png_to_jpeg = tf.image.encode_jpeg(image, format='rgb', quality=100)

        # Initializes function that decodes RGB JPEG data.
        self._decode_jpeg_data = tf.placeholder(dtype=tf.string)
        self._decode_jpeg = tf.image.decode_jpeg(self._decode_jpeg_data, channels=3)

    def png_to_jpeg(self, image_data):
        # Convert the image data from png to jpg
        return self._sess.run(self._png_to_jpeg,
                              feed_dict={self._png_data: image_data})

    def decode_jpeg(self, image_data):
        # Decode the image data as a jpeg image
        image = self._sess.run(self._decode_jpeg,
                               feed_dict={self._decode_jpeg_data: image_data})
        assert len(image.shape) == 3, "JPEG needs to have height x width x channels"
        assert image.shape[2] == 3, "JPEG needs to have 3 channels (RGB)"
        return image

def _is_png(filename):
    """Determine if a file contains a PNG format image.
    Args:
      filename: string, path of the image file.
    Returns:
      boolean indicating if the image is a PNG.
    """
    _, file_extension = os.path.splitext(filename)
    return file_extension.lower() == '.png'

def _process_image(filename, coder):
    """Process a single image file.
    Args:
      filename: string, path to an image file e.g., '/path/to/example.JPG'.
      coder: instance of ImageCoder to provide TensorFlow image coding utils.
    Returns:
      image_buffer: string, JPEG encoding of RGB image.
      height: integer, image height in pixels.
      width: integer, image width in pixels.
    """
    # Read the image file.
    image_data = tf.gfile.FastGFile(filename, 'r').read()

    # Clean the dirty data.
    if _is_png(filename):
        image_data = coder.png_to_jpeg(image_data)

    # Decode the RGB JPEG.
    image = coder.decode_jpeg(image_data)

    # Check that image converted to RGB
    assert len(image.shape) == 3
    height = image.shape[0]
    width = image.shape[1]
    assert image.shape[2] == 3

    return image_data, height, width

def _process_image_files_batch(coder, thread_index, ranges, name, output_directory,
                               dataset, num_shards, error_queue):
    """Processes and saves list of images as TFRecord in 1 thread.
    Args:
      coder: instance of ImageCoder to provide TensorFlow image coding utils.
      thread_index: integer, unique batch to run index is within [0, len(ranges)).
      ranges: list of pairs of integers specifying ranges of each batches to
        analyze in parallel.
      name: string, unique identifier specifying the data set (e.g. `train` or `test`)
      output_directory: string, file path to store the tfrecord files.
      dataset: list, a list of image example dicts
      num_shards: integer number of shards for this data set.
      error_queue: Queue, a queue to place image examples that failed.
    """
    # Each thread produces N shards where N = int(num_shards / num_threads).
    # For instance, if num_shards = 128, and the num_threads = 2, then the first
    # thread would produce shards [0, 64).
    num_threads = len(ranges)
    assert not num_shards % num_threads
    num_shards_per_batch = int(num_shards / num_threads)

    shard_ranges = np.linspace(ranges[thread_index][0],
                               ranges[thread_index][1],
                               num_shards_per_batch + 1).astype(int)
    num_files_in_thread = ranges[thread_index][1] - ranges[thread_index][0]

    counter = 0
    error_counter = 0
    for s in xrange(num_shards_per_batch):
        # Generate a sharded version of the file name, e.g. 'train-00002-of-00010'
        shard = thread_index * num_shards_per_batch + s
        output_filename = '%s-%.5d-of-%.5d' % (name, shard, num_shards)
        output_file = os.path.join(output_directory, output_filename)
        writer = tf.python_io.TFRecordWriter(output_file)

        shard_counter = 0
        files_in_shard = np.arange(shard_ranges[s], shard_ranges[s + 1], dtype=int)
        for i in files_in_shard:

            image_example = dataset[i]

            filename = str(image_example['filename'])

            try:
                if 'encoded' in image_example:
                    image_buffer = image_example['encoded']
                    height = image_example['height']
                    width = image_example['width']
                    colorspace = image_example['colorspace']
                    image_format = image_example['format']
                    num_channels = image_example['channels']
                    example = _convert_to_example(image_example, image_buffer, height,
                                                  width, colorspace, num_channels,
                                                  image_format)

                else:
                    image_buffer, height, width = _process_image(filename, coder)
                    example = _convert_to_example(image_example, image_buffer, height,
                                                  width)

                writer.write(example.SerializeToString())
                shard_counter += 1
                counter += 1
            except Exception as e:
                error_counter += 1
                error_msg = repr(e)
                image_example['error_msg'] = error_msg
                error_queue.put(image_example)

            if not counter % 1000:
                print('%s [thread %d]: Processed %d of %d images in thread batch, with %d errors.' %
                      (datetime.now(), thread_index, counter, num_files_in_thread, error_counter))
                sys.stdout.flush()

        print('%s [thread %d]: Wrote %d images to %s, with %d errors.' %
              (datetime.now(), thread_index, shard_counter, output_file, error_counter))
        sys.stdout.flush()
        shard_counter = 0

    print('%s [thread %d]: Wrote %d images to %d shards, with %d errors.' %
          (datetime.now(), thread_index, counter, num_files_in_thread, error_counter))
    sys.stdout.flush()

def create(dataset, dataset_name, output_directory, num_shards, num_threads, shuffle=True):
    """Create the tfrecord files to be used to train or test a model.

    Args:
      dataset : [{
        "filename" : <REQUIRED: path to the image file>,
        "id" : <REQUIRED: id of the image>,
        "class" : {
          "label" : <[0, num_classes)>,
          "text" : <text description of class>
        },
        "object" : {
          "bbox" : {
            "xmin" : [],
            "xmax" : [],
            "ymin" : [],
            "ymax" : [],
            "label" : []
          }
        }
      }]

      dataset_name: a name for the dataset

      output_directory: path to a directory to write the tfrecord files

      num_shards: the number of tfrecord files to create

      num_threads: the number of threads to use

      shuffle : bool, should the image examples be shuffled or not prior to creating the tfrecords.

    Returns:
      list : a list of image examples that failed to process.
    """

    # Images in the tfrecords set must be shuffled properly
    if shuffle:
        random.shuffle(dataset)

    # Break all images into batches with a [ranges[i][0], ranges[i][1]].
    spacing = np.linspace(0, len(dataset), num_threads + 1).astype(np.int)
    ranges = []
    threads = []
    for i in xrange(len(spacing) - 1):
        ranges.append([spacing[i], spacing[i+1]])

    # Launch a thread for each batch.
    print('Launching %d threads for spacings: %s' % (num_threads, ranges))
    sys.stdout.flush()

    # Create a mechanism for monitoring when all threads are finished.
    coord = tf.train.Coordinator()

    # Create a generic TensorFlow-based utility for converting all image codings.
    coder = ImageCoder()

    # A Queue to hold the image examples that fail to process.
    error_queue = Queue()

    threads = []
    for thread_index in xrange(len(ranges)):
        args = (coder, thread_index, ranges, dataset_name, output_directory, dataset,
                num_shards, error_queue)
        t = threading.Thread(target=_process_image_files_batch, args=args)
        t.start()
        threads.append(t)

    # Wait for all the threads to terminate.
    coord.join(threads)
    print('%s: Finished writing all %d images in data set.' %
          (datetime.now(), len(dataset)))

    # Collect the errors
    errors = []
    while not error_queue.empty():
        errors.append(error_queue.get())
    print ('%d examples failed.' % (len(errors),))

    return errors
<|endoftext|>


<commit_before># Copyright (c) 2006 rPath, Inc.
#
# This program is distributed under the terms of the Common Public License,
# version 1.0. A copy of this license should have been distributed with this
# source file in a file called LICENSE. If it is not present, the license
# is always available at http://www.opensource.org/licenses/cpl.php.
#
# This program is distributed in the hope that it will be useful, but
# without any waranty; without even the implied warranty of merchantability
# or fitness for a particular purpose. See the Common Public License for
# full details.
#
""" Defines an on-demand importer that only actually loads modules when their
    attributes are accessed.  NOTE: if the ondemand module is viewed using
    introspection, like dir(), isinstance, etc, it will appear as a
    ModuleProxy, not a module, and will not have the correct attributes.
    Barring introspection, however, the module will behave as normal.
"""
import sys
import imp
import os
import types

def makeImportedModule(name, data, scope):
    """ Returns a ModuleProxy that has access to a closure w/
        information about the module to load, but is otherwise
        empty.  On an attempted access of any member of the module,
        the module is loaded.
    """

    def _loadModule():
        """ Load the given module, and insert it into the parent
            scope, and also the original importing scope.
        """
        mod = sys.modules.get(name, None)
        if mod is None or not isinstance(mod, types.ModuleType):
            mod = imp.load_module(name, *data)
            sys.modules[name] = mod
        scope[name] = mod

        frame = sys._getframe(2)
        global_scope = frame.f_globals
        local_scope = frame.f_locals

        if name in local_scope:
            local_scope[name] = mod
        elif name in global_scope:
            global_scope[name] = mod

        return mod

    class ModuleProxy(object):
        # we don't add any docs for the module in case the
        # user tries accessing '__doc__'
        def __hasattr__(self, key):
            mod = _loadModule()
            return hasattr(mod, key)

        def __getattr__(self, key):
            mod = _loadModule()
            if key == '__all__':
                # the caller tried to use __all__ to implement import *.
                # Unforunately, that probably means the caller is
                # an import statement which has a handle on an import
                # object.  We can't rely on being able to overwrite that
                # variable, so just update the proxy's dict.
                self.__dict__.update(mod.__dict__)
            return getattr(mod, key)

        def __setattr__(self, key, value):
            mod = _loadModule()
            return setattr(mod, key, value)

        def __repr__(self):
            return "<moduleProxy '%s' from '%s'>" % (name, data[1])

    return ModuleProxy()

class OnDemandLoader(object):
    """ The loader takes a name and info about the module to load and
        "loads" it - in this case returning loading a proxy that
        will only load the class when an attribute is accessed.
    """
    def __init__(self, name, data, scope):
        self.name = name
        self.data = data
        self.scope = scope

    def load_module(self, fullname):
        mod = makeImportedModule(self.name, self.data, self.scope)
        sys.modules[fullname] = mod
        return mod

class OnDemandImporter(object):
    """ The on-demand importer imports a module proxy that
        inserts the desired module into the calling scope only when
        an attribute from the module is actually used.
    """

    def find_module(self, fullname, path=None):
        origName = fullname
        if not path:
            mod = sys.modules.get(fullname, False)
            if mod is None or mod and isinstance(mod, types.ModuleType):
                return mod

        frame = sys._getframe(1)
        global_scope = frame.f_globals
        # this is the scope in which import <fullname> was called

        if '.' in fullname:
            head, fullname = fullname.rsplit('.', 1)

            # this import protocol works such that if I am going to be
            # able to import fullname, then everything in front of the
            # last . in fullname must already be loaded into sys.modules.
            mod = sys.modules.get(head,None)
            if mod is None:
                return None

            if hasattr(mod, '__path__'):
                path = mod.__path__

        try:
            data = imp.find_module(fullname, path)
            return OnDemandLoader(origName, data, global_scope)
        except ImportError:
            # don't return an import error.  That will stop
            # the automated search mechanism from working.
            return None

def install():
    sys.meta_path.append(OnDemandImporter())
<commit_msg>fixes for lazy imports, so that a module doesn't overwrite a builtin<commit_after># Copyright (c) 2006 rPath, Inc.
#
# This program is distributed under the terms of the Common Public License,
# version 1.0. A copy of this license should have been distributed with this
# source file in a file called LICENSE. If it is not present, the license
# is always available at http://www.opensource.org/licenses/cpl.php.
#
# This program is distributed in the hope that it will be useful, but
# without any waranty; without even the implied warranty of merchantability
# or fitness for a particular purpose. See the Common Public License for
# full details.
#
""" Defines an on-demand importer that only actually loads modules when their
    attributes are accessed.  NOTE: if the ondemand module is viewed using
    introspection, like dir(), isinstance, etc, it will appear as a
    ModuleProxy, not a module, and will not have the correct attributes.
    Barring introspection, however, the module will behave as normal.
"""
import sys
import imp
import os
import types

def makeImportedModule(name, data, scope):
    """ Returns a ModuleProxy that has access to a closure w/
        information about the module to load, but is otherwise
        empty.  On an attempted access of any member of the module,
        the module is loaded.
    """

    def _loadModule():
        """ Load the given module, and insert it into the parent
            scope, and also the original importing scope.
        """
        mod = sys.modules.get(name, None)
        if mod is None or not isinstance(mod, types.ModuleType):
            mod = imp.load_module(name, *data)
            sys.modules[name] = mod

        scope[name] = mod

        frame = sys._getframe(2)
        global_scope = frame.f_globals
        local_scope = frame.f_locals

        if name in local_scope:
            if name.__class__.__name__ == 'ModuleProxy':
                local_scope[name] = mod
        elif name in global_scope:
            if name.__class__.__name__ == 'ModuleProxy':
                global_scope[name] = mod

        return mod

    class ModuleProxy(object):
        # we don't add any docs for the module in case the
        # user tries accessing '__doc__'
        def __hasattr__(self, key):
            mod = _loadModule()
            return hasattr(mod, key)

        def __getattr__(self, key):
            mod = _loadModule()
            if key == '__all__':
                # the caller tried to use __all__ to implement import *.
                # Unforunately, that probably means the caller is
                # an import statement which has a handle on an import
                # object.  We can't rely on being able to overwrite that
                # variable, so just update the proxy's dict.
                self.__dict__.update(mod.__dict__)
            return getattr(mod, key)

        def __setattr__(self, key, value):
            mod = _loadModule()
            return setattr(mod, key, value)

        def __repr__(self):
            return "<moduleProxy '%s' from '%s'>" % (name, data[1])

    return ModuleProxy()

class OnDemandLoader(object):
    """ The loader takes a name and info about the module to load and
        "loads" it - in this case returning loading a proxy that
        will only load the class when an attribute is accessed.
    """
    def __init__(self, name, data, scope):
        self.name = name
        self.data = data
        self.scope = scope

    def load_module(self, fullname):
	if fullname in __builtins__:
            mod = imp.load_module(self.name, *self.data)
	    sys.modules[fullname] = mod
        else:
            mod = makeImportedModule(self.name, self.data, self.scope)
            sys.modules[fullname] = mod
        return mod

class OnDemandImporter(object):
    """ The on-demand importer imports a module proxy that
        inserts the desired module into the calling scope only when
        an attribute from the module is actually used.
    """

    def find_module(self, fullname, path=None):
        origName = fullname
        if not path:
            mod = sys.modules.get(fullname, False)
            if mod is None or mod and isinstance(mod, types.ModuleType):
                return mod

        frame = sys._getframe(1)
        global_scope = frame.f_globals
        # this is the scope in which import <fullname> was called

        if '.' in fullname:
            head, fullname = fullname.rsplit('.', 1)

            # this import protocol works such that if I am going to be
            # able to import fullname, then everything in front of the
            # last . in fullname must already be loaded into sys.modules.
            mod = sys.modules.get(head,None)
            if mod is None:
                return None

            if hasattr(mod, '__path__'):
                path = mod.__path__

        try:
            data = imp.find_module(fullname, path)
            return OnDemandLoader(origName, data, global_scope)
        except ImportError:
            # don't return an import error.  That will stop
            # the automated search mechanism from working.
            return None

def install():
    sys.meta_path.append(OnDemandImporter())
<|endoftext|>


<commit_before>from twisted.internet.protocol import Factory
from twisted.protocols.basic import LineReceiver
from twisted.application.internet import UNIXServer

import os
import json

import events

class Scrollback:
    def __init__(self, length):
        self.length = length
        self.data = []

    def put(self, line):
        self.data.append(line)
        if len(self.data) > self.length:
            self.data.pop(0)

    def get(self, max_items = None):
        if max_items is None:
            return self.data[:]
        else:
            return self.data[-max_items:]

class UserServerProtocol(LineReceiver):
    delimiter = '\n'

    tab_last = None
    tab_index = 0

    attached_user = None

    def connectionMade(self):
        self.register(self.console_helper, events.Console)
        self.register(self.handle_attach,  events.UserAttach)
        self.register(self.handle_detach,  events.UserDetach)

    def connectionLost(self, reason):
        if self.attached_user:
            self.dispatch(events.UserDetach(user=self.attached_user))

        self.unregister(self.console_helper, events.Console)
        self.unregister(self.handle_attach,  events.UserAttach)
        self.unregister(self.handle_detach,  events.UserDetach)

    def lineReceived(self, line):
        msg = json.loads(str(line))
        ty = msg["type"]

        if ty == "attach":
            self.attached_user = msg['user']
            self.dispatch(events.UserAttach(user=msg['user']))

        elif ty == "input":
            self.dispatch(events.UserInput(user=msg['user'], line=msg['line']))

        elif ty == "get_lines":
            for e in self.factory.scrollback.get(msg['line_count']):
                self.console_helper(e)

        elif ty == "get_users":
            for u in self.factory.users:
                self.send_helper("user_status", user=u, online=True)

        elif ty == "get_stats":
            self.send_helper("stats", stats=self.factory.stats)

        elif ty == "tab":
            beginning = msg["line"].split(" ")
            end = beginning.pop().lower()

            candidates = filter(lambda p: p.lower().startswith(end), self.factory.players)
            if len(candidates) == 0:
                send = msg["data"]
            else:
                i = msg['index'] % len(candidates)
                beginning.append(candidates[i])
                send = " ".join(beginning)

            self.send_helper("tab", line=send)

        else:
            self.factory.parent.console("unknown packet: %s" % str(msg))

    def send_helper(self, ty, **k):
        k["type"] = ty
        self.sendLine(json.dumps(k))

    def console_helper(self, event):
        self.send_helper("console", **{k: getattr(event, k) for k in event.contains})

    def handle_attach(self, event):
        self.send_helper("user_status", user=event.user, online=True)

    def handle_detach(self, event):
        self.send_helper("user_status", user=event.user, online=False)

class UserServerFactory(Factory):
    players = []

    def __init__(self, parent):
        self.parent     = parent
        self.scrollback = Scrollback(200)
        self.users      = set()

        self.parent.events.register(self.handle_console, events.Console)
        self.parent.events.register(self.handle_attach,  events.UserAttach)
        self.parent.events.register(self.handle_detach,  events.UserDetach)

        self.parent.events.register(self.handle_player_count, events.StatPlayerCount)
        self.parent.events.register(self.handle_players,      events.StatPlayers)
        self.parent.events.register(self.handle_memory,       events.StatMemory)
        self.parent.events.register(self.handle_tick_time,    events.StatTickTime)

        self.stats = {k: '___' for k in ('tick_time', 'memory_current', 'memory_max', 'players_current', 'players_max')}

    def buildProtocol(self, addr):
        p = UserServerProtocol()
        p.register   = self.parent.events.register
        p.unregister = self.parent.events.unregister
        p.dispatch   = self.parent.events.dispatch
        p.factory    = self
        return p

    def handle_console(self, event):
        self.scrollback.put(event)

    def handle_attach(self, event):
        self.users.add(event.user)

    def handle_detach(self, event):
        self.users.discard(event.user)

    #stat handlers
    def handle_player_count(self, event):
        self.stats['players_current'] = event.players_current
        self.stats['players_max']     = event.players_max

    def handle_players(self, event):
        self.players = event.players

    def handle_memory(self, event):
        self.stats['memory_current'] = event.memory_current
        self.stats['memory_max']     = event.memory_max

    def handle_tick_time(self, event):
        self.stats['tick_time'] = event.tick_time

class UserServer(UNIXServer):
    def __init__(self, parent, socket):
        self.parent = parent
        if os.path.exists(socket):
            os.remove(socket)
        factory = UserServerFactory(parent)
        UNIXServer.__init__(self, socket, factory)
<commit_msg>user_server: fix handler unregistering<commit_after>from twisted.internet.protocol import Factory
from twisted.protocols.basic import LineReceiver
from twisted.application.internet import UNIXServer

import os
import json

import events

class Scrollback:
    def __init__(self, length):
        self.length = length
        self.data = []

    def put(self, line):
        self.data.append(line)
        if len(self.data) > self.length:
            self.data.pop(0)

    def get(self, max_items = None):
        if max_items is None:
            return self.data[:]
        else:
            return self.data[-max_items:]

class UserServerProtocol(LineReceiver):
    delimiter = '\n'

    tab_last = None
    tab_index = 0

    attached_user = None

    def connectionMade(self):
        self._handlers = []
        for callback, ty in (
            (self.console_helper, events.Console),
            (self.handle_attach,  events.UserAttach),
            (self.handle_detach,  events.UserDetach)):
            self._handlers.append(self.register(callback, ty))

    def connectionLost(self, reason):
        if self.attached_user:
            self.dispatch(events.UserDetach(user=self.attached_user))

        for i in self._handlers:
            self.unregister(i)
        self._handlers = []

    def lineReceived(self, line):
        msg = json.loads(str(line))
        ty = msg["type"]

        if ty == "attach":
            self.attached_user = msg['user']
            self.dispatch(events.UserAttach(user=msg['user']))

        elif ty == "input":
            self.dispatch(events.UserInput(user=msg['user'], line=msg['line']))

        elif ty == "get_lines":
            for e in self.factory.scrollback.get(msg['line_count']):
                self.console_helper(e)

        elif ty == "get_users":
            for u in self.factory.users:
                self.send_helper("user_status", user=u, online=True)

        elif ty == "get_stats":
            self.send_helper("stats", stats=self.factory.stats)

        elif ty == "tab":
            beginning = msg["line"].split(" ")
            end = beginning.pop().lower()

            candidates = filter(lambda p: p.lower().startswith(end), self.factory.players)
            if len(candidates) == 0:
                send = msg["data"]
            else:
                i = msg['index'] % len(candidates)
                beginning.append(candidates[i])
                send = " ".join(beginning)

            self.send_helper("tab", line=send)

        else:
            self.factory.parent.console("unknown packet: %s" % str(msg))

    def send_helper(self, ty, **k):
        k["type"] = ty
        self.sendLine(json.dumps(k))

    def console_helper(self, event):
        self.send_helper("console", **{k: getattr(event, k) for k in event.contains})

    def handle_attach(self, event):
        self.send_helper("user_status", user=event.user, online=True)

    def handle_detach(self, event):
        self.send_helper("user_status", user=event.user, online=False)

class UserServerFactory(Factory):
    players = []

    def __init__(self, parent):
        self.parent     = parent
        self.scrollback = Scrollback(200)
        self.users      = set()

        self.parent.events.register(self.handle_console, events.Console)
        self.parent.events.register(self.handle_attach,  events.UserAttach)
        self.parent.events.register(self.handle_detach,  events.UserDetach)

        self.parent.events.register(self.handle_player_count, events.StatPlayerCount)
        self.parent.events.register(self.handle_players,      events.StatPlayers)
        self.parent.events.register(self.handle_memory,       events.StatMemory)
        self.parent.events.register(self.handle_tick_time,    events.StatTickTime)

        self.stats = {k: '___' for k in ('tick_time', 'memory_current', 'memory_max', 'players_current', 'players_max')}

    def buildProtocol(self, addr):
        p = UserServerProtocol()
        p.register   = self.parent.events.register
        p.unregister = self.parent.events.unregister
        p.dispatch   = self.parent.events.dispatch
        p.factory    = self
        return p

    def handle_console(self, event):
        self.scrollback.put(event)

    def handle_attach(self, event):
        self.users.add(event.user)

    def handle_detach(self, event):
        self.users.discard(event.user)

    #stat handlers
    def handle_player_count(self, event):
        self.stats['players_current'] = event.players_current
        self.stats['players_max']     = event.players_max

    def handle_players(self, event):
        self.players = event.players

    def handle_memory(self, event):
        self.stats['memory_current'] = event.memory_current
        self.stats['memory_max']     = event.memory_max

    def handle_tick_time(self, event):
        self.stats['tick_time'] = event.tick_time

class UserServer(UNIXServer):
    def __init__(self, parent, socket):
        self.parent = parent
        if os.path.exists(socket):
            os.remove(socket)
        factory = UserServerFactory(parent)
        UNIXServer.__init__(self, socket, factory)
<|endoftext|>


<commit_before>from os import environ

import dj_database_url

import apistar
from apistar import App, http, routing, test
from apistar.backends import DjangoBackend
from apistar.commands import django_makemigrations, django_migrate
from apistar.test import CommandLineRunner

def list_stars(orm: DjangoBackend):
    Star = orm.Star
    return {
        'stars': Star.objects.values('name', 'id')
    }

def create_star(orm: DjangoBackend, name: http.QueryParam):
    star = orm.Star(name=name)
    star.save()
    return {
        'star': {'name': star.name, 'id': star.id}
    }

app = App(
    routes=[
        routing.Route('/api/stars', 'GET', list_stars),
        routing.Route('/api/stars/create', 'GET', create_star),
    ],
    settings={
        'DATABASES': {
            'default': dj_database_url.config(
                default=environ.get('DB_URL', 'sqlite:///test.db')
            )
        },
        'INSTALLED_APPS': ['django_project'],
    },
    commands=[django_makemigrations, django_migrate]
)

client = test.TestClient(app)
runner = CommandLineRunner(app)

def test_list_create(monkeypatch):

    def mock_get_current_app():
        return app

    monkeypatch.setattr(apistar.main, 'get_current_app', mock_get_current_app)

    result = runner.invoke(['django_makemigrations'])
    assert 'makemigrations' in result.output
    result = runner.invoke(['django_migrate'])
    assert 'migrate' in result.output

    response = client.get('http://example.com/api/stars/create?name=mars')
    assert response.status_code == 200
    created_star = response.json()
    assert created_star['star']['name'] == 'mars'

    response = client.get('http://example.com/api/stars')
    assert response.status_code == 200
    assert response.json() == {'stars': [created_star['star']]}
<commit_msg>Fix monkeypatch<commit_after>from os import environ

import dj_database_url

import apistar
from apistar import App, http, routing, test
from apistar.backends import DjangoBackend
from apistar.commands import django_makemigrations, django_migrate
from apistar.test import CommandLineRunner

def list_stars(orm: DjangoBackend):
    Star = orm.Star
    return {
        'stars': Star.objects.values('name', 'id')
    }

def create_star(orm: DjangoBackend, name: http.QueryParam):
    star = orm.Star(name=name)
    star.save()
    return {
        'star': {'name': star.name, 'id': star.id}
    }

app = App(
    routes=[
        routing.Route('/api/stars', 'GET', list_stars),
        routing.Route('/api/stars/create', 'GET', create_star),
    ],
    settings={
        'DATABASES': {
            'default': dj_database_url.config(
                default=environ.get('DB_URL', 'sqlite:///test.db')
            )
        },
        'INSTALLED_APPS': ['django_project']
    },
    commands=[django_makemigrations, django_migrate]
)

client = test.TestClient(app)
runner = CommandLineRunner(app)

def test_list_create(monkeypatch):

    def mock_get_current_app():
        return app

    monkeypatch.setattr(apistar.cli, 'get_current_app', mock_get_current_app)

    result = runner.invoke(['django_makemigrations'])
    assert 'makemigrations' in result.output
    result = runner.invoke(['django_migrate'])
    assert 'migrate' in result.output

    response = client.get('http://example.com/api/stars/create?name=mars')
    assert response.status_code == 200
    created_star = response.json()
    assert created_star['star']['name'] == 'mars'

    response = client.get('http://example.com/api/stars')
    assert response.status_code == 200
    assert response.json() == {'stars': [created_star['star']]}
<|endoftext|>


<commit_before>import os
import re
from six.moves.http_client import HTTPSConnection
import socket
import ssl
try:
  import json
except ImportError:
  import simplejson as json

class CertificateError(ValueError):
  """
  Raised when an error occurs when attempting to verify an SSL certificate.
  """
  pass

def _dnsname_to_pat(dn):
  pats = []
  for frag in dn.split(r'.'):
    if frag == '*':
      # When '*' is a fragment by itself, it matches a non-empty dotless
      # fragment.
      pats.append('[^.]+')
    else:
      # Otherwise, '*' matches any dotless fragment.
      frag = re.escape(frag)
      pats.append(frag.replace(r'\*', '[^.]*'))
  return re.compile(r'\A' + r'\.'.join(pats) + r'\Z', re.IGNORECASE)

def match_hostname(cert, hostname):
  """
  This is a backport of the match_hostname() function from Python 3.2,
  essential when using SSL.
  Verifies that *cert* (in decoded format as returned by
  SSLSocket.getpeercert()) matches the *hostname*.  RFC 2818 rules
  are mostly followed, but IP addresses are not accepted for *hostname*.

  CertificateError is raised on failure. On success, the function
  returns nothing.
  """
  if not cert:
    raise ValueError("empty or no certificate")
  dnsnames = []
  san = cert.get('subjectAltName', ())
  for key, value in san:
    if key == 'DNS':
      if _dnsname_to_pat(value).match(hostname):
        return
      dnsnames.append(value)
  if not san:
    # The subject is only checked when subjectAltName is empty
    for sub in cert.get('subject', ()):
      for key, value in sub:
        # XXX according to RFC 2818, the most specific Common Name
        # must be used.
        if key == 'commonName':
          if _dnsname_to_pat(value).match(hostname):
            return
          dnsnames.append(value)
  if len(dnsnames) > 1:
    raise CertificateError("hostname %r "
      "doesn't match either of %s"
      % (hostname, ', '.join(map(repr, dnsnames))))
  elif len(dnsnames) == 1:
    raise CertificateError("hostname %r "
      "doesn't match %r"
      % (hostname, dnsnames[0]))
  else:
    raise CertificateError("no appropriate commonName or "
      "subjectAltName fields were found")

class VerifiedHTTPSConnection(HTTPSConnection):
  """
  A connection that includes SSL certificate verification.
  """
  def connect(self):
    self.connection_kwargs = {}
    # for > py2.5
    if hasattr(self, 'timeout'):
      self.connection_kwargs.update(timeout = self.timeout)

    # for >= py2.7
    if hasattr(self, 'source_address'):
      self.connection_kwargs.update(source_address = self.source_address)

    sock = socket.create_connection((self.host, self.port), **self.connection_kwargs)

    # for >= py2.7
    if getattr(self, '_tunnel_host', None):
      self.sock = sock
      self._tunnel()

    cert_path = os.path.join(os.path.dirname(__file__), 'cacert.pem')

    self.sock = ssl.wrap_socket(
      sock,
      self.key_file,
      self.cert_file,
      cert_reqs=ssl.CERT_REQUIRED,
      ca_certs=cert_path)

    try:
      match_hostname(self.sock.getpeercert(), self.host)
    except CertificateError:
      self.sock.shutdown(socket.SHUT_RDWR)
      self.sock.close()
      raise

def json_to_py(j):
  o = json.loads(j.decode('utf-8'))
  if isinstance(o, dict):
  	return dict_to_object(o)
  else:
    return dict_to_object({ "response": o }).response

def dict_to_object(d):
	"""Recursively converts a dict to an object"""
	top = type('CreateSendModel', (object,), d)
	seqs = tuple, list, set, frozenset
	for i, j in d.items():
	  if isinstance(j, dict):
	    setattr(top, i, dict_to_object(j))
	  elif isinstance(j, seqs):
	    setattr(top, i, type(j)(dict_to_object(sj) if isinstance(sj, dict) else sj for sj in j))
	  else:
	    setattr(top, i, j)
	return top

def get_faker(expected_url, filename, status=None, body = None):

  class Faker(object):
    """Represents a fake web request, including the expected URL, an open
    function which reads the expected response from a fixture file, and the
    expected response status code."""
    def __init__(self, expected_url, filename, status, body = None):
      self.url = self.createsend_url(expected_url)
      self.filename = filename
      self.status = status
      self.body = body

    def open(self):
      if self.filename:
        return open("%s/../test/fixtures/%s" % (os.path.dirname(__file__), self.filename), mode='rb').read()
      else:
        return ''

    def createsend_url(self, url):
      if url.startswith("http"):
        return url
      else:
        return "https://api.createsend.com/api/v3.1/%s" % url

  return Faker(expected_url, filename, status, body)
<commit_msg>run 2 to 3<commit_after>import os
import re
from six.moves.http_client import HTTPSConnection
import socket
import ssl
try:
  import json
except ImportError:
  import simplejson as json

class CertificateError(ValueError):
  """
  Raised when an error occurs when attempting to verify an SSL certificate.
  """
  pass

def _dnsname_to_pat(dn):
  pats = []
  for frag in dn.split(r'.'):
    if frag == '*':
      # When '*' is a fragment by itself, it matches a non-empty dotless
      # fragment.
      pats.append('[^.]+')
    else:
      # Otherwise, '*' matches any dotless fragment.
      frag = re.escape(frag)
      pats.append(frag.replace(r'\*', '[^.]*'))
  return re.compile(r'\A' + r'\.'.join(pats) + r'\Z', re.IGNORECASE)

def match_hostname(cert, hostname):
  """
  This is a backport of the match_hostname() function from Python 3.2,
  essential when using SSL.
  Verifies that *cert* (in decoded format as returned by
  SSLSocket.getpeercert()) matches the *hostname*.  RFC 2818 rules
  are mostly followed, but IP addresses are not accepted for *hostname*.

  CertificateError is raised on failure. On success, the function
  returns nothing.
  """
  if not cert:
    raise ValueError("empty or no certificate")
  dnsnames = []
  san = cert.get('subjectAltName', ())
  for key, value in san:
    if key == 'DNS':
      if _dnsname_to_pat(value).match(hostname):
        return
      dnsnames.append(value)
  if not san:
    # The subject is only checked when subjectAltName is empty
    for sub in cert.get('subject', ()):
      for key, value in sub:
        # XXX according to RFC 2818, the most specific Common Name
        # must be used.
        if key == 'commonName':
          if _dnsname_to_pat(value).match(hostname):
            return
          dnsnames.append(value)
  if len(dnsnames) > 1:
    raise CertificateError("hostname %r "
      "doesn't match either of %s"
      % (hostname, ', '.join(map(repr, dnsnames))))
  elif len(dnsnames) == 1:
    raise CertificateError("hostname %r "
      "doesn't match %r"
      % (hostname, dnsnames[0]))
  else:
    raise CertificateError("no appropriate commonName or "
      "subjectAltName fields were found")

class VerifiedHTTPSConnection(HTTPSConnection):
  """
  A connection that includes SSL certificate verification.
  """
  def connect(self):
    self.connection_kwargs = {}
    # for > py2.5
    if hasattr(self, 'timeout'):
      self.connection_kwargs.update(timeout = self.timeout)

    # for >= py2.7
    if hasattr(self, 'source_address'):
      self.connection_kwargs.update(source_address = self.source_address)

    sock = socket.create_connection((self.host, self.port), **self.connection_kwargs)

    # for >= py2.7
    if getattr(self, '_tunnel_host', None):
      self.sock = sock
      self._tunnel()

    cert_path = os.path.join(os.path.dirname(__file__), 'cacert.pem')

    self.sock = ssl.wrap_socket(
      sock,
      self.key_file,
      self.cert_file,
      cert_reqs=ssl.CERT_REQUIRED,
      ca_certs=cert_path)

    try:
      match_hostname(self.sock.getpeercert(), self.host)
    except CertificateError:
      self.sock.shutdown(socket.SHUT_RDWR)
      self.sock.close()
      raise

def json_to_py(j):
  o = json.loads(j.decode('utf-8'))
  if isinstance(o, dict):
  	return dict_to_object(o)
  else:
    return dict_to_object({ "response": o }).response

def dict_to_object(d):
	"""Recursively converts a dict to an object"""
	top = type('CreateSendModel', (object,), d)
	seqs = tuple, list, set, frozenset
	for i, j in list(d.items()):
	  if isinstance(j, dict):
	    setattr(top, i, dict_to_object(j))
	  elif isinstance(j, seqs):
	    setattr(top, i, type(j)(dict_to_object(sj) if isinstance(sj, dict) else sj for sj in j))
	  else:
	    setattr(top, i, j)
	return top

def get_faker(expected_url, filename, status=None, body = None):

  class Faker(object):
    """Represents a fake web request, including the expected URL, an open
    function which reads the expected response from a fixture file, and the
    expected response status code."""
    def __init__(self, expected_url, filename, status, body = None):
      self.url = self.createsend_url(expected_url)
      self.filename = filename
      self.status = status
      self.body = body

    def open(self):
      if self.filename:
        return open("%s/../test/fixtures/%s" % (os.path.dirname(__file__), self.filename), mode='rb').read()
      else:
        return ''

    def createsend_url(self, url):
      if url.startswith("http"):
        return url
      else:
        return "https://api.createsend.com/api/v3.1/%s" % url

  return Faker(expected_url, filename, status, body)
<|endoftext|>


<commit_before># Django settings for the ocl project.

DEBUG = True
TEMPLATE_DEBUG = DEBUG

ADMINS = (
    # ('Your Name', 'your_email@example.com'),
)

MANAGERS = ADMINS

DATABASES = {
    'default': {
        'ENGINE': 'django_mongodb_engine',
        #'HOST': '65.99.230.144',
        'NAME': 'ocl',
    }
}

# Hosts/domain names that are valid for this site; required if DEBUG is False
# See https://docs.djangoproject.com/en/1.3/ref/settings/#allowed-hosts
ALLOWED_HOSTS = []

# Local time zone for this installation. Choices can be found here:
# http://en.wikipedia.org/wiki/List_of_tz_zones_by_name
# although not all choices may be available on all operating systems.
# On Unix systems, a value of None will cause Django to use the same
# timezone as the operating system.
# If running in a Windows environment this must be set to the same as your
# system time zone.
TIME_ZONE = 'America/New_York'

# Language code for this installation. All choices can be found here:
# http://www.i18nguy.com/unicode/language-identifiers.html
LANGUAGE_CODE = 'en-us'

# If you set this to False, Django will make some optimizations so as not
# to load the internationalization machinery.
USE_I18N = True

# If you set this to False, Django will not format dates, numbers and
# calendars according to the current locale
USE_L10N = True
DEFAULT_LOCALE = 'en'

# Absolute filesystem path to the directory that will hold user-uploaded files.
# Example: "/home/media/media.lawrence.com/media/"
MEDIA_ROOT = ''

# URL that handles the media served from MEDIA_ROOT. Make sure to use a
# trailing slash.
# Examples: "http://media.lawrence.com/media/", "http://example.com/media/"
MEDIA_URL = ''

# Absolute path to the directory static files should be collected to.
# Don't put anything in this directory yourself; store your static files
# in apps' "static/" subdirectories and in STATICFILES_DIRS.
# Example: "/home/media/media.lawrence.com/static/"
STATIC_ROOT = ''
# In the deployment environment, comment out the above line, and uncomment the one below
#STATIC_ROOT = '/usr/local/wsgi/static/'

# URL prefix for static files.
# Example: "http://media.lawrence.com/static/"
STATIC_URL = '/static/'

# URL prefix for admin static files -- CSS, JavaScript and images.
# Make sure to use a trailing slash.
# Examples: "http://foo.com/static/admin/", "/static/admin/".
ADMIN_MEDIA_PREFIX = '/static/admin/'

# Additional locations of static files
STATICFILES_DIRS = (
    # Put strings here, like "/home/html/static" or "C:/www/django/static".
    # Always use forward slashes, even on Windows.
    # Don't forget to use absolute paths, not relative paths.
)

# List of finder classes that know how to find static files in
# various locations.
STATICFILES_FINDERS = (
    'django.contrib.staticfiles.finders.FileSystemFinder',
    'django.contrib.staticfiles.finders.AppDirectoriesFinder',
#    'django.contrib.staticfiles.finders.DefaultStorageFinder',
)

# Make this unique, and don't share it with anybody.
SECRET_KEY = '+p+lx2*o3ywq+z)%f7929b6)93)^mcc9-0eu9ynq77gc+pe=ck'

# List of callables that know how to import templates from various sources.
TEMPLATE_LOADERS = (
    'django.template.loaders.filesystem.Loader',
    'django.template.loaders.app_directories.Loader',
#     'django.template.loaders.eggs.Loader',
)

MIDDLEWARE_CLASSES = (
    'django.middleware.common.CommonMiddleware',
    'django.contrib.sessions.middleware.SessionMiddleware',
    'django.middleware.csrf.CsrfViewMiddleware',
    'django.contrib.auth.middleware.AuthenticationMiddleware',
    'django.contrib.messages.middleware.MessageMiddleware',
)

ROOT_URLCONF = 'urls'

TEMPLATE_DIRS = (
    # Put strings here, like "/home/html/django_templates" or "C:/www/django/templates".
    # Always use forward slashes, even on Windows.
    # Don't forget to use absolute paths, not relative paths.
)

INSTALLED_APPS = (
    'django.contrib.auth',
    'django.contrib.contenttypes',
    'django.contrib.sessions',
    'django.contrib.messages',
    'django.contrib.staticfiles',
    # Uncomment the next line to enable the admin:
    'django.contrib.admin',
    # Uncomment the next line to enable admin documentation:
    # 'django.contrib.admindocs',
    # Third-party apps:
    'djangotoolbox',
    'django_mongodb_engine',
    'rest_framework',
    'rest_framework.authtoken',
    #'south',
    # Project-specific apps:
    #'accounts',
    'concepts',
    'conceptcollections',
    'orgs',
    'oclapi',
    'sources',
    'users',
)

# A sample logging configuration. The only tangible logging
# performed by this configuration is to send an email to
# the site admins on every HTTP 500 error.
# See http://docs.djangoproject.com/en/dev/topics/logging for
# more details on how to customize your logging configuration.
LOGGING = {
    'version': 1,
    'disable_existing_loggers': False,
    'handlers': {
        'mail_admins': {
            'level': 'ERROR',
            'class': 'django.utils.log.AdminEmailHandler'
        }
    },
    'loggers': {
        'django.request': {
            'handlers': ['mail_admins'],
            'level': 'ERROR',
            'propagate': True,
        },
    }
}

# Django Rest Framework configuration
REST_FRAMEWORK = {
    # Default to token-based authentication; fall back on session-based
    # A user gets a unique token upon account creation (residing in the authtoken_token data store).
    # To pass an authentication token along with your request, include the following header:
    # Authorization: Token [TOKEN_VALUE]
    # e.g.
    # Authorization: Token ad73f481096c3b6202bce395820199
    'DEFAULT_AUTHENTICATION_CLASSES': (
        'rest_framework.authentication.TokenAuthentication',
        'rest_framework.authentication.SessionAuthentication',
    ),
    # Use hyperlinked styles by default.
    # Only used if the `serializer_class` attribute is not set on a view.
    'DEFAULT_MODEL_SERIALIZER_CLASS':
        'rest_framework.serializers.HyperlinkedModelSerializer',

    # Use Django's standard `django.contrib.auth` permissions,
    # or allow read-only access for unauthenticated users.
    'DEFAULT_PERMISSION_CLASSES': [
        #'rest_framework.permissions.DjangoModelPermissionsOrAnonReadOnly',
        'rest_framework.permissions.IsAuthenticated',
    ]
}

# Model that stores auxiliary user profile attributes.
# A user must have a profile in order to access the system.
# (A profile is created automatically for any user created using the 'POST /users' endpoint.)
AUTH_PROFILE_MODULE = 'users.UserProfile'
<commit_msg>Added parameters to enable pagination<commit_after># Django settings for the ocl project.

DEBUG = True
TEMPLATE_DEBUG = DEBUG

ADMINS = (
    # ('Your Name', 'your_email@example.com'),
)

MANAGERS = ADMINS

DATABASES = {
    'default': {
        'ENGINE': 'django_mongodb_engine',
        #'HOST': '65.99.230.144',
        'NAME': 'ocl',
    }
}

# Hosts/domain names that are valid for this site; required if DEBUG is False
# See https://docs.djangoproject.com/en/1.3/ref/settings/#allowed-hosts
ALLOWED_HOSTS = []

# Local time zone for this installation. Choices can be found here:
# http://en.wikipedia.org/wiki/List_of_tz_zones_by_name
# although not all choices may be available on all operating systems.
# On Unix systems, a value of None will cause Django to use the same
# timezone as the operating system.
# If running in a Windows environment this must be set to the same as your
# system time zone.
TIME_ZONE = 'America/New_York'

# Language code for this installation. All choices can be found here:
# http://www.i18nguy.com/unicode/language-identifiers.html
LANGUAGE_CODE = 'en-us'

# If you set this to False, Django will make some optimizations so as not
# to load the internationalization machinery.
USE_I18N = True

# If you set this to False, Django will not format dates, numbers and
# calendars according to the current locale
USE_L10N = True
DEFAULT_LOCALE = 'en'

# Absolute filesystem path to the directory that will hold user-uploaded files.
# Example: "/home/media/media.lawrence.com/media/"
MEDIA_ROOT = ''

# URL that handles the media served from MEDIA_ROOT. Make sure to use a
# trailing slash.
# Examples: "http://media.lawrence.com/media/", "http://example.com/media/"
MEDIA_URL = ''

# Absolute path to the directory static files should be collected to.
# Don't put anything in this directory yourself; store your static files
# in apps' "static/" subdirectories and in STATICFILES_DIRS.
# Example: "/home/media/media.lawrence.com/static/"
STATIC_ROOT = ''
# In the deployment environment, comment out the above line, and uncomment the one below
#STATIC_ROOT = '/usr/local/wsgi/static/'

# URL prefix for static files.
# Example: "http://media.lawrence.com/static/"
STATIC_URL = '/static/'

# URL prefix for admin static files -- CSS, JavaScript and images.
# Make sure to use a trailing slash.
# Examples: "http://foo.com/static/admin/", "/static/admin/".
ADMIN_MEDIA_PREFIX = '/static/admin/'

# Additional locations of static files
STATICFILES_DIRS = (
    # Put strings here, like "/home/html/static" or "C:/www/django/static".
    # Always use forward slashes, even on Windows.
    # Don't forget to use absolute paths, not relative paths.
)

# List of finder classes that know how to find static files in
# various locations.
STATICFILES_FINDERS = (
    'django.contrib.staticfiles.finders.FileSystemFinder',
    'django.contrib.staticfiles.finders.AppDirectoriesFinder',
#    'django.contrib.staticfiles.finders.DefaultStorageFinder',
)

# Make this unique, and don't share it with anybody.
SECRET_KEY = '+p+lx2*o3ywq+z)%f7929b6)93)^mcc9-0eu9ynq77gc+pe=ck'

# List of callables that know how to import templates from various sources.
TEMPLATE_LOADERS = (
    'django.template.loaders.filesystem.Loader',
    'django.template.loaders.app_directories.Loader',
#     'django.template.loaders.eggs.Loader',
)

MIDDLEWARE_CLASSES = (
    'django.middleware.common.CommonMiddleware',
    'django.contrib.sessions.middleware.SessionMiddleware',
    'django.middleware.csrf.CsrfViewMiddleware',
    'django.contrib.auth.middleware.AuthenticationMiddleware',
    'django.contrib.messages.middleware.MessageMiddleware',
)

ROOT_URLCONF = 'urls'

TEMPLATE_DIRS = (
    # Put strings here, like "/home/html/django_templates" or "C:/www/django/templates".
    # Always use forward slashes, even on Windows.
    # Don't forget to use absolute paths, not relative paths.
)

INSTALLED_APPS = (
    'django.contrib.auth',
    'django.contrib.contenttypes',
    'django.contrib.sessions',
    'django.contrib.messages',
    'django.contrib.staticfiles',
    # Uncomment the next line to enable the admin:
    'django.contrib.admin',
    # Uncomment the next line to enable admin documentation:
    # 'django.contrib.admindocs',
    # Third-party apps:
    'djangotoolbox',
    'django_mongodb_engine',
    'rest_framework',
    'rest_framework.authtoken',
    #'south',
    # Project-specific apps:
    #'accounts',
    'concepts',
    'conceptcollections',
    'orgs',
    'oclapi',
    'sources',
    'users',
)

# A sample logging configuration. The only tangible logging
# performed by this configuration is to send an email to
# the site admins on every HTTP 500 error.
# See http://docs.djangoproject.com/en/dev/topics/logging for
# more details on how to customize your logging configuration.
LOGGING = {
    'version': 1,
    'disable_existing_loggers': False,
    'handlers': {
        'mail_admins': {
            'level': 'ERROR',
            'class': 'django.utils.log.AdminEmailHandler'
        }
    },
    'loggers': {
        'django.request': {
            'handlers': ['mail_admins'],
            'level': 'ERROR',
            'propagate': True,
        },
    }
}

# Django Rest Framework configuration
REST_FRAMEWORK = {
    # Default to token-based authentication; fall back on session-based
    # A user gets a unique token upon account creation (residing in the authtoken_token data store).
    # To pass an authentication token along with your request, include the following header:
    # Authorization: Token [TOKEN_VALUE]
    # e.g.
    # Authorization: Token ad73f481096c3b6202bce395820199
    'DEFAULT_AUTHENTICATION_CLASSES': (
        'rest_framework.authentication.TokenAuthentication',
        'rest_framework.authentication.SessionAuthentication',
    ),
    # Use hyperlinked styles by default.
    # Only used if the `serializer_class` attribute is not set on a view.
    'DEFAULT_MODEL_SERIALIZER_CLASS':
        'rest_framework.serializers.HyperlinkedModelSerializer',

    # Use Django's standard `django.contrib.auth` permissions,
    # or allow read-only access for unauthenticated users.
    'DEFAULT_PERMISSION_CLASSES': [
        #'rest_framework.permissions.DjangoModelPermissionsOrAnonReadOnly',
        'rest_framework.permissions.IsAuthenticated',
    ],
    'PAGINATE_BY': 10,             # Default to 10
    'PAGINATE_BY_PARAM': 'limit',  # Allow client to override, using `?limit=xxx`.
    'MAX_PAGINATE_BY': 100         # Maximum limit allowed when using `?limit=xxx`.
}

# Model that stores auxiliary user profile attributes.
# A user must have a profile in order to access the system.
# (A profile is created automatically for any user created using the 'POST /users' endpoint.)
AUTH_PROFILE_MODULE = 'users.UserProfile'
<|endoftext|>


<commit_before># coding=utf-8
from __future__ import unicode_literals

import datetime
import re
import os
import logging

# Get an instance of a logger
logger = logging.getLogger(__name__)

from bs4 import BeautifulSoup
from contracts.crawler import AbstractCrawler

from .models import Type, Creator, Document
from . import models

class FirstSeriesCrawler(AbstractCrawler):

    data_directory = '../../law_data'

    # document_id is represented by 8 numbers: year#### (e.g. 19971111)

    class DocumentNotFound(Exception):
        pass

    def __init__(self):
        super(FirstSeriesCrawler, self).__init__()

    def retrieve_document(self, document_id):
        """
        Retrieves the source of the document.
        """
        file_name = '%s/%s.dat' % (self.data_directory, document_id)
        try:
            f = open(file_name, "rb")
            try:
                html = f.read()
            finally:
                f.close()
        except IOError:
            # online retrieval
            html = self.goToPage(models.dre_url_formater.format(document_id=document_id))

            soup = BeautifulSoup(html)

            html = soup.find("div", dict(id='centro_total'))

            if not html.find('div', {'id': 'doc_data'}):
                raise self.DocumentNotFound

            logger.info("saving document_id %d", document_id)
            f = open(file_name, "wb")
            try:
                f.write(html)
            finally:
                f.close()
        return html

    def last_document_id(self, year):
        regex = re.compile(r"%04d(\d+).dat" % year)
        files = [int(re.findall(regex, f)[0]) for f in os.listdir('%s/' % self.data_directory) if re.match(regex, f)]
        files = sorted(files, key=lambda x: int(x), reverse=True)
        if len(files):
            return files[0]
        else:
            return 1

    def get_documents(self, year):
        logger.info("get_documents(year=%d)", year)
        document_number = self.last_document_id(year)
        fails = 0
        while True:
            try:
                self.retrieve_document(int("%04d%04d" % (year, document_number)))
                fails = 0
            except self.DocumentNotFound:
                fails += 1
                logger.info("DocumentNotFound: %d", fails)

            if fails == 50:
                break

            document_number += 1

    def retrieve_all(self):
        first_year = 1910
        last_year = datetime.datetime.now().date().year

        for year in range(first_year, last_year + 1):
            self.get_documents(year)

def clean_data(data_string):
    soup = BeautifulSoup(data_string)

    doc_data = soup.find('div', {'id': 'doc_data'})

    data = {}
    for element in doc_data.findAll('p'):
        if element.find('span'):
            identifier = element.find('span').extract()
        else:
            continue
        string = ' '.join(element.text.split()).strip()  # remove double white spaces

        if identifier.text.startswith('DATA'):
            data['date'] = clean_date(string)
        elif identifier.text.startswith('DIPLOMA'):
            data['type'], data['number'] = clean_type_and_number(string)
        elif identifier.text.startswith('NÚMERO'):
            data['series'], data['series_number'], data['series_other'] = clean_series(string)
        elif identifier.text.startswith('EMISSOR'):
            data['creator'] = clean_creator(string)
        elif identifier.text.startswith('PÁGINAS'):
            data['series_pages'] = clean_pages(string)
        elif identifier.text.startswith('SUMÁRIO'):
            data['summary'] = clean_summary(element)
        elif element.text == '':
            pass
        else:
            logger.exception('Unknown element in data')
            raise IndexError(element.text)

    doc_text = soup.find('div', {'id': 'doc_texto'})
    data['text'] = clean_text(doc_text)

    doc_boxes = soup.find('div', {'id': 'doc_caixas'})
    if doc_boxes is not None:
        doc_pdf = doc_boxes.find('div', {'class': 'cx_pdf'})
        pdf_anchor = doc_pdf.find('a')
        data['pdf_url'] = pdf_anchor['href']

    return data

def clean_date(string):
    if ',' in string:
        string = string.split(', ')[-1]
    elif string[0] in [str(x) for x in range(1, 9)] or string[0:2] in [str(x) for x in range(10, 30)]:
        pass
    else:
        string = ' '.join(string.split(' ')[1:])

    if 'Nota:' in string:
        string = string.split(' Nota:')[0]

    date = datetime.datetime.strptime(string, '%d de %B de %Y').date()
    return date

def clean_type_and_number(string):
    if 'n.º' in string:
        matches = re.search('(.*) n\.º (.*) \(?', string)
        type_name = matches.group(1).strip()
        number = matches.group(2).strip()
    else:
        type_name = string.strip()
        number = None

    # synonymous and typos check
    if type_name == 'Declaração de Rectificação':
        type_name = 'Declaração de Retificação'
    if type_name == 'Declaração de rectificação':
        type_name = 'Declaração de Retificação'
    if type_name == 'Decreto do Presidente de República':
        type_name = 'Decreto do Presidente da República'
    if type_name == 'Resolução da  Assembleia da República':
        type_name = 'Resolução da Assembleia da República'

    type, created = Type.objects.get_or_create(name=type_name)

    return type, number

def clean_series(string):
    strings = string.split(' SÉRIE ')

    series_number = strings[0]
    series_string = strings[1]

    if series_string[0] == 'I':
        series = 1
        series_other = series_string[1:]
    elif series_string[0:1] == 'II':
        series = 2
        series_other = series_string[2:]
    else:
        raise IndexError(series_string)

    return series, series_number, series_other

def clean_creator(string):
    creator_name = string.strip()

    creator, created = Creator.objects.get_or_create(name=creator_name[:254])
    return creator

def clean_pages(string):
    return string

def clean_summary(element):
    search = re.search("
(.*)
", str(element).strip())
    return search.group(1).strip()

def clean_text(doc_text):
    if 'TEXTO' in str(doc_text.find('p')):
        doc_text.find('p').extract()
    else:
        return None

    text = str(doc_text).strip()
    text = text.replace('
', '')
    text = text.replace('
', '')
    return text

class Populator:

    data_directory = '../../law_data'

    def __init__(self):
        import locale
        locale.setlocale(locale.LC_TIME, "pt_PT.UTF-8")

    class DocumentNotFound(Exception):
        pass

    def get_document(self, document_id):
        file_name = '%s/%d.dat' % (self.data_directory, document_id)

        try:
            f = open(file_name, "r")
            try:
                return f.read()
            finally:
                f.close()
        except IOError:
            raise self.DocumentNotFound

    @staticmethod
    def populate_from_document(document_id, data):
        logger.info("populate_from_document(%d)", document_id)
        data = clean_data(data)
        data['dre_doc_id'] = document_id

        try:
            document = Document.objects.get(dre_doc_id=document_id)
        except Document.DoesNotExist:
            document = Document()

        for attr, value in data.items():
            setattr(document, attr, value)
        document.save()

        return document

    def get_cached_documents_id_list(self, year):
        regex = re.compile(r"(%04d\w+).dat" % year)
        files = [int(re.findall(regex, f)[0]) for f in os.listdir('%s/' % self.data_directory) if re.match(regex, f)]
        files = sorted(files, key=lambda x: int(x))
        return files

    def populate_documents(self, year):
        for document_id in self.get_cached_documents_id_list(year):
            self.populate_document(document_id)

    def populate_all(self, first_year=1910):
        last_year = datetime.datetime.now().date().year
        for year in range(first_year, last_year + 1):
            self.populate_documents(year)

    def populate_document(self, document_id):
        data = self.get_document(document_id)
        self.populate_from_document(document_id, data)
<commit_msg>Simplified exception DocumentNotFound in law.crawler.<commit_after># coding=utf-8
from __future__ import unicode_literals

import datetime
import re
import os
import logging

# Get an instance of a logger
logger = logging.getLogger(__name__)

from bs4 import BeautifulSoup
from contracts.crawler import AbstractCrawler

from .models import Type, Creator, Document
from . import models

class DocumentNotFound(Exception):
    pass

class FirstSeriesCrawler(AbstractCrawler):

    data_directory = '../../law_data'

    # document_id is represented by 8 numbers: year#### (e.g. 19971111)

    def __init__(self):
        super(FirstSeriesCrawler, self).__init__()

    def retrieve_document(self, document_id):
        """
        Retrieves the source of the document.
        """
        file_name = '%s/%s.dat' % (self.data_directory, document_id)
        try:
            f = open(file_name, "rb")
            try:
                html = f.read()
            finally:
                f.close()
        except IOError:
            # online retrieval
            html = self.goToPage(models.dre_url_formater.format(document_id=document_id))

            soup = BeautifulSoup(html)

            html = soup.find("div", dict(id='centro_total'))

            if not html.find('div', {'id': 'doc_data'}):
                raise DocumentNotFound

            logger.info("saving document_id %d", document_id)
            f = open(file_name, "wb")
            try:
                f.write(html)
            finally:
                f.close()
        return html

    def last_document_id(self, year):
        regex = re.compile(r"%04d(\d+).dat" % year)
        files = [int(re.findall(regex, f)[0]) for f in os.listdir('%s/' % self.data_directory) if re.match(regex, f)]
        files = sorted(files, key=lambda x: int(x), reverse=True)
        if len(files):
            return files[0]
        else:
            return 1

    def get_documents(self, year):
        logger.info("get_documents(year=%d)", year)
        document_number = self.last_document_id(year)
        fails = 0
        while True:
            try:
                self.retrieve_document(int("%04d%04d" % (year, document_number)))
                fails = 0
            except DocumentNotFound:
                fails += 1
                logger.info("DocumentNotFound: %d", fails)

            if fails == 50:
                break

            document_number += 1

    def retrieve_all(self):
        first_year = 1910
        last_year = datetime.datetime.now().date().year

        for year in range(first_year, last_year + 1):
            self.get_documents(year)

def clean_data(data_string):
    soup = BeautifulSoup(data_string)

    doc_data = soup.find('div', {'id': 'doc_data'})

    data = {}
    for element in doc_data.findAll('p'):
        if element.find('span'):
            identifier = element.find('span').extract()
        else:
            continue
        string = ' '.join(element.text.split()).strip()  # remove double white spaces

        if identifier.text.startswith('DATA'):
            data['date'] = clean_date(string)
        elif identifier.text.startswith('DIPLOMA'):
            data['type'], data['number'] = clean_type_and_number(string)
        elif identifier.text.startswith('NÚMERO'):
            data['series'], data['series_number'], data['series_other'] = clean_series(string)
        elif identifier.text.startswith('EMISSOR'):
            data['creator'] = clean_creator(string)
        elif identifier.text.startswith('PÁGINAS'):
            data['series_pages'] = clean_pages(string)
        elif identifier.text.startswith('SUMÁRIO'):
            data['summary'] = clean_summary(element)
        elif element.text == '':
            pass
        else:
            logger.exception('Unknown element in data')
            raise IndexError(element.text)

    doc_text = soup.find('div', {'id': 'doc_texto'})
    data['text'] = clean_text(doc_text)

    doc_boxes = soup.find('div', {'id': 'doc_caixas'})
    if doc_boxes is not None:
        doc_pdf = doc_boxes.find('div', {'class': 'cx_pdf'})
        pdf_anchor = doc_pdf.find('a')
        data['pdf_url'] = pdf_anchor['href']

    return data

def clean_date(string):
    if ',' in string:
        string = string.split(', ')[-1]
    elif string[0] in [str(x) for x in range(1, 9)] or string[0:2] in [str(x) for x in range(10, 30)]:
        pass
    else:
        string = ' '.join(string.split(' ')[1:])

    if 'Nota:' in string:
        string = string.split(' Nota:')[0]

    date = datetime.datetime.strptime(string, '%d de %B de %Y').date()
    return date

def clean_type_and_number(string):
    if 'n.º' in string:
        matches = re.search('(.*) n\.º (.*) \(?', string)
        type_name = matches.group(1).strip()
        number = matches.group(2).strip()
    else:
        type_name = string.strip()
        number = None

    # synonymous and typos check
    if type_name == 'Declaração de Rectificação':
        type_name = 'Declaração de Retificação'
    if type_name == 'Declaração de rectificação':
        type_name = 'Declaração de Retificação'
    if type_name == 'Decreto do Presidente de República':
        type_name = 'Decreto do Presidente da República'
    if type_name == 'Resolução da  Assembleia da República':
        type_name = 'Resolução da Assembleia da República'

    type, created = Type.objects.get_or_create(name=type_name)

    return type, number

def clean_series(string):
    strings = string.split(' SÉRIE ')

    series_number = strings[0]
    series_string = strings[1]

    if series_string[0] == 'I':
        series = 1
        series_other = series_string[1:]
    elif series_string[0:1] == 'II':
        series = 2
        series_other = series_string[2:]
    else:
        raise IndexError(series_string)

    return series, series_number, series_other

def clean_creator(string):
    creator_name = string.strip()

    creator, created = Creator.objects.get_or_create(name=creator_name[:254])
    return creator

def clean_pages(string):
    return string

def clean_summary(element):
    search = re.search("
(.*)
", str(element).strip())
    return search.group(1).strip()

def clean_text(doc_text):
    if 'TEXTO' in str(doc_text.find('p')):
        doc_text.find('p').extract()
    else:
        return None

    text = str(doc_text).strip()
    text = text.replace('
', '')
    text = text.replace('
', '')
    return text

class Populator:

    data_directory = '../../law_data'

    def __init__(self):
        import locale
        locale.setlocale(locale.LC_TIME, "pt_PT.UTF-8")

    def get_document(self, document_id):
        file_name = '%s/%d.dat' % (self.data_directory, document_id)

        try:
            f = open(file_name, "r")
            try:
                return f.read()
            finally:
                f.close()
        except IOError:
            raise DocumentNotFound

    @staticmethod
    def populate_from_document(document_id, data):
        logger.info("populate_from_document(%d)", document_id)
        data = clean_data(data)
        data['dre_doc_id'] = document_id

        try:
            document = Document.objects.get(dre_doc_id=document_id)
        except Document.DoesNotExist:
            document = Document()

        for attr, value in data.items():
            setattr(document, attr, value)
        document.save()

        return document

    def get_cached_documents_id_list(self, year):
        regex = re.compile(r"(%04d\w+).dat" % year)
        files = [int(re.findall(regex, f)[0]) for f in os.listdir('%s/' % self.data_directory) if re.match(regex, f)]
        files = sorted(files, key=lambda x: int(x))
        return files

    def populate_documents(self, year):
        for document_id in self.get_cached_documents_id_list(year):
            self.populate_document(document_id)

    def populate_all(self, first_year=1910):
        last_year = datetime.datetime.now().date().year
        for year in range(first_year, last_year + 1):
            self.populate_documents(year)

    def populate_document(self, document_id):
        data = self.get_document(document_id)
        self.populate_from_document(document_id, data)
<|endoftext|>