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

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

"""This application demonstrates face detection, label detection,
explicit content, and shot change detection using the Google Cloud API.

Usage Examples:

    python analyze.py faces gs://demomaker/google_gmail.mp4
    python analyze.py labels gs://cloud-ml-sandbox/video/chicago.mp4
    python analyze.py labels_file resources/cat.mp4
    python analyze.py shots gs://demomaker/gbikes_dinosaur.mp4
    python analyze.py explicit_content gs://demomaker/gbikes_dinosaur.mp4

"""

import argparse
import base64
import io

from google.cloud import videointelligence

def analyze_explicit_content(path):
    """ Detects explicit content from the GCS path to a video. """
    video_client = videointelligence.VideoIntelligenceServiceClient()
    features = [videointelligence.enums.Feature.EXPLICIT_CONTENT_DETECTION]

    operation = video_client.annotate_video(path, features=features)
    print('\nProcessing video for explicit content annotations:')

    result = operation.result(timeout=90)
    print('\nFinished processing.')

    likely_string = ("Unknown", "Very unlikely", "Unlikely", "Possible",
                     "Likely", "Very likely")

    # first result is retrieved because a single video was processed
    for frame in result.annotation_results[0].explicit_annotation.frames:
        frame_time = frame.time_offset.seconds + frame.time_offset.nanos / 1e9
        print('Time: {}s'.format(frame_time))
        print('\tpornography: {}'.format(
            likely_string[frame.pornography_likelihood]))

def analyze_faces(path):
    """ Detects faces given a GCS path. """
    video_client = videointelligence.VideoIntelligenceServiceClient()
    features = [videointelligence.enums.Feature.FACE_DETECTION]

    config = videointelligence.types.FaceDetectionConfig(
        include_bounding_boxes=True)
    context = videointelligence.types.VideoContext(
        face_detection_config=config)

    operation = video_client.annotate_video(
        path, features=features, video_context=context)
    print('\nProcessing video for face annotations:')

    result = operation.result(timeout=600)
    print('\nFinished processing.')

    # first result is retrieved because a single video was processed
    faces = result.annotation_results[0].face_annotations
    for face_id, face in enumerate(faces):
        print('Face {}'.format(face_id))
        print('Thumbnail size: {}'.format(len(face.thumbnail)))

        for segment_id, segment in enumerate(face.segments):
            start_time = (segment.segment.start_time_offset.seconds +
                          segment.segment.start_time_offset.nanos / 1e9)
            end_time = (segment.segment.end_time_offset.seconds +
                        segment.segment.end_time_offset.nanos / 1e9)
            positions = '{}s to {}s'.format(start_time, end_time)
            print('\tSegment {}: {}'.format(segment_id, positions))

        # There are typically many frames for each face,
        # here we print information on only the first frame.
        frame = face.frames[0]
        time_offset = (frame.time_offset.seconds +
                       frame.time_offset.nanos / 1e9)
        box = frame.normalized_bounding_boxes[0]
        print('First frame time offset: {}s'.format(time_offset))
        print('First frame normalized bounding box:')
        print('\tleft: {}'.format(box.left))
        print('\ttop: {}'.format(box.top))
        print('\tright: {}'.format(box.right))
        print('\tbottom: {}'.format(box.bottom))
        print('\n')

def analyze_labels(path):
    """ Detects labels given a GCS path. """
    video_client = videointelligence.VideoIntelligenceServiceClient()
    features = [videointelligence.enums.Feature.LABEL_DETECTION]

    mode = videointelligence.enums.LabelDetectionMode.SHOT_AND_FRAME_MODE
    config = videointelligence.types.LabelDetectionConfig(
        label_detection_mode=mode)
    context = videointelligence.types.VideoContext(
        label_detection_config=config)

    operation = video_client.annotate_video(
        path, features=features, video_context=context)
    print('\nProcessing video for label annotations:')

    result = operation.result(timeout=90)
    print('\nFinished processing.')

    # Process video/segment level label annotations
    segment_labels = result.annotation_results[0].segment_label_annotations
    for i, segment_label in enumerate(segment_labels):
        print('Video label description: {}'.format(
            segment_label.entity.description))
        for category_entity in segment_label.category_entities:
            print('\tLabel category description: {}'.format(
                category_entity.description))

        for i, segment in enumerate(segment_label.segments):
            start_time = (segment.segment.start_time_offset.seconds +
                          segment.segment.start_time_offset.nanos / 1e9)
            end_time = (segment.segment.end_time_offset.seconds +
                        segment.segment.end_time_offset.nanos / 1e9)
            positions = '{}s to {}s'.format(start_time, end_time)
            confidence = segment.confidence
            print('\tSegment {}: {}'.format(i, positions))
            print('\tConfidence: {}'.format(confidence))
        print('\n')

    # Process shot level label annotations
    shot_labels = result.annotation_results[0].shot_label_annotations
    for i, shot_label in enumerate(shot_labels):
        print('Shot label description: {}'.format(
            shot_label.entity.description))
        for category_entity in shot_label.category_entities:
            print('\tLabel category description: {}'.format(
                category_entity.description))

        for i, shot in enumerate(shot_label.segments):
            start_time = (shot.segment.start_time_offset.seconds +
                          shot.segment.start_time_offset.nanos / 1e9)
            end_time = (shot.segment.end_time_offset.seconds +
                        shot.segment.end_time_offset.nanos / 1e9)
            positions = '{}s to {}s'.format(start_time, end_time)
            confidence = shot.confidence
            print('\tSegment {}: {}'.format(i, positions))
            print('\tConfidence: {}'.format(confidence))
        print('\n')

    # Process frame level label annotations
    frame_labels = result.annotation_results[0].frame_label_annotations
    for i, frame_label in enumerate(frame_labels):
        print('Frame label description: {}'.format(
            frame_label.entity.description))
        for category_entity in frame_label.category_entities:
            print('\tLabel category description: {}'.format(
                category_entity.description))

        # Each frame_label_annotation has many frames,
        # here we print information only about the first frame.
        frame = frame_label.frames[0]
        time_offset = (frame.time_offset.seconds +
                       frame.time_offset.nanos / 1e9)
        print('\tFirst frame time offset: {}s'.format(time_offset))
        print('\tFirst frame confidence: {}'.format(frame.confidence))
        print('\n')

def analyze_labels_file(path):
    """ Detects labels given a file path. """
    video_client = videointelligence.VideoIntelligenceServiceClient()
    features = [videointelligence.enums.Feature.LABEL_DETECTION]

    with io.open(path, "rb") as movie:
        content_base64 = base64.b64encode(movie.read())

    operation = video_client.annotate_video(
        '', features=features, input_content=content_base64)
    print('\nProcessing video for label annotations:')

    result = operation.result(timeout=90)
    print('\nFinished processing.')

    # Process video/segment level label annotations
    segment_labels = result.annotation_results[0].segment_label_annotations
    for i, segment_label in enumerate(segment_labels):
        print('Video label description: {}'.format(
            segment_label.entity.description))
        for category_entity in segment_label.category_entities:
            print('\tLabel category description: {}'.format(
                category_entity.description))

        for i, segment in enumerate(segment_label.segments):
            start_time = (segment.segment.start_time_offset.seconds +
                          segment.segment.start_time_offset.nanos / 1e9)
            end_time = (segment.segment.end_time_offset.seconds +
                        segment.segment.end_time_offset.nanos / 1e9)
            positions = '{}s to {}s'.format(start_time, end_time)
            confidence = segment.confidence
            print('\tSegment {}: {}'.format(i, positions))
            print('\tConfidence: {}'.format(confidence))
        print('\n')

    # Process shot level label annotations
    shot_labels = result.annotation_results[0].shot_label_annotations
    for i, shot_label in enumerate(shot_labels):
        print('Shot label description: {}'.format(
            shot_label.entity.description))
        for category_entity in shot_label.category_entities:
            print('\tLabel category description: {}'.format(
                category_entity.description))

        for i, shot in enumerate(shot_label.segments):
            start_time = (shot.segment.start_time_offset.seconds +
                          shot.segment.start_time_offset.nanos / 1e9)
            end_time = (shot.segment.end_time_offset.seconds +
                        shot.segment.end_time_offset.nanos / 1e9)
            positions = '{}s to {}s'.format(start_time, end_time)
            confidence = shot.confidence
            print('\tSegment {}: {}'.format(i, positions))
            print('\tConfidence: {}'.format(confidence))
        print('\n')

    # Process frame level label annotations
    frame_labels = result.annotation_results[0].frame_label_annotations
    for i, frame_label in enumerate(frame_labels):
        print('Frame label description: {}'.format(
            frame_label.entity.description))
        for category_entity in frame_label.category_entities:
            print('\tLabel category description: {}'.format(
                category_entity.description))

        # Each frame_label_annotation has many frames,
        # here we print information only about the first frame.
        frame = frame_label.frames[0]
        time_offset = frame.time_offset.seconds + frame.time_offset.nanos / 1e9
        print('\tFirst frame time offset: {}s'.format(time_offset))
        print('\tFirst frame confidence: {}'.format(frame.confidence))
        print('\n')

def analyze_shots(path):
    """ Detects camera shot changes. """
    video_client = videointelligence.VideoIntelligenceServiceClient()
    features = [videointelligence.enums.Feature.SHOT_CHANGE_DETECTION]
    operation = video_client.annotate_video(path, features=features)
    print('\nProcessing video for shot change annotations:')

    result = operation.result(timeout=90)
    print('\nFinished processing.')

    # first result is retrieved because a single video was processed
    for i, shot in enumerate(result.annotation_results[0].shot_annotations):
        start_time = (shot.start_time_offset.seconds +
                      shot.start_time_offset.nanos / 1e9)
        end_time = (shot.end_time_offset.seconds +
                    shot.end_time_offset.nanos / 1e9)
        print('\tShot {}: {} to {}'.format(i, start_time, end_time))

if __name__ == '__main__':
    parser = argparse.ArgumentParser(
        description=__doc__,
        formatter_class=argparse.RawDescriptionHelpFormatter)
    subparsers = parser.add_subparsers(dest='command')
    analyze_faces_parser = subparsers.add_parser(
        'faces', help=analyze_faces.__doc__)
    analyze_faces_parser.add_argument('path')
    analyze_labels_parser = subparsers.add_parser(
        'labels', help=analyze_labels.__doc__)
    analyze_labels_parser.add_argument('path')
    analyze_labels_file_parser = subparsers.add_parser(
        'labels_file', help=analyze_labels_file.__doc__)
    analyze_labels_file_parser.add_argument('path')
    analyze_explicit_content_parser = subparsers.add_parser(
        'explicit_content', help=analyze_explicit_content.__doc__)
    analyze_explicit_content_parser.add_argument('path')
    analyze_shots_parser = subparsers.add_parser(
        'shots', help=analyze_shots.__doc__)
    analyze_shots_parser.add_argument('path')

    args = parser.parse_args()

    if args.command == 'faces':
        analyze_faces(args.path)
    if args.command == 'labels':
        analyze_labels(args.path)
    if args.command == 'labels_file':
        analyze_labels_file(args.path)
    if args.command == 'shots':
        analyze_shots(args.path)
    if args.command == 'explicit_content':
        analyze_explicit_content(args.path)
<commit_msg>Do not base64 encode the input content. (#1241)<commit_after>#!/usr/bin/env python

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

"""This application demonstrates face detection, label detection,
explicit content, and shot change detection using the Google Cloud API.

Usage Examples:

    python analyze.py faces gs://demomaker/google_gmail.mp4
    python analyze.py labels gs://cloud-ml-sandbox/video/chicago.mp4
    python analyze.py labels_file resources/cat.mp4
    python analyze.py shots gs://demomaker/gbikes_dinosaur.mp4
    python analyze.py explicit_content gs://demomaker/gbikes_dinosaur.mp4

"""

import argparse
import io

from google.cloud import videointelligence

def analyze_explicit_content(path):
    """ Detects explicit content from the GCS path to a video. """
    video_client = videointelligence.VideoIntelligenceServiceClient()
    features = [videointelligence.enums.Feature.EXPLICIT_CONTENT_DETECTION]

    operation = video_client.annotate_video(path, features=features)
    print('\nProcessing video for explicit content annotations:')

    result = operation.result(timeout=90)
    print('\nFinished processing.')

    likely_string = ("Unknown", "Very unlikely", "Unlikely", "Possible",
                     "Likely", "Very likely")

    # first result is retrieved because a single video was processed
    for frame in result.annotation_results[0].explicit_annotation.frames:
        frame_time = frame.time_offset.seconds + frame.time_offset.nanos / 1e9
        print('Time: {}s'.format(frame_time))
        print('\tpornography: {}'.format(
            likely_string[frame.pornography_likelihood]))

def analyze_faces(path):
    """ Detects faces given a GCS path. """
    video_client = videointelligence.VideoIntelligenceServiceClient()
    features = [videointelligence.enums.Feature.FACE_DETECTION]

    config = videointelligence.types.FaceDetectionConfig(
        include_bounding_boxes=True)
    context = videointelligence.types.VideoContext(
        face_detection_config=config)

    operation = video_client.annotate_video(
        path, features=features, video_context=context)
    print('\nProcessing video for face annotations:')

    result = operation.result(timeout=600)
    print('\nFinished processing.')

    # first result is retrieved because a single video was processed
    faces = result.annotation_results[0].face_annotations
    for face_id, face in enumerate(faces):
        print('Face {}'.format(face_id))
        print('Thumbnail size: {}'.format(len(face.thumbnail)))

        for segment_id, segment in enumerate(face.segments):
            start_time = (segment.segment.start_time_offset.seconds +
                          segment.segment.start_time_offset.nanos / 1e9)
            end_time = (segment.segment.end_time_offset.seconds +
                        segment.segment.end_time_offset.nanos / 1e9)
            positions = '{}s to {}s'.format(start_time, end_time)
            print('\tSegment {}: {}'.format(segment_id, positions))

        # There are typically many frames for each face,
        # here we print information on only the first frame.
        frame = face.frames[0]
        time_offset = (frame.time_offset.seconds +
                       frame.time_offset.nanos / 1e9)
        box = frame.normalized_bounding_boxes[0]
        print('First frame time offset: {}s'.format(time_offset))
        print('First frame normalized bounding box:')
        print('\tleft: {}'.format(box.left))
        print('\ttop: {}'.format(box.top))
        print('\tright: {}'.format(box.right))
        print('\tbottom: {}'.format(box.bottom))
        print('\n')

def analyze_labels(path):
    """ Detects labels given a GCS path. """
    video_client = videointelligence.VideoIntelligenceServiceClient()
    features = [videointelligence.enums.Feature.LABEL_DETECTION]

    mode = videointelligence.enums.LabelDetectionMode.SHOT_AND_FRAME_MODE
    config = videointelligence.types.LabelDetectionConfig(
        label_detection_mode=mode)
    context = videointelligence.types.VideoContext(
        label_detection_config=config)

    operation = video_client.annotate_video(
        path, features=features, video_context=context)
    print('\nProcessing video for label annotations:')

    result = operation.result(timeout=90)
    print('\nFinished processing.')

    # Process video/segment level label annotations
    segment_labels = result.annotation_results[0].segment_label_annotations
    for i, segment_label in enumerate(segment_labels):
        print('Video label description: {}'.format(
            segment_label.entity.description))
        for category_entity in segment_label.category_entities:
            print('\tLabel category description: {}'.format(
                category_entity.description))

        for i, segment in enumerate(segment_label.segments):
            start_time = (segment.segment.start_time_offset.seconds +
                          segment.segment.start_time_offset.nanos / 1e9)
            end_time = (segment.segment.end_time_offset.seconds +
                        segment.segment.end_time_offset.nanos / 1e9)
            positions = '{}s to {}s'.format(start_time, end_time)
            confidence = segment.confidence
            print('\tSegment {}: {}'.format(i, positions))
            print('\tConfidence: {}'.format(confidence))
        print('\n')

    # Process shot level label annotations
    shot_labels = result.annotation_results[0].shot_label_annotations
    for i, shot_label in enumerate(shot_labels):
        print('Shot label description: {}'.format(
            shot_label.entity.description))
        for category_entity in shot_label.category_entities:
            print('\tLabel category description: {}'.format(
                category_entity.description))

        for i, shot in enumerate(shot_label.segments):
            start_time = (shot.segment.start_time_offset.seconds +
                          shot.segment.start_time_offset.nanos / 1e9)
            end_time = (shot.segment.end_time_offset.seconds +
                        shot.segment.end_time_offset.nanos / 1e9)
            positions = '{}s to {}s'.format(start_time, end_time)
            confidence = shot.confidence
            print('\tSegment {}: {}'.format(i, positions))
            print('\tConfidence: {}'.format(confidence))
        print('\n')

    # Process frame level label annotations
    frame_labels = result.annotation_results[0].frame_label_annotations
    for i, frame_label in enumerate(frame_labels):
        print('Frame label description: {}'.format(
            frame_label.entity.description))
        for category_entity in frame_label.category_entities:
            print('\tLabel category description: {}'.format(
                category_entity.description))

        # Each frame_label_annotation has many frames,
        # here we print information only about the first frame.
        frame = frame_label.frames[0]
        time_offset = (frame.time_offset.seconds +
                       frame.time_offset.nanos / 1e9)
        print('\tFirst frame time offset: {}s'.format(time_offset))
        print('\tFirst frame confidence: {}'.format(frame.confidence))
        print('\n')

def analyze_labels_file(path):
    """Detect labels given a file path."""
    video_client = videointelligence.VideoIntelligenceServiceClient()
    features = [videointelligence.enums.Feature.LABEL_DETECTION]

    with io.open(path, 'rb') as movie:
        input_content = movie.read()

    operation = video_client.annotate_video(
        '', features=features, input_content=input_content)
    print('\nProcessing video for label annotations:')

    result = operation.result(timeout=90)
    print('\nFinished processing.')

    # Process video/segment level label annotations
    segment_labels = result.annotation_results[0].segment_label_annotations
    for i, segment_label in enumerate(segment_labels):
        print('Video label description: {}'.format(
            segment_label.entity.description))
        for category_entity in segment_label.category_entities:
            print('\tLabel category description: {}'.format(
                category_entity.description))

        for i, segment in enumerate(segment_label.segments):
            start_time = (segment.segment.start_time_offset.seconds +
                          segment.segment.start_time_offset.nanos / 1e9)
            end_time = (segment.segment.end_time_offset.seconds +
                        segment.segment.end_time_offset.nanos / 1e9)
            positions = '{}s to {}s'.format(start_time, end_time)
            confidence = segment.confidence
            print('\tSegment {}: {}'.format(i, positions))
            print('\tConfidence: {}'.format(confidence))
        print('\n')

    # Process shot level label annotations
    shot_labels = result.annotation_results[0].shot_label_annotations
    for i, shot_label in enumerate(shot_labels):
        print('Shot label description: {}'.format(
            shot_label.entity.description))
        for category_entity in shot_label.category_entities:
            print('\tLabel category description: {}'.format(
                category_entity.description))

        for i, shot in enumerate(shot_label.segments):
            start_time = (shot.segment.start_time_offset.seconds +
                          shot.segment.start_time_offset.nanos / 1e9)
            end_time = (shot.segment.end_time_offset.seconds +
                        shot.segment.end_time_offset.nanos / 1e9)
            positions = '{}s to {}s'.format(start_time, end_time)
            confidence = shot.confidence
            print('\tSegment {}: {}'.format(i, positions))
            print('\tConfidence: {}'.format(confidence))
        print('\n')

    # Process frame level label annotations
    frame_labels = result.annotation_results[0].frame_label_annotations
    for i, frame_label in enumerate(frame_labels):
        print('Frame label description: {}'.format(
            frame_label.entity.description))
        for category_entity in frame_label.category_entities:
            print('\tLabel category description: {}'.format(
                category_entity.description))

        # Each frame_label_annotation has many frames,
        # here we print information only about the first frame.
        frame = frame_label.frames[0]
        time_offset = frame.time_offset.seconds + frame.time_offset.nanos / 1e9
        print('\tFirst frame time offset: {}s'.format(time_offset))
        print('\tFirst frame confidence: {}'.format(frame.confidence))
        print('\n')

def analyze_shots(path):
    """ Detects camera shot changes. """
    video_client = videointelligence.VideoIntelligenceServiceClient()
    features = [videointelligence.enums.Feature.SHOT_CHANGE_DETECTION]
    operation = video_client.annotate_video(path, features=features)
    print('\nProcessing video for shot change annotations:')

    result = operation.result(timeout=90)
    print('\nFinished processing.')

    # first result is retrieved because a single video was processed
    for i, shot in enumerate(result.annotation_results[0].shot_annotations):
        start_time = (shot.start_time_offset.seconds +
                      shot.start_time_offset.nanos / 1e9)
        end_time = (shot.end_time_offset.seconds +
                    shot.end_time_offset.nanos / 1e9)
        print('\tShot {}: {} to {}'.format(i, start_time, end_time))

if __name__ == '__main__':
    parser = argparse.ArgumentParser(
        description=__doc__,
        formatter_class=argparse.RawDescriptionHelpFormatter)
    subparsers = parser.add_subparsers(dest='command')
    analyze_faces_parser = subparsers.add_parser(
        'faces', help=analyze_faces.__doc__)
    analyze_faces_parser.add_argument('path')
    analyze_labels_parser = subparsers.add_parser(
        'labels', help=analyze_labels.__doc__)
    analyze_labels_parser.add_argument('path')
    analyze_labels_file_parser = subparsers.add_parser(
        'labels_file', help=analyze_labels_file.__doc__)
    analyze_labels_file_parser.add_argument('path')
    analyze_explicit_content_parser = subparsers.add_parser(
        'explicit_content', help=analyze_explicit_content.__doc__)
    analyze_explicit_content_parser.add_argument('path')
    analyze_shots_parser = subparsers.add_parser(
        'shots', help=analyze_shots.__doc__)
    analyze_shots_parser.add_argument('path')

    args = parser.parse_args()

    if args.command == 'faces':
        analyze_faces(args.path)
    if args.command == 'labels':
        analyze_labels(args.path)
    if args.command == 'labels_file':
        analyze_labels_file(args.path)
    if args.command == 'shots':
        analyze_shots(args.path)
    if args.command == 'explicit_content':
        analyze_explicit_content(args.path)
<|endoftext|>


<commit_before>from keras.preprocessing.image import ImageDataGenerator
from keras.models import Sequential
from keras.layers import Convolution2D, MaxPooling2D
from keras.layers.convolutional import Conv2D
from keras.layers import Activation, Dropout, Flatten, Dense
import h5py
import configs
import numpy as np
import os
from keras.utils import plot_model

class VisionDataProcessor():

    def __init__(self):
        self.data_generator = ImageDataGenerator(
        shear_range = configs.shear_range,
        zoom_range = configs.zoom_range,
        zca_whitening = configs.zca_whitening,
        rotation_range = configs.rotation_range,
        width_shift_range = configs.width_shift_range,
        height_shift_range = configs.height_shift_range,
        vertical_flip = configs.vertical_flip,
        horizontal_flip = configs.horizontal_flip,
        )

        self.train_generator = self.create_data_generator_from_directory(configs.test_dir)
        self.validation_generator = self.create_data_generator_from_directory(configs.val_dir)

    def create_data_generator_from_directory(self, directory):
        generated_generator = self.data_generator.flow_from_directory(
        directory = directory,
        target_size = (configs.img_width, configs.img_height),
        batch_size = configs.batch_size,
        color_mode = configs.color_mode,
        class_mode = configs.class_mode,
        save_to_dir = 'visualize'
        )

        return generated_generator

    def generate_h5py_files(self):
        test_file = h5py.File(configs.test_fname, "w")
        self.test_image_data = test_file.create_dataset(
            "image_data",
            (configs.nb_test_images, configs.img_width, configs.img_height),
            dtype = "float32"
        )
        self.test_class_data = test_file.create_dataset(
            "class_data",
            (configs.nb_test_images, configs.nb_classes),
            dtype = "float32"
        )

        val_file = h5py.File(configs.val_fname, "w")
        self.val_image_data = val_file.create_dataset(
            "image_data",
            (configs.nb_val_images, configs.img_width, configs.img_height),
            dtype = 'float32'
        )
        self.val_class_data = val_file.create_dataset(
            "class_data",
            (configs.nb_val_images, configs.nb_classes),
            dtype = "float32"
        )

    def load_h5py_with_generator(self,counter_limit,data_generator):
        #Can't figure out how to remove the extra dimension
        print(type(self.validation_generator))
        counter = 0
        while counter != counter_limit:
            data_pack = data_generator.next() # tuple
            image_data = data_pack[0]
            class_data = data_pack[1]

            for index in range(len(image_data)):
                print(image_data[index][0].shape)
                print(self.test_image_data[counter][0].shape)
                self.test_image_data[counter] = image_data[index]
                self.test_class_data[counter] = class_data[index]
                counter += 1
                if counter == configs.nb_test_images:
                    break
            import sys
            sys.stdout.write("  " + str(counter) + "/" + str(configs.nb_test_images) + "\r")
            sys.stdout.flush()
        test_file.close()

    def create_simple_keras_model(self):
        input_shape = (1, configs.img_width, configs.img_height)

        model = Sequential()

        model.add(Conv2D(32, (3, 3),
            padding='same',
            data_format='channels_first',
            input_shape=input_shape))
            # now: model.output_shape == (None, 64, 32, 32)

        model.add(Activation('relu'))

        model.add(MaxPooling2D(pool_size=(2, 2)))

        model.add(Flatten())
        model.add(Dense(75))
        model.add(Dense(40))
        model.add(Dense(10))

        model.add(Activation('relu'))

        model.add(Dropout(0.75))

        model.add(Dense(configs.nb_classes))
        model.add(Activation('sigmoid'))

        if configs.print_summary:
            model.summary()
        model.compile(optimizer='rmsprop',
              loss='sparse_categorical_crossentropy',
              metrics=['accuracy'])

        self.model = model

    def fit_simple_keras_model(self):
        self.model.fit_generator(
            self.train_generator,
            steps_per_epoch=int(nb_train_images/configs.batch_size),
            epochs=configs.nb_epoch,
            validation_data=self.validation_generator,
            validation_steps=(nb_val_images/configs.batch_size)

if __name__ == '__main__':
    dataprocessor = VisionDataProcessor()
    dataprocessor.create_simple_keras_model()
    dataprocessor.fit_simple_keras_model()
    print("Done.")
<commit_msg>Hanging bracket fix<commit_after>from keras.preprocessing.image import ImageDataGenerator
from keras.models import Sequential
from keras.layers import Convolution2D, MaxPooling2D
from keras.layers.convolutional import Conv2D
from keras.layers import Activation, Dropout, Flatten, Dense
import h5py
import configs
import numpy as np
import os
from keras.utils import plot_model

class VisionDataProcessor():

    def __init__(self):
        self.data_generator = ImageDataGenerator(
        shear_range = configs.shear_range,
        zoom_range = configs.zoom_range,
        zca_whitening = configs.zca_whitening,
        rotation_range = configs.rotation_range,
        width_shift_range = configs.width_shift_range,
        height_shift_range = configs.height_shift_range,
        vertical_flip = configs.vertical_flip,
        horizontal_flip = configs.horizontal_flip,
        )

        self.train_generator = self.create_data_generator_from_directory(configs.test_dir)
        self.validation_generator = self.create_data_generator_from_directory(configs.val_dir)

    def create_data_generator_from_directory(self, directory):
        generated_generator = self.data_generator.flow_from_directory(
        directory = directory,
        target_size = (configs.img_width, configs.img_height),
        batch_size = configs.batch_size,
        color_mode = configs.color_mode,
        class_mode = configs.class_mode,
        save_to_dir = 'visualize'
        )

        return generated_generator

    def generate_h5py_files(self):
        test_file = h5py.File(configs.test_fname, "w")
        self.test_image_data = test_file.create_dataset(
            "image_data",
            (configs.nb_test_images, configs.img_width, configs.img_height),
            dtype = "float32"
        )
        self.test_class_data = test_file.create_dataset(
            "class_data",
            (configs.nb_test_images, configs.nb_classes),
            dtype = "float32"
        )

        val_file = h5py.File(configs.val_fname, "w")
        self.val_image_data = val_file.create_dataset(
            "image_data",
            (configs.nb_val_images, configs.img_width, configs.img_height),
            dtype = 'float32'
        )
        self.val_class_data = val_file.create_dataset(
            "class_data",
            (configs.nb_val_images, configs.nb_classes),
            dtype = "float32"
        )

    def load_h5py_with_generator(self,counter_limit,data_generator):
        #Can't figure out how to remove the extra dimension
        print(type(self.validation_generator))
        counter = 0
        while counter != counter_limit:
            data_pack = data_generator.next() # tuple
            image_data = data_pack[0]
            class_data = data_pack[1]

            for index in range(len(image_data)):
                print(image_data[index][0].shape)
                print(self.test_image_data[counter][0].shape)
                self.test_image_data[counter] = image_data[index]
                self.test_class_data[counter] = class_data[index]
                counter += 1
                if counter == configs.nb_test_images:
                    break
            import sys
            sys.stdout.write("  " + str(counter) + "/" + str(configs.nb_test_images) + "\r")
            sys.stdout.flush()
        test_file.close()

    def create_simple_keras_model(self):
        input_shape = (1, configs.img_width, configs.img_height)

        model = Sequential()

        model.add(Conv2D(32, (3, 3),
            padding='same',
            data_format='channels_first',
            input_shape=input_shape))
            # now: model.output_shape == (None, 64, 32, 32)

        model.add(Activation('relu'))

        model.add(MaxPooling2D(pool_size=(2, 2)))

        model.add(Flatten())
        model.add(Dense(75))
        model.add(Dense(40))
        model.add(Dense(10))

        model.add(Activation('relu'))

        model.add(Dropout(0.75))

        model.add(Dense(configs.nb_classes))
        model.add(Activation('sigmoid'))

        if configs.print_summary:
            model.summary()
        model.compile(optimizer='rmsprop',
              loss='sparse_categorical_crossentropy',
              metrics=['accuracy'])

        self.model = model

    def fit_simple_keras_model(self):
        self.model.fit_generator(
            self.train_generator,
            steps_per_epoch=int(nb_train_images/configs.batch_size),
            epochs=configs.nb_epoch,
            validation_data=self.validation_generator,
            validation_steps=(nb_val_images/configs.batch_size)
            )

if __name__ == '__main__':
    dataprocessor = VisionDataProcessor()
    dataprocessor.create_simple_keras_model()
    dataprocessor.fit_simple_keras_model()
    print("Done.")
<|endoftext|>


<commit_before># encoding: UTF-8
"""
    Author: vigarbuaa
"""

import hashlib
import hmac
import sys
import time
from copy import copy
from threading import Lock
from datetime import datetime
from urllib.parse import urlencode
from vnpy.api.rest import Request, RestClient
from vnpy.api.websocket import WebsocketClient

from vnpy.trader.constant import (
    Direction,
    Exchange,
    OrderType,
    Product,
    Status,
    Offset,
    Interval
)
from vnpy.trader.gateway import BaseGateway
from vnpy.trader.object import (
    TickData,
    OrderData,
    TradeData,
    PositionData, AccountData,
       ContractData,
    OrderRequest,
    CancelRequest,
    SubscribeRequest,
)

REST_HOST = "https://api.alpaca.markets"
PAPER_REST_HOST = "https://paper-api.alpaca.markets"
KEY = ""
SECRET = ""

STATUS_ALPACA2VT = {
    "new": Status.SUBMITTING,
    "open": Status.NOTTRADED,
    "part_filled": Status.PARTTRADED,
    "filled": Status.ALLTRADED,
    "cancelled": Status.CANCELLED,
    "cancelling": Status.CANCELLED,
    "failure": Status.REJECTED,
}

DIRECTION_VT2ALPACA = {Direction.LONG: "buy", Direction.SHORT: "sell"}
DIRECTION_ALPACA2VT = {"buy": Direction.LONG, "sell": Direction.SHORT,
                       "long": Direction.LONG, "short": Direction.SHORT}

ORDERTYPE_VT2ALPACA = {
    OrderType.LIMIT: "limit",
    OrderType.MARKET: "market"
}
ORDERTYPE_ALPACA2VT = {v: k for k, v in ORDERTYPE_VT2ALPACA.items()}

class AlpacaGateway(BaseGateway):
    """
    VN Trader Gateway for Alpaca connection.
    """

    default_setting = {
        "key": "",
        "secret": "",
        "session": 3,
        "proxy_host": "127.0.0.1",
        "proxy_port": 1080,
    }

    def __init__(self, event_engine):
        """Constructor"""
        super(AlpacaGateway, self).__init__(event_engine, "ALPACA")

        self.rest_api = AlpacaRestApi(self)
        #self.ws_api = AlpacaWebsocketApi(self)
        self.order_map = {}

    def connect(self, setting: dict):
        """"""
        key = setting["key"]
        secret = setting["secret"]
        session = setting["session"]
        proxy_host = setting["proxy_host"]
        proxy_port = setting["proxy_port"]

        self.rest_api.connect(key, secret, session, proxy_host, proxy_port)

        #self.ws_api.connect(key, secret, proxy_host, proxy_port)

    def subscribe(self, req: SubscribeRequest):
        """"""
        #self.ws_api.subscribe(req)
        pass

    def send_order(self, req: OrderRequest):
        """"""
        return self.rest_api.send_order(req)

    def cancel_order(self, req: CancelRequest):
        """"""
        #self.ws_api.cancel_order(req)
        pass

    def query_account(self):
        """"""
        self.rest_api.query_account()

    def query_position(self):
        """"""
        self.rest_api.query_position()

    def close(self):
        """"""
        self.rest_api.stop()
        self.ws_api.stop()

class AlpacaRestApi(RestClient):
    """
    Alpaca REST API
    """

    def __init__(self, gateway: BaseGateway):
        """"""
        super(AlpacaRestApi, self).__init__()

        self.gateway = gateway
        self.gateway_name = gateway.gateway_name

        self.key = ""
        self.secret = ""

        self.order_count = 1_000_000
        self.order_count_lock = Lock()
        self.connect_time = 0
        print(
            self.gateway_name,
            self.key,
            self.secret,
            self.order_count,
            self.connect_time)

    def query_account(self):
        print("call query_account")
        path = f"/v1/account"
        self.add_request(
            method="GET",
            path=path,
            callback=self.on_query_account
        )

    def on_query_account(self, data, request):
        print("debug on_query_account: ", data, request)
        account = AccountData(
            accountid=data['id'],
            balance=float(data['cash']),
            frozen=float(data['cash']) - float(data['buying_power']),
            gateway_name=self.gateway_name
        )
        self.gateway.on_account(account)

    def query_position(self):
        print("call query_position")
        path = f"/v1/positions"
        self.add_request(
            method="GET",
            path=path,
            callback=self.on_query_position
        )

    def on_query_position(self, data, request):
        print("debug on_query_position: ", data,"!!!", request)
        for d in data:
            position = PositionData(
                symbol=d['symbol'],
                exchange=Exchange.ALPACA,
                direction=DIRECTION_ALPACA2VT[d['side']],
                volume=d['qty'],
                price=d['avg_entry_price'],
                pnl=d['unrealized_pl'],
                gateway_name=self.gateway_name,
            )
            self.gateway.on_position(position)

    def sign(self, request):
        """
        Generate Alpaca signature.
        """
        #msg = request.method + \
        #    "/api/v1/{}{}{}".format(path, nonce, request.data)
        headers = {
            "APCA-API-KEY-ID": self.key,
            "APCA-API-SECRET-KEY": self.secret,
        }

        request.headers = headers
        request.allow_redirects=False
        return request

    def _new_order_id(self):
        with self.order_count_lock:
            self.order_count += 1
            return self.order_count

    def connect(
        self,
        key: str,
        secret: str,
        session_num: int,
        proxy_host: str,
        proxy_port: int,
    ):
        """
               Initialize connection to REST server.
        """
        self.key = key
        self.secret = secret

        self.init(PAPER_REST_HOST, proxy_host, proxy_port)
        print("rest connect: ", PAPER_REST_HOST, proxy_host, proxy_port)
        self.start(session_num)
        self.connect_time = (
            int(datetime.now().strftime("%y%m%d%H%M%S")) * self.order_count
        )
        print("rest client connected",self.connect_time)
        self.gateway.write_log("ALPACA REST API启动成功")
        self.query_account()
        self.query_position()

    def on_send_order(self, status_code: int, request: Request):
        print("debug on send order: ", status_code, request)
        pass

    def on_failed(self, status_code: int, request: Request):
        """
        Callback to handle request failed.
        """
        msg = f"请求失败，状态码：{status_code}，信息：{request.response.text}"
        print('debug on_failed', msg)
        self.gateway.write_log(msg)

    def on_error(
        self, exception_type: type, exception_value: Exception, tb, request: Request
    ):
        """
        Callback to handler request exception.
        """
        msg = f"触发异常，状态码：{exception_type}，信息：{exception_value}"
        print('debug on_error', msg)
        self.gateway.write_log(msg)

        sys.stderr.write(
            self.exception_detail(exception_type, exception_value, tb, request)
        )

    def send_order(self, req: OrderRequest):
        orderid = str(self.connect_time + self._new_order_id())
        raw_dict={
            "symbol": req.symbol,
            "qty": int(req.volume),
            "side": DIRECTION_VT2ALPACA[req.direction],
            "type": ORDERTYPE_VT2ALPACA[req.type],
            "time_in_force":'day',
        }
        if raw_dict['type'] == "limit":
            raw_dict['limit_price'] = float(req.price)

        data = raw_dict
        order = req.create_order_data(orderid, self.gateway_name)
        self.add_request(
            "POST",
            "/v1/orders",
            callback=self.on_send_order,
            #data=data,
            extra=order,
            on_failed=self.on_failed,
            on_error=self.on_error,
            json_str=data,
        )
        self.gateway.on_order(order)
        return order.vt_orderid

class AlpacaWebsocketApi(WebsocketClient):
    """"""

    def __init__(self, gateway):
        """"""
        super(AlpacaWebsocketApi, self).__init__()

        self.gateway = gateway
        self.gateway_name = gateway.gateway_name
        self.order_id = 1_000_000
        # self.date = int(datetime.now().strftime('%y%m%d%H%M%S')) * self.orderId
        self.key = ""
        self.secret = ""

        self.callbacks = {
            "trade": self.on_tick,
            "orderBook10": self.on_depth,
            "execution": self.on_trade,
            "order": self.on_order,
            "position": self.on_position,
            "margin": self.on_account,
            "instrument": self.on_contract,
        }

        self.ticks = {}
        self.accounts = {}
        self.orders = {}
        self.trades = set()
        self.tickDict = {}
        self.bidDict = {}
        self.askDict = {}
        self.orderLocalDict = {}
        self.channelDict = {}       # ChannelID : (Channel, Symbol)

    def connect(
        self, key: str, secret: str, proxy_host: str, proxy_port: int
    ):
        """"""
        self.key = key
        self.secret = secret.encode()

        self.init(WEBSOCKET_HOST, proxy_host, proxy_port)

        self.start()

    def subscribe(self, req: SubscribeRequest):
        pass

    def send_order(self, req: OrderRequest):
        pass

    # ----------------------------------------------------------------------
    def cancel_order(self, req: CancelRequest):
        """"""
        pass

    def on_connected(self):
        """"""
        self.gateway.write_log("Websocket API连接成功")
        self.authenticate()

    def on_disconnected(self):
        """"""
        self.gateway.write_log("Websocket API连接断开")

    # need debug 20190404
    def on_packet(self, packet: dict):
        """"""
        print("debug on_packet: " , packet)

    # ----------------------------------------------------------------------
    def on_response(self, data):
        """"""
        pass
    # ----------------------------------------------------------------------
    def on_update(self, data):
        """"""
        pass

    # ----------------------------------------------------------------------
    def on_wallet(self, data):
        """"""
        pass

    # ----------------------------------------------------------------------
    def on_trade_update(self, data):
        """"""
        pass

    def on_error(self, exception_type: type, exception_value: Exception, tb):
        """"""
        print("on_error: " ,type,Exception,tb)
        pass

    # debug OK , 0405
    def authenticate(self):
        pass

    def subscribe_topic(self):
        pass

    def on_tick(self, d):
        """"""
        pass

    def on_depth(self, d):
        """"""
        pass

    def on_trade(self, d):
        """"""
        pass

    def generateDateTime(self, s):
        """生成时间"""
        dt = datetime.fromtimestamp(s / 1000.0)
        time = dt.strftime("%H:%M:%S.%f")
        return time

    def on_order(self, data):
        """"""
        pass

    def on_position(self, d):
        """"""
        pass

    def on_account(self, d):
        """"""
        pass

    def on_contract(self, d):
        """"""
        pass<commit_msg>format<commit_after># encoding: UTF-8
"""
    Author: vigarbuaa
"""

import hashlib
import hmac
import sys
import time
from copy import copy
from threading import Lock
from datetime import datetime
from urllib.parse import urlencode
from vnpy.api.rest import Request, RestClient
from vnpy.api.websocket import WebsocketClient

from vnpy.trader.constant import (
    Direction,
    Exchange,
    OrderType,
    Product,
    Status,
    Offset,
    Interval
)
from vnpy.trader.gateway import BaseGateway
from vnpy.trader.object import (
    TickData,
    OrderData,
    TradeData,
    PositionData, AccountData,
    ContractData,
    OrderRequest,
    CancelRequest,
    SubscribeRequest,
)

REST_HOST = "https://api.alpaca.markets"
PAPER_REST_HOST = "https://paper-api.alpaca.markets"
KEY = ""
SECRET = ""

STATUS_ALPACA2VT = {
    "new": Status.SUBMITTING,
    "open": Status.NOTTRADED,
    "part_filled": Status.PARTTRADED,
    "filled": Status.ALLTRADED,
    "cancelled": Status.CANCELLED,
    "cancelling": Status.CANCELLED,
    "failure": Status.REJECTED,
}

DIRECTION_VT2ALPACA = {Direction.LONG: "buy", Direction.SHORT: "sell"}
DIRECTION_ALPACA2VT = {"buy": Direction.LONG, "sell": Direction.SHORT,
                       "long": Direction.LONG, "short": Direction.SHORT}

ORDERTYPE_VT2ALPACA = {
    OrderType.LIMIT: "limit",
    OrderType.MARKET: "market"
}
ORDERTYPE_ALPACA2VT = {v: k for k, v in ORDERTYPE_VT2ALPACA.items()}

class AlpacaGateway(BaseGateway):
    """
    VN Trader Gateway for Alpaca connection.
    """

    default_setting = {
        "key": "",
        "secret": "",
        "session": 3,
        "proxy_host": "127.0.0.1",
        "proxy_port": 1080,
    }

    def __init__(self, event_engine):
        """Constructor"""
        super(AlpacaGateway, self).__init__(event_engine, "ALPACA")

        self.rest_api = AlpacaRestApi(self)
        #self.ws_api = AlpacaWebsocketApi(self)
        self.order_map = {}

    def connect(self, setting: dict):
        """"""
        key = setting["key"]
        secret = setting["secret"]
        session = setting["session"]
        proxy_host = setting["proxy_host"]
        proxy_port = setting["proxy_port"]

        self.rest_api.connect(key, secret, session, proxy_host, proxy_port)

        #self.ws_api.connect(key, secret, proxy_host, proxy_port)

    def subscribe(self, req: SubscribeRequest):
        """"""
        # self.ws_api.subscribe(req)
        pass

    def send_order(self, req: OrderRequest):
        """"""
        return self.rest_api.send_order(req)

    def cancel_order(self, req: CancelRequest):
        """"""
        self.rest_api.cancel_order(req)

    def query_account(self):
        """"""
        self.rest_api.query_account()

    def query_position(self):
        """"""
        self.rest_api.query_position()

    def close(self):
        """"""
        self.rest_api.stop()
        self.ws_api.stop()

class AlpacaRestApi(RestClient):
    """
    Alpaca REST API
    """

    def __init__(self, gateway: BaseGateway):
        """"""
        super(AlpacaRestApi, self).__init__()

        self.gateway = gateway
        self.gateway_name = gateway.gateway_name

        self.key = ""
        self.secret = ""

        self.order_count = 1_000_000
        self.order_count_lock = Lock()
        self.connect_time = 0
        self.order_dict={}

    def query_account(self):
        print("call query_account")
        path = f"/v1/account"
        self.add_request(
            method="GET",
            path=path,
            callback=self.on_query_account
        )

    def on_query_account(self, data, request):
        print("debug on_query_account: ", data, request)
        account = AccountData(
            accountid=data['id'],
            balance=float(data['cash']),
            frozen=float(data['cash']) - float(data['buying_power']),
            gateway_name=self.gateway_name
        )
        self.gateway.on_account(account)

    def query_position(self):
        print("call query_position")
        path = f"/v1/positions"
        self.add_request(
            method="GET",
            path=path,
            callback=self.on_query_position
        )

    def on_query_position(self, data, request):
        print("debug on_query_position: ", data, "!!!", request)
        for d in data:
            position = PositionData(
                symbol=d['symbol'],
                exchange=Exchange.ALPACA,
                direction=DIRECTION_ALPACA2VT[d['side']],
                volume=d['qty'],
                price=d['avg_entry_price'],
                pnl=d['unrealized_pl'],
                gateway_name=self.gateway_name,
            )
            self.gateway.on_position(position)

    def sign(self, request):
        """
        Generate Alpaca signature.
        """
        headers = {
            "APCA-API-KEY-ID": self.key,
            "APCA-API-SECRET-KEY": self.secret,
        }

        request.headers = headers
        request.allow_redirects = False
        return request

    def _new_order_id(self):
        with self.order_count_lock:
            self.order_count += 1
            return self.order_count

    def connect(
        self,
        key: str,
        secret: str,
        session_num: int,
        proxy_host: str,
        proxy_port: int,
    ):
        """
               Initialize connection to REST server.
        """
        self.key = key
        self.secret = secret

        self.init(PAPER_REST_HOST, proxy_host, proxy_port)
        print("rest connect: ", PAPER_REST_HOST, proxy_host, proxy_port)
        self.start(session_num)
        self.connect_time = (
            int(datetime.now().strftime("%y%m%d%H%M%S")) * self.order_count
        )
        print("rest client connected", self.connect_time)
        self.gateway.write_log("ALPACA REST API启动成功")
        self.query_account()
        self.query_position()

    def on_send_order(self, status_code: int, request: Request):
        print("debug on send order: ", status_code, request)
        order = request.extra
        self.gateway.on_order(order)

    def on_failed_order(self, status_code: int, request: Request):
        """
        Callback to handle request failed.
        """
        order = request.extra
        order.status = Status.REJECTED
        self.gateway.on_order(order)
        msg = f"请求失败，状态码：{status_code}，信息：{request.response.text}"
        print('debug on_failed', msg)
        self.gateway.write_log(msg)

    def on_error_order(
        self, exception_type: type, exception_value: Exception, tb, request: Request
    ):
        """
        Callback to handler request exception.
        """
        order = request.extra
        order.status = Status.REJECTED
        self.gateway.on_order(order)
        msg = f"触发异常，状态码：{exception_type}，信息：{exception_value}"
        print('debug on_error', msg)
        self.gateway.write_log(msg)

        sys.stderr.write(
            self.exception_detail(exception_type, exception_value, tb, request)
        )

    def cancel_order(self, req: CancelRequest):
        """"""
        order_id = req.orderid
        print("come to cancel_order", order_id)

    def send_order(self, req: OrderRequest):
        orderid = str(self.connect_time + self._new_order_id())
        raw_dict = {
            "symbol": req.symbol,
            "qty": int(req.volume),
            "side": DIRECTION_VT2ALPACA[req.direction],
            "type": ORDERTYPE_VT2ALPACA[req.type],
            "time_in_force": 'day',
        }
        if raw_dict['type'] == "limit":
            raw_dict['limit_price'] = float(req.price)

        data = raw_dict
        order = req.create_order_data(orderid, self.gateway_name)
        self.add_request(
            "POST",
            "/v1/orders",
            callback=self.on_send_order,
            # data=data,
            extra=order,
            on_failed=self.on_failed_order,
            on_error=self.on_error_order,
            json_str=data,
        )
        return order.vt_orderid

class AlpacaWebsocketApi(WebsocketClient):
    """"""

    def __init__(self, gateway):
        """"""
        super(AlpacaWebsocketApi, self).__init__()

        self.gateway = gateway
        self.gateway_name = gateway.gateway_name
        self.order_id = 1_000_000
        # self.date = int(datetime.now().strftime('%y%m%d%H%M%S')) * self.orderId
        self.key = ""
        self.secret = ""

        self.callbacks = {
            "trade": self.on_tick,
            "orderBook10": self.on_depth,
            "execution": self.on_trade,
            "order": self.on_order,
            "position": self.on_position,
            "margin": self.on_account,
            "instrument": self.on_contract,
        }

        self.ticks = {}
        self.accounts = {}
        self.orders = {}
        self.trades = set()
        self.tickDict = {}
        self.bidDict = {}
        self.askDict = {}
        self.orderLocalDict = {}
        self.channelDict = {}       # ChannelID : (Channel, Symbol)

    def connect(
        self, key: str, secret: str, proxy_host: str, proxy_port: int
    ):
        """"""
        self.key = key
        self.secret = secret.encode()
        self.init(WEBSOCKET_HOST, proxy_host, proxy_port)
        self.start()

    def subscribe(self, req: SubscribeRequest):
        pass

    def send_order(self, req: OrderRequest):
        pass

    # ----------------------------------------------------------------------
    def cancel_order(self, req: CancelRequest):
        """"""
        pass

    def on_connected(self):
        """"""
        self.gateway.write_log("Websocket API连接成功")
        self.authenticate()

    def on_disconnected(self):
        """"""
        self.gateway.write_log("Websocket API连接断开")

    # need debug 20190404
    def on_packet(self, packet: dict):
        """"""
        print("debug on_packet: ", packet)

    # ----------------------------------------------------------------------
    def on_response(self, data):
        """"""
        pass
    # ----------------------------------------------------------------------

    def on_update(self, data):
        """"""
        pass

    # ----------------------------------------------------------------------
    def on_wallet(self, data):
        """"""
        pass

    # ----------------------------------------------------------------------
    def on_trade_update(self, data):
        """"""
        pass

    def on_error(self, exception_type: type, exception_value: Exception, tb):
        """"""
        print("on_error: ", type, Exception, tb)
        pass

    # debug OK , 0405
    def authenticate(self):
        pass

    def subscribe_topic(self):
        pass

    def on_tick(self, d):
        """"""
        pass

    def on_depth(self, d):
        """"""
        pass

    def on_trade(self, d):
        """"""
        pass

    def generateDateTime(self, s):
        """生成时间"""
        dt = datetime.fromtimestamp(s / 1000.0)
        time = dt.strftime("%H:%M:%S.%f")
        return time

    def on_order(self, data):
        """"""
        pass

    def on_position(self, d):
        """"""
        pass

    def on_account(self, d):
        """"""
        pass

    def on_contract(self, d):
        """"""
        pass
<|endoftext|>


<commit_before># -*- coding: utf-8 -*-

__all__ = ['MapCubeAnimator']

from sunpy.visualization import imageanimator

class MapCubeAnimator(imageanimator.BaseFuncAnimator):
    """
    Create an interactive viewer for a MapCube

    The following keyboard shortcuts are defined in the viewer:

    - 'left': previous step on active slider
    - 'right': next step on active slider
    - 'top': change the active slider up one
    - 'bottom': change the active slider down one
    - 'p': play/pause active slider

    Parameters
    ----------
    mapcube: sunpy.map.MapCube
        A MapCube

    annotate: bool
        Annotate the figure with scale and titles

    fig: mpl.figure
        Figure to use

    interval: int
        Animation interval in ms

    colorbar: bool
        Plot colorbar

    Notes
    -----
    Extra keywords are passed to `mapcube[0].plot()` i.e. the `plot()` routine of
    the first map in the cube.
    """
    def __init__(self, mapcube, annotate=True, **kwargs):

        self.mapcube = mapcube
        self.annotate = annotate
        slider_functions = [self.updatefig]
        slider_ranges = [[0,len(mapcube._maps)]]

        imageanimator.BaseFuncAnimator.__init__(self, mapcube._maps, slider_functions,
                                        slider_ranges, **kwargs)

        self._annotate_plot(0)

    def updatefig(self, val, im, slider):
        i = int(val)
        im.set_array(self.data[i].data)
        im.set_cmap(self.mapcube[i].cmap)
        im.set_norm(self.mapcube[i].norm)
        # Having this line in means the plot will resize for non-homogenous
        # maps. However it also means that if you zoom in on the plot bad
        # things happen.
        # im.set_extent(self.mapcube[i].xrange + self.mapcube[i].yrange)
        if self.annotate:
            self._annotate_plot(i)

    def _annotate_plot(self, ind):
        """
        Annotate the image.

        This may overwrite some stuff in `GenericMap.plot()`
        """
        # Normal plot
        self.axes.set_title("%s %s" % (self.data[ind].name, self.data[ind].date))

        # x-axis label
        if self.data[ind].coordinate_system['x'] == 'HG':
            xlabel = 'Longitude [%s]' % self.data[ind].units['x']
        else:
            xlabel = 'X-position [%s]' % self.data[ind].units['x']

        # y-axis label
        if self.data[ind].coordinate_system['y'] == 'HG':
            ylabel = 'Latitude [%s]' % self.data[ind].units['y']
        else:
            ylabel = 'Y-position [%s]' % self.data[ind].units['y']

        self.axes.set_xlabel(xlabel)
        self.axes.set_ylabel(ylabel)

    def plot_start_image(self, ax):
        im = self.mapcube[0].plot(annotate=self.annotate, axes=ax,
                                       **self.imshow_kwargs)
        return im<commit_msg>Fixed reference to .norm in visualization.mapcubeanimator<commit_after># -*- coding: utf-8 -*-

__all__ = ['MapCubeAnimator']

from sunpy.visualization import imageanimator

class MapCubeAnimator(imageanimator.BaseFuncAnimator):
    """
    Create an interactive viewer for a MapCube

    The following keyboard shortcuts are defined in the viewer:

    - 'left': previous step on active slider
    - 'right': next step on active slider
    - 'top': change the active slider up one
    - 'bottom': change the active slider down one
    - 'p': play/pause active slider

    Parameters
    ----------
    mapcube: sunpy.map.MapCube
        A MapCube

    annotate: bool
        Annotate the figure with scale and titles

    fig: mpl.figure
        Figure to use

    interval: int
        Animation interval in ms

    colorbar: bool
        Plot colorbar

    Notes
    -----
    Extra keywords are passed to `mapcube[0].plot()` i.e. the `plot()` routine of
    the first map in the cube.
    """
    def __init__(self, mapcube, annotate=True, **kwargs):

        self.mapcube = mapcube
        self.annotate = annotate
        slider_functions = [self.updatefig]
        slider_ranges = [[0,len(mapcube._maps)]]

        imageanimator.BaseFuncAnimator.__init__(self, mapcube._maps, slider_functions,
                                        slider_ranges, **kwargs)

        self._annotate_plot(0)

    def updatefig(self, val, im, slider):
        i = int(val)
        im.set_array(self.data[i].data)
        im.set_cmap(self.mapcube[i].cmap)
        im.set_mpl_color_normalizer(self.mapcube[i].mpl_color_normalizer)
        # Having this line in means the plot will resize for non-homogenous
        # maps. However it also means that if you zoom in on the plot bad
        # things happen.
        # im.set_extent(self.mapcube[i].xrange + self.mapcube[i].yrange)
        if self.annotate:
            self._annotate_plot(i)

    def _annotate_plot(self, ind):
        """
        Annotate the image.

        This may overwrite some stuff in `GenericMap.plot()`
        """
        # Normal plot
        self.axes.set_title("%s %s" % (self.data[ind].name, self.data[ind].date))

        # x-axis label
        if self.data[ind].coordinate_system['x'] == 'HG':
            xlabel = 'Longitude [%s]' % self.data[ind].units['x']
        else:
            xlabel = 'X-position [%s]' % self.data[ind].units['x']

        # y-axis label
        if self.data[ind].coordinate_system['y'] == 'HG':
            ylabel = 'Latitude [%s]' % self.data[ind].units['y']
        else:
            ylabel = 'Y-position [%s]' % self.data[ind].units['y']

        self.axes.set_xlabel(xlabel)
        self.axes.set_ylabel(ylabel)

    def plot_start_image(self, ax):
        im = self.mapcube[0].plot(annotate=self.annotate, axes=ax,
                                       **self.imshow_kwargs)
        return im
<|endoftext|>


<commit_before>#!/usr/bin/python
# -*- coding: utf-8  -*-
"""
image_resizer.py - a script that reduces the size of non-free images

usage:

    python pwb.py oversize_image
"""
#
# (C) w:fa:User:Huji, 2021
#
# Distributed under the terms of the MIT license.
#

import pywikibot
from pywikibot import pagegenerators
import pandas as pd
from persiantools import digits

class OversizeImageBot:

    def __init__(self):
        self.site = pywikibot.Site()
        self.cat = 'رده:محتویات غیر آزاد'
        self.out = 'User:Huji/oversize_images'
        self.detail = 'User:Huji/oversize_images/details'
        self.summary = 'روزآمدسازی آمار'
        self.df = pd.DataFrame(
            columns=['user', 'file', 'ts', 'width', 'height'],
            dtype=str)

    def short_date(self, ts):
        return '%s-%s-%s' % (
            digits.en_to_fa(str(ts.year)),
            digits.en_to_fa(('0' + str(ts.month))[-2:]),
            digits.en_to_fa(('0' + str(ts.day))[-2:])
        )

    def tally(self, title, fileinfo):
        self.df = self.df.append({
            'user': fileinfo.user,
            'file': title,
            'ts': self.short_date(fileinfo.timestamp),
            'width': digits.en_to_fa(str(fileinfo.width)),
            'height': digits.en_to_fa(str(fileinfo.height))
        }, ignore_index=True)

    def aggregate(self):
        tab = self.df[['user', 'file']].groupby('user').count(). \
            sort_values('file', ascending=False)

        wikitab = '{| class=wikitable\n'
        wikitab += '! کاربر !! شمار پرونده\n'

        for idx, row in tab.iterrows():
            wikitab += '|-\n'
            wikitab += '| [[User:' + row.name + '|]]\n'
            wikitab += '| ' + str(row['file']) + '\n'

        wikitab += '|}'

        return wikitab

    def tabulate(self):
        wikitab = '{| class=wikitable\n'
        wikitab += '! کاربر !! پرونده !! تاریخ بارگذاری !! ابعاد\n'

        for idx, row in self.df.sort_values('user').iterrows():
            wikitab += '|-\n'
            wikitab += '| [[User:' + row['user'] + '|]]\n'
            wikitab += '| [[:' + row['file'] + ']]\n'
            wikitab += '| ' + row['ts'] + '\n'
            wikitab += '| %s×%s\n' % (row['width'], row['height'])

        wikitab += '|}'

        return wikitab

    def treat(self, filepage):
        print(filepage)
        title = filepage.title()
        fi = filepage.latest_file_info
        if fi.width * fi.height > 1000000:
            self.tally(title, fi)

    def run(self):
        cat = pywikibot.Category(self.site, self.cat)
        gen = pagegenerators.CategorizedPageGenerator(cat)

        for page in pagegenerators.PreloadingGenerator(gen):
            self.treat(page)

        tab = self.aggregate()
        page = pywikibot.Page(self.site, self.out)
        page.text = tab
        page.save(summary=self.summary, minor=False, botflag=False)

        tab = self.tabulate()
        page = pywikibot.Page(self.site, self.detail)
        page.text = tab
        page.save(summary=self.summary, minor=False, botflag=False)

if __name__ == "__main__":
    bot = OversizeImageBot()
    bot.run()
<commit_msg>Allow sorting oversize image table<commit_after>#!/usr/bin/python
# -*- coding: utf-8  -*-
"""
image_resizer.py - a script that reduces the size of non-free images

usage:

    python pwb.py oversize_image
"""
#
# (C) w:fa:User:Huji, 2021
#
# Distributed under the terms of the MIT license.
#

import pywikibot
from pywikibot import pagegenerators
import pandas as pd
from persiantools import digits

class OversizeImageBot:

    def __init__(self):
        self.site = pywikibot.Site()
        self.cat = 'رده:محتویات غیر آزاد'
        self.out = 'User:Huji/oversize_images'
        self.detail = 'User:Huji/oversize_images/details'
        self.summary = 'روزآمدسازی آمار'
        self.df = pd.DataFrame(
            columns=['user', 'file', 'ts', 'width', 'height'],
            dtype=str)

    def short_date(self, ts):
        return '%s-%s-%s' % (
            digits.en_to_fa(str(ts.year)),
            digits.en_to_fa(('0' + str(ts.month))[-2:]),
            digits.en_to_fa(('0' + str(ts.day))[-2:])
        )

    def tally(self, title, fileinfo):
        self.df = self.df.append({
            'user': fileinfo.user,
            'file': title,
            'ts': self.short_date(fileinfo.timestamp),
            'width': digits.en_to_fa(str(fileinfo.width)),
            'height': digits.en_to_fa(str(fileinfo.height))
        }, ignore_index=True)

    def aggregate(self):
        tab = self.df[['user', 'file']].groupby('user').count(). \
            sort_values('file', ascending=False)

        wikitab = '{| class="wikitable sortable"\n'
        wikitab += '! کاربر !! شمار پرونده\n'

        for idx, row in tab.iterrows():
            wikitab += '|-\n'
            wikitab += '| [[User:' + row.name + '|]]\n'
            wikitab += '| ' + str(row['file']) + '\n'

        wikitab += '|}'

        return wikitab

    def tabulate(self):
        wikitab = '{| class=wikitable\n'
        wikitab += '! کاربر !! پرونده !! تاریخ بارگذاری !! ابعاد\n'

        for idx, row in self.df.sort_values('user').iterrows():
            wikitab += '|-\n'
            wikitab += '| [[User:' + row['user'] + '|]]\n'
            wikitab += '| [[:' + row['file'] + ']]\n'
            wikitab += '| ' + row['ts'] + '\n'
            wikitab += '| %s×%s\n' % (row['width'], row['height'])

        wikitab += '|}'

        return wikitab

    def treat(self, filepage):
        print(filepage)
        title = filepage.title()
        fi = filepage.latest_file_info
        if fi.width * fi.height > 1000000:
            self.tally(title, fi)

    def run(self):
        cat = pywikibot.Category(self.site, self.cat)
        gen = pagegenerators.CategorizedPageGenerator(cat)

        for page in pagegenerators.PreloadingGenerator(gen):
            self.treat(page)

        tab = self.aggregate()
        page = pywikibot.Page(self.site, self.out)
        page.text = tab
        page.save(summary=self.summary, minor=False, botflag=False)

        tab = self.tabulate()
        page = pywikibot.Page(self.site, self.detail)
        page.text = tab
        page.save(summary=self.summary, minor=False, botflag=False)

if __name__ == "__main__":
    bot = OversizeImageBot()
    bot.run()
<|endoftext|>


<commit_before>import os
import shutil
import subprocess
import tempfile

class SCSynthDef(object):

    ### INITIALIZER ###

    def __init__(self, name, code):
        self._name = name
        self._code = code

    ### PUBLIC METHODS ###

    def compile(self):
        directory_path = tempfile.mkdtemp()
        input_ = []
        input_.append('(')
        input_.append('a = SynthDef(')
        input_.append(r'    \{}, {{'.format(self.name))
        for line in self.code.splitlines():
            input_.append('    ' + line)
        input_.append('});')
        input_.append('a.writeDefFile("{}");'.format(directory_path))
        input_.append('0.exit;')
        input_.append(')')
        input_ = '\n'.join(input_)
        sc_file_name = '{}.sc'.format(self.name)
        sc_file_path = os.path.join(directory_path, sc_file_name)
        synthdef_file_name = '{}.scsyndef'.format(self.name)
        synthdef_file_path = os.path.join(directory_path, synthdef_file_name)
        with open(sc_file_path, 'w') as f:
            f.write(input_)
        command = 'sclang {}'.format(sc_file_path)
        subprocess.call(command, shell=True)
        with open(synthdef_file_path, 'rb') as f:
            result = f.read()
        shutil.rmtree(directory_path)
        return bytearray(result)

    ### PUBLIC PROPERTIES ###

    @property
    def code(self):
        return self._code

    @property
    def name(self):
        return self._name
<commit_msg>Adjusted subprocess creation in SCSynthDef.<commit_after>import os
import shutil
import subprocess
import tempfile

class SCSynthDef(object):

    ### INITIALIZER ###

    def __init__(self, name, code):
        self._name = name
        self._code = code

    ### PUBLIC METHODS ###

    def compile(self):
        directory_path = tempfile.mkdtemp()
        input_ = []
        input_.append('(')
        input_.append('a = SynthDef(')
        input_.append(r'    \{}, {{'.format(self.name))
        for line in self.code.splitlines():
            input_.append('    ' + line)
        input_.append('});')
        input_.append('a.writeDefFile("{}");'.format(directory_path))
        input_.append('0.exit;')
        input_.append(')')
        input_ = '\n'.join(input_)
        sc_file_name = '{}.sc'.format(self.name)
        sc_file_path = os.path.join(directory_path, sc_file_name)
        synthdef_file_name = '{}.scsyndef'.format(self.name)
        synthdef_file_path = os.path.join(directory_path, synthdef_file_name)
        with open(sc_file_path, 'w') as f:
            f.write(input_)
        command = ['sclang', sc_file_path]
        subprocess.call(command)
        with open(synthdef_file_path, 'rb') as f:
            result = f.read()
        shutil.rmtree(directory_path)
        return bytearray(result)

    ### PUBLIC PROPERTIES ###

    @property
    def code(self):
        return self._code

    @property
    def name(self):
        return self._name
<|endoftext|>


<commit_before># -*- coding: utf-8 -*-

"""Tests for vumi.components.tagpool."""

import json

from twisted.trial.unittest import TestCase
from twisted.internet.defer import inlineCallbacks

from vumi.components.tagpool import TagpoolManager, TagpoolError
from vumi.tests.utils import PersistenceMixin

class TestTxTagpoolManager(TestCase, PersistenceMixin):

    @inlineCallbacks
    def setUp(self):
        self._persist_setUp()
        self.redis = yield self.get_redis_manager()
        self.tpm = TagpoolManager(self.redis)

    def tearDown(self):
        return self._persist_tearDown()

    def pool_key_generator(self, pool):
        def tkey(x):
            return "tagpools:%s:%s" % (pool, x)
        return tkey

    @inlineCallbacks
    def test_declare_tags(self):
        tag1, tag2 = ("poolA", "tag1"), ("poolA", "tag2")
        yield self.tpm.declare_tags([tag1, tag2])
        self.assertEqual((yield self.tpm.acquire_tag("poolA")), tag1)
        self.assertEqual((yield self.tpm.acquire_tag("poolA")), tag2)
        self.assertEqual((yield self.tpm.acquire_tag("poolA")), None)
        tag3 = ("poolA", "tag3")
        yield self.tpm.declare_tags([tag2, tag3])
        self.assertEqual((yield self.tpm.acquire_tag("poolA")), tag3)

    @inlineCallbacks
    def test_purge_pool(self):
        tag1, tag2 = ("poolA", "tag1"), ("poolA", "tag2")
        yield self.tpm.declare_tags([tag1, tag2])
        yield self.tpm.purge_pool('poolA')
        self.assertEqual((yield self.tpm.acquire_tag('poolA')), None)

    @inlineCallbacks
    def test_purge_inuse_pool(self):
        tag1, tag2 = ("poolA", "tag1"), ("poolA", "tag2")
        yield self.tpm.declare_tags([tag1, tag2])
        self.assertEqual((yield self.tpm.acquire_tag('poolA')), tag1)
        try:
            yield self.tpm.purge_pool('poolA')
        except TagpoolError:
            pass
        else:
            self.fail("Expected TagpoolError to be raised.")

    @inlineCallbacks
    def test_list_pools(self):
        tag1, tag2 = ("poolA", "tag1"), ("poolB", "tag2")
        yield self.tpm.declare_tags([tag1, tag2])
        self.assertEqual((yield self.tpm.list_pools()),
                         set(['poolA', 'poolB']))

    @inlineCallbacks
    def test_acquire_tag(self):
        tkey = self.pool_key_generator("poolA")
        tag1, tag2 = ("poolA", "tag1"), ("poolA", "tag2")
        yield self.tpm.declare_tags([tag1, tag2])
        self.assertEqual((yield self.tpm.acquire_tag("poolA")), tag1)
        self.assertEqual((yield self.tpm.acquire_tag("poolB")), None)
        redis = self.redis
        self.assertEqual((yield redis.lrange(tkey("free:list"), 0, -1)),
                         ["tag2"])
        self.assertEqual((yield redis.smembers(tkey("free:set"))),
                          set(["tag2"]))
        self.assertEqual((yield redis.smembers(tkey("inuse:set"))),
                          set(["tag1"]))

    @inlineCallbacks
    def test_acquire_specific_tag(self):
        tkey = self.pool_key_generator("poolA")
        tags = [("poolA", "tag%d" % i) for i in range(10)]
        tag5 = tags[5]
        yield self.tpm.declare_tags(tags)
        self.assertEqual((yield self.tpm.acquire_specific_tag(tag5)), tag5)
        self.assertEqual((yield self.tpm.acquire_specific_tag(tag5)), None)
        free_local_tags = [t[1] for t in tags]
        free_local_tags.remove("tag5")
        redis = self.redis
        self.assertEqual((yield redis.lrange(tkey("free:list"), 0, -1)),
                         free_local_tags)
        self.assertEqual((yield redis.smembers(tkey("free:set"))),
                         set(free_local_tags))
        self.assertEqual((yield redis.smembers(tkey("inuse:set"))),
                         set(["tag5"]))

    @inlineCallbacks
    def test_release_tag(self):
        tkey = self.pool_key_generator("poolA")
        tag1, tag2, tag3 = [("poolA", "tag%d" % i) for i in (1, 2, 3)]
        yield self.tpm.declare_tags([tag1, tag2, tag3])
        yield self.tpm.acquire_tag("poolA")
        yield self.tpm.acquire_tag("poolA")
        yield self.tpm.release_tag(tag1)
        redis = self.redis
        self.assertEqual((yield redis.lrange(tkey("free:list"), 0, -1)),
                         ["tag3", "tag1"])
        self.assertEqual((yield redis.smembers(tkey("free:set"))),
                         set(["tag1", "tag3"]))
        self.assertEqual((yield redis.smembers(tkey("inuse:set"))),
                         set(["tag2"]))

    @inlineCallbacks
    def test_metadata(self):
        mkey = self.pool_key_generator("poolA")("metadata")
        metadata = {
            "transport_type": "sms",
            "default_msg_fields": {
                "transport_name": "sphex",
                "helper_metadata": {
                    "even_more_nested": "foo",
                },
            },
        }
        yield self.tpm.set_metadata("poolA", metadata)
        self.assertEqual((yield self.tpm.get_metadata("poolA")), metadata)
        tt_json = yield self.redis.hget(mkey, "transport_type")
        transport_type = json.loads(tt_json)
        self.assertEqual(transport_type, "sms")

        short_md = {"foo": "bar"}
        yield self.tpm.set_metadata("poolA", short_md)
        self.assertEqual((yield self.tpm.get_metadata("poolA")), short_md)

class TestTagpoolManager(TestTxTagpoolManager):
    sync_persistence = True
<commit_msg>Clear redis keys at start of tagpool tests to avoid spurious failure when other tests are broken.<commit_after># -*- coding: utf-8 -*-

"""Tests for vumi.components.tagpool."""

import json

from twisted.trial.unittest import TestCase
from twisted.internet.defer import inlineCallbacks

from vumi.components.tagpool import TagpoolManager, TagpoolError
from vumi.tests.utils import PersistenceMixin

class TestTxTagpoolManager(TestCase, PersistenceMixin):

    @inlineCallbacks
    def setUp(self):
        self._persist_setUp()
        self.redis = yield self.get_redis_manager()
        yield self.redis._purge_all()  # Just in case
        self.tpm = TagpoolManager(self.redis)

    def tearDown(self):
        return self._persist_tearDown()

    def pool_key_generator(self, pool):
        def tkey(x):
            return "tagpools:%s:%s" % (pool, x)
        return tkey

    @inlineCallbacks
    def test_declare_tags(self):
        tag1, tag2 = ("poolA", "tag1"), ("poolA", "tag2")
        yield self.tpm.declare_tags([tag1, tag2])
        self.assertEqual((yield self.tpm.acquire_tag("poolA")), tag1)
        self.assertEqual((yield self.tpm.acquire_tag("poolA")), tag2)
        self.assertEqual((yield self.tpm.acquire_tag("poolA")), None)
        tag3 = ("poolA", "tag3")
        yield self.tpm.declare_tags([tag2, tag3])
        self.assertEqual((yield self.tpm.acquire_tag("poolA")), tag3)

    @inlineCallbacks
    def test_purge_pool(self):
        tag1, tag2 = ("poolA", "tag1"), ("poolA", "tag2")
        yield self.tpm.declare_tags([tag1, tag2])
        yield self.tpm.purge_pool('poolA')
        self.assertEqual((yield self.tpm.acquire_tag('poolA')), None)

    @inlineCallbacks
    def test_purge_inuse_pool(self):
        tag1, tag2 = ("poolA", "tag1"), ("poolA", "tag2")
        yield self.tpm.declare_tags([tag1, tag2])
        self.assertEqual((yield self.tpm.acquire_tag('poolA')), tag1)
        try:
            yield self.tpm.purge_pool('poolA')
        except TagpoolError:
            pass
        else:
            self.fail("Expected TagpoolError to be raised.")

    @inlineCallbacks
    def test_list_pools(self):
        tag1, tag2 = ("poolA", "tag1"), ("poolB", "tag2")
        yield self.tpm.declare_tags([tag1, tag2])
        self.assertEqual((yield self.tpm.list_pools()),
                         set(['poolA', 'poolB']))

    @inlineCallbacks
    def test_acquire_tag(self):
        tkey = self.pool_key_generator("poolA")
        tag1, tag2 = ("poolA", "tag1"), ("poolA", "tag2")
        yield self.tpm.declare_tags([tag1, tag2])
        self.assertEqual((yield self.tpm.acquire_tag("poolA")), tag1)
        self.assertEqual((yield self.tpm.acquire_tag("poolB")), None)
        redis = self.redis
        self.assertEqual((yield redis.lrange(tkey("free:list"), 0, -1)),
                         ["tag2"])
        self.assertEqual((yield redis.smembers(tkey("free:set"))),
                          set(["tag2"]))
        self.assertEqual((yield redis.smembers(tkey("inuse:set"))),
                          set(["tag1"]))

    @inlineCallbacks
    def test_acquire_specific_tag(self):
        tkey = self.pool_key_generator("poolA")
        tags = [("poolA", "tag%d" % i) for i in range(10)]
        tag5 = tags[5]
        yield self.tpm.declare_tags(tags)
        self.assertEqual((yield self.tpm.acquire_specific_tag(tag5)), tag5)
        self.assertEqual((yield self.tpm.acquire_specific_tag(tag5)), None)
        free_local_tags = [t[1] for t in tags]
        free_local_tags.remove("tag5")
        redis = self.redis
        self.assertEqual((yield redis.lrange(tkey("free:list"), 0, -1)),
                         free_local_tags)
        self.assertEqual((yield redis.smembers(tkey("free:set"))),
                         set(free_local_tags))
        self.assertEqual((yield redis.smembers(tkey("inuse:set"))),
                         set(["tag5"]))

    @inlineCallbacks
    def test_release_tag(self):
        tkey = self.pool_key_generator("poolA")
        tag1, tag2, tag3 = [("poolA", "tag%d" % i) for i in (1, 2, 3)]
        yield self.tpm.declare_tags([tag1, tag2, tag3])
        yield self.tpm.acquire_tag("poolA")
        yield self.tpm.acquire_tag("poolA")
        yield self.tpm.release_tag(tag1)
        redis = self.redis
        self.assertEqual((yield redis.lrange(tkey("free:list"), 0, -1)),
                         ["tag3", "tag1"])
        self.assertEqual((yield redis.smembers(tkey("free:set"))),
                         set(["tag1", "tag3"]))
        self.assertEqual((yield redis.smembers(tkey("inuse:set"))),
                         set(["tag2"]))

    @inlineCallbacks
    def test_metadata(self):
        mkey = self.pool_key_generator("poolA")("metadata")
        metadata = {
            "transport_type": "sms",
            "default_msg_fields": {
                "transport_name": "sphex",
                "helper_metadata": {
                    "even_more_nested": "foo",
                },
            },
        }
        yield self.tpm.set_metadata("poolA", metadata)
        self.assertEqual((yield self.tpm.get_metadata("poolA")), metadata)
        tt_json = yield self.redis.hget(mkey, "transport_type")
        transport_type = json.loads(tt_json)
        self.assertEqual(transport_type, "sms")

        short_md = {"foo": "bar"}
        yield self.tpm.set_metadata("poolA", short_md)
        self.assertEqual((yield self.tpm.get_metadata("poolA")), short_md)

class TestTagpoolManager(TestTxTagpoolManager):
    sync_persistence = True
<|endoftext|>


<commit_before>"""Utilities to manipulate JSON objects.
"""
#-----------------------------------------------------------------------------
#  Copyright (C) 2010-2011  The IPython Development Team
#
#  Distributed under the terms of the BSD License.  The full license is in
#  the file COPYING.txt, distributed as part of this software.
#-----------------------------------------------------------------------------

#-----------------------------------------------------------------------------
# Imports
#-----------------------------------------------------------------------------
# stdlib
import math
import re
import types
from datetime import datetime

try:
    # base64.encodestring is deprecated in Python 3.x
    from base64 import encodebytes
except ImportError:
    # Python 2.x
    from base64 import encodestring as encodebytes

from IPython.utils import py3compat
from IPython.utils.encoding import DEFAULT_ENCODING
next_attr_name = '__next__' if py3compat.PY3 else 'next'

#-----------------------------------------------------------------------------
# Globals and constants
#-----------------------------------------------------------------------------

# timestamp formats
ISO8601="%Y-%m-%dT%H:%M:%S.%f"
ISO8601_PAT=re.compile(r"^\d{4}-\d{2}-\d{2}T\d{2}:\d{2}:\d{2}\.\d+(Z?[\+\-]\d{2}:?\d{2})?$")

#-----------------------------------------------------------------------------
# Classes and functions
#-----------------------------------------------------------------------------

def rekey(dikt):
    """Rekey a dict that has been forced to use str keys where there should be
    ints by json."""
    for k in dikt.iterkeys():
        if isinstance(k, basestring):
            ik=fk=None
            try:
                ik = int(k)
            except ValueError:
                try:
                    fk = float(k)
                except ValueError:
                    continue
            if ik is not None:
                nk = ik
            else:
                nk = fk
            if nk in dikt:
                raise KeyError("already have key %r"%nk)
            dikt[nk] = dikt.pop(k)
    return dikt

def extract_dates(obj):
    """extract ISO8601 dates from unpacked JSON"""
    if isinstance(obj, dict):
        obj = dict(obj) # don't clobber
        for k,v in obj.iteritems():
            obj[k] = extract_dates(v)
    elif isinstance(obj, (list, tuple)):
        obj = [ extract_dates(o) for o in obj ]
    elif isinstance(obj, basestring):
        if ISO8601_PAT.match(obj):
            # FIXME: add actual timezone support
            notz = obj.split('Z',1)[0]
            obj = datetime.strptime(notz, ISO8601)
    return obj

def squash_dates(obj):
    """squash datetime objects into ISO8601 strings"""
    if isinstance(obj, dict):
        obj = dict(obj) # don't clobber
        for k,v in obj.iteritems():
            obj[k] = squash_dates(v)
    elif isinstance(obj, (list, tuple)):
        obj = [ squash_dates(o) for o in obj ]
    elif isinstance(obj, datetime):
        obj = obj.isoformat()
    return obj

def date_default(obj):
    """default function for packing datetime objects in JSON."""
    if isinstance(obj, datetime):
        return obj.isoformat()
    else:
        raise TypeError("%r is not JSON serializable"%obj)

# constants for identifying png/jpeg data
PNG = b'\x89PNG\r\n\x1a\n'
# front of PNG base64-encoded
PNG64 = b'iVBORw0KG'
JPEG = b'\xff\xd8'
# front of JPEG base64-encoded
JPEG64 = b'/9'

def encode_images(format_dict):
    """b64-encodes images in a displaypub format dict

    Perhaps this should be handled in json_clean itself?

    Parameters
    ----------

    format_dict : dict
        A dictionary of display data keyed by mime-type

    Returns
    -------

    format_dict : dict
        A copy of the same dictionary,
        but binary image data ('image/png' or 'image/jpeg')
        is base64-encoded.

    """
    encoded = format_dict.copy()

    pngdata = format_dict.get('image/png')
    if isinstance(pngdata, bytes):
        # make sure we don't double-encode
        if not pngdata.startswith(PNG64):
            pngdata = encodebytes(pngdata)
        encoded['image/png'] = pngdata.decode('ascii')

    jpegdata = format_dict.get('image/jpeg')
    if isinstance(jpegdata, bytes):
        # make sure we don't double-encode
        if not jpegdata.startswith(JPEG64):
            jpegdata = encodebytes(jpegdata)
        encoded['image/jpeg'] = jpegdata.decode('ascii')

    return encoded

def json_clean(obj):
    """Clean an object to ensure it's safe to encode in JSON.

    Atomic, immutable objects are returned unmodified.  Sets and tuples are
    converted to lists, lists are copied and dicts are also copied.

    Note: dicts whose keys could cause collisions upon encoding (such as a dict
    with both the number 1 and the string '1' as keys) will cause a ValueError
    to be raised.

    Parameters
    ----------
    obj : any python object

    Returns
    -------
    out : object

      A version of the input which will not cause an encoding error when
      encoded as JSON.  Note that this function does not *encode* its inputs,
      it simply sanitizes it so that there will be no encoding errors later.

    Examples
    --------
    >>> json_clean(4)
    4
    >>> json_clean(range(10))
    [0, 1, 2, 3, 4, 5, 6, 7, 8, 9]
    >>> sorted(json_clean(dict(x=1, y=2)).items())
    [('x', 1), ('y', 2)]
    >>> sorted(json_clean(dict(x=1, y=2, z=[1,2,3])).items())
    [('x', 1), ('y', 2), ('z', [1, 2, 3])]
    >>> json_clean(True)
    True
    """
    # types that are 'atomic' and ok in json as-is.  bool doesn't need to be
    # listed explicitly because bools pass as int instances
    atomic_ok = (unicode, int, types.NoneType)

    # containers that we need to convert into lists
    container_to_list = (tuple, set, types.GeneratorType)

    if isinstance(obj, float):
        # cast out-of-range floats to their reprs
        if math.isnan(obj) or math.isinf(obj):
            return repr(obj)
        return obj

    if isinstance(obj, atomic_ok):
        return obj

    if isinstance(obj, bytes):
        return obj.decode(DEFAULT_ENCODING, 'replace')

    if isinstance(obj, container_to_list) or (
        hasattr(obj, '__iter__') and hasattr(obj, next_attr_name)):
        obj = list(obj)

    if isinstance(obj, list):
        return [json_clean(x) for x in obj]

    if isinstance(obj, dict):
        # First, validate that the dict won't lose data in conversion due to
        # key collisions after stringification.  This can happen with keys like
        # True and 'true' or 1 and '1', which collide in JSON.
        nkeys = len(obj)
        nkeys_collapsed = len(set(map(str, obj)))
        if nkeys != nkeys_collapsed:
            raise ValueError('dict can not be safely converted to JSON: '
                             'key collision would lead to dropped values')
        # If all OK, proceed by making the new dict that will be json-safe
        out = {}
        for k,v in obj.iteritems():
            out[str(k)] = json_clean(v)
        return out

    # If we get here, we don't know how to handle the object, so we just get
    # its repr and return that.  This will catch lambdas, open sockets, class
    # objects, and any other complicated contraption that json can't encode
    return repr(obj)
<commit_msg>allow just the trailing Z for timezone<commit_after>"""Utilities to manipulate JSON objects.
"""
#-----------------------------------------------------------------------------
#  Copyright (C) 2010-2011  The IPython Development Team
#
#  Distributed under the terms of the BSD License.  The full license is in
#  the file COPYING.txt, distributed as part of this software.
#-----------------------------------------------------------------------------

#-----------------------------------------------------------------------------
# Imports
#-----------------------------------------------------------------------------
# stdlib
import math
import re
import types
from datetime import datetime

try:
    # base64.encodestring is deprecated in Python 3.x
    from base64 import encodebytes
except ImportError:
    # Python 2.x
    from base64 import encodestring as encodebytes

from IPython.utils import py3compat
from IPython.utils.encoding import DEFAULT_ENCODING
next_attr_name = '__next__' if py3compat.PY3 else 'next'

#-----------------------------------------------------------------------------
# Globals and constants
#-----------------------------------------------------------------------------

# timestamp formats
ISO8601="%Y-%m-%dT%H:%M:%S.%f"
ISO8601_PAT=re.compile(r"^\d{4}-\d{2}-\d{2}T\d{2}:\d{2}:\d{2}\.\d+Z?([\+\-]\d{2}:?\d{2})?$")

#-----------------------------------------------------------------------------
# Classes and functions
#-----------------------------------------------------------------------------

def rekey(dikt):
    """Rekey a dict that has been forced to use str keys where there should be
    ints by json."""
    for k in dikt.iterkeys():
        if isinstance(k, basestring):
            ik=fk=None
            try:
                ik = int(k)
            except ValueError:
                try:
                    fk = float(k)
                except ValueError:
                    continue
            if ik is not None:
                nk = ik
            else:
                nk = fk
            if nk in dikt:
                raise KeyError("already have key %r"%nk)
            dikt[nk] = dikt.pop(k)
    return dikt

def extract_dates(obj):
    """extract ISO8601 dates from unpacked JSON"""
    if isinstance(obj, dict):
        obj = dict(obj) # don't clobber
        for k,v in obj.iteritems():
            obj[k] = extract_dates(v)
    elif isinstance(obj, (list, tuple)):
        obj = [ extract_dates(o) for o in obj ]
    elif isinstance(obj, basestring):
        if ISO8601_PAT.match(obj):
            # FIXME: add actual timezone support
            notz = obj.split('Z',1)[0]
            obj = datetime.strptime(notz, ISO8601)
    return obj

def squash_dates(obj):
    """squash datetime objects into ISO8601 strings"""
    if isinstance(obj, dict):
        obj = dict(obj) # don't clobber
        for k,v in obj.iteritems():
            obj[k] = squash_dates(v)
    elif isinstance(obj, (list, tuple)):
        obj = [ squash_dates(o) for o in obj ]
    elif isinstance(obj, datetime):
        obj = obj.isoformat()
    return obj

def date_default(obj):
    """default function for packing datetime objects in JSON."""
    if isinstance(obj, datetime):
        return obj.isoformat()
    else:
        raise TypeError("%r is not JSON serializable"%obj)

# constants for identifying png/jpeg data
PNG = b'\x89PNG\r\n\x1a\n'
# front of PNG base64-encoded
PNG64 = b'iVBORw0KG'
JPEG = b'\xff\xd8'
# front of JPEG base64-encoded
JPEG64 = b'/9'

def encode_images(format_dict):
    """b64-encodes images in a displaypub format dict

    Perhaps this should be handled in json_clean itself?

    Parameters
    ----------

    format_dict : dict
        A dictionary of display data keyed by mime-type

    Returns
    -------

    format_dict : dict
        A copy of the same dictionary,
        but binary image data ('image/png' or 'image/jpeg')
        is base64-encoded.

    """
    encoded = format_dict.copy()

    pngdata = format_dict.get('image/png')
    if isinstance(pngdata, bytes):
        # make sure we don't double-encode
        if not pngdata.startswith(PNG64):
            pngdata = encodebytes(pngdata)
        encoded['image/png'] = pngdata.decode('ascii')

    jpegdata = format_dict.get('image/jpeg')
    if isinstance(jpegdata, bytes):
        # make sure we don't double-encode
        if not jpegdata.startswith(JPEG64):
            jpegdata = encodebytes(jpegdata)
        encoded['image/jpeg'] = jpegdata.decode('ascii')

    return encoded

def json_clean(obj):
    """Clean an object to ensure it's safe to encode in JSON.

    Atomic, immutable objects are returned unmodified.  Sets and tuples are
    converted to lists, lists are copied and dicts are also copied.

    Note: dicts whose keys could cause collisions upon encoding (such as a dict
    with both the number 1 and the string '1' as keys) will cause a ValueError
    to be raised.

    Parameters
    ----------
    obj : any python object

    Returns
    -------
    out : object

      A version of the input which will not cause an encoding error when
      encoded as JSON.  Note that this function does not *encode* its inputs,
      it simply sanitizes it so that there will be no encoding errors later.

    Examples
    --------
    >>> json_clean(4)
    4
    >>> json_clean(range(10))
    [0, 1, 2, 3, 4, 5, 6, 7, 8, 9]
    >>> sorted(json_clean(dict(x=1, y=2)).items())
    [('x', 1), ('y', 2)]
    >>> sorted(json_clean(dict(x=1, y=2, z=[1,2,3])).items())
    [('x', 1), ('y', 2), ('z', [1, 2, 3])]
    >>> json_clean(True)
    True
    """
    # types that are 'atomic' and ok in json as-is.  bool doesn't need to be
    # listed explicitly because bools pass as int instances
    atomic_ok = (unicode, int, types.NoneType)

    # containers that we need to convert into lists
    container_to_list = (tuple, set, types.GeneratorType)

    if isinstance(obj, float):
        # cast out-of-range floats to their reprs
        if math.isnan(obj) or math.isinf(obj):
            return repr(obj)
        return obj

    if isinstance(obj, atomic_ok):
        return obj

    if isinstance(obj, bytes):
        return obj.decode(DEFAULT_ENCODING, 'replace')

    if isinstance(obj, container_to_list) or (
        hasattr(obj, '__iter__') and hasattr(obj, next_attr_name)):
        obj = list(obj)

    if isinstance(obj, list):
        return [json_clean(x) for x in obj]

    if isinstance(obj, dict):
        # First, validate that the dict won't lose data in conversion due to
        # key collisions after stringification.  This can happen with keys like
        # True and 'true' or 1 and '1', which collide in JSON.
        nkeys = len(obj)
        nkeys_collapsed = len(set(map(str, obj)))
        if nkeys != nkeys_collapsed:
            raise ValueError('dict can not be safely converted to JSON: '
                             'key collision would lead to dropped values')
        # If all OK, proceed by making the new dict that will be json-safe
        out = {}
        for k,v in obj.iteritems():
            out[str(k)] = json_clean(v)
        return out

    # If we get here, we don't know how to handle the object, so we just get
    # its repr and return that.  This will catch lambdas, open sockets, class
    # objects, and any other complicated contraption that json can't encode
    return repr(obj)
<|endoftext|>