<commit_before>"""
    All methods must return media_ids that can be
    passed into e.g. like() or comment() functions.
"""

import time
import random
from . import limits

# filters

def filter_not_liked(media_items, log=False):
    # TODO: convert these print function to logging

    not_liked_medias = []
    for m in media_items:
        if 'pk' in m.keys():
            if 'has_liked' in m.keys():
                if not m['has_liked']:
                    if m['like_count'] <= limits.MAX_LIKES_TO_LIKE:
                        not_liked_medias.append(m['pk'])
    if log:
        print ("  Recieved: %d. Already liked: %d." % (
                            len(media_items),
                            len(media_items) - len(not_liked_medias)
                            )
        )
    return not_liked_medias

def filter_media(media_items, log=False):
    # TODO: should return not liked and not commented medias
    # remove filter_not_liked when implemented
    return filter_not_liked(media_items, log=False)

def filter_users(user_items, log=False):
    # TODO: filter users from blacklist and already subscribed
    pass

# getters

def get_your_medias(self):
    self.getSelfUserFeed()
    return [item["pk"] for item in self.LastJson["items"]]

def get_timeline_medias(self):
    if not self.getTimelineFeed():
        self.logger.info("  Error while getting timeline feed")
        return False
    return filter_media(self.LastJson["items"])

def get_user_medias(self, user_id):
    self.getUserFeed(user_id)
    if self.LastJson["status"] == 'fail':
        self.logger.info("  This is a closed account")
        return False
    return filter_media(self.LastJson["items"])

def get_hashtag_medias(self, hashtag):
    if not self.getHashtagFeed(hashtag):
         self.logger.info("Error while getting hashtag feed")
         return False
    return filter_media(self.LastJson["items"])

def get_geotag_medias(self, geotag):
    # TODO: returns list of medias from geotag
    pass

def get_timeline_users(self):
    # TODO: returns list userids who just posted on your timeline feed
    pass

def get_hashtag_users(self, hashtag):
    users = []
    self.getHashtagFeed(hashtag)
    for i in self.LastJson['items']:
        users.append(i['user']['pk'])
    return users

def get_geotag_users(self, geotag):
    # TODO: returns list userids who just posted on this geotag
    pass

def get_userid_from_username(self, username):
    self.searchUsername(username)
    if "user" in self.LastJson:
        return str(self.LastJson["user"]["pk"])
    return None # Not found

def get_user_followers(self, user_id):
    followers = self.getTotalFollowers(user_id)
    return [item['pk'] for item in followers] if followers else False

def get_user_following(self, user_id):
    following = self.getTotalFollowings(user_id)
    return [item['pk'] for item in following] if following else False

def get_media_likers(self, media_id):
    self.getMediaLikers(media_id)
    return [item['pk'] for item in self.LastJson["users"]]

def get_media_comments(self, media_id):
    # TODO:
    pass

def get_media_commenters(self, media_id):
    self.getMediaComments(media_id)
    if 'comments' not in self.LastJson:
        return []
    return [item["user"]["username"] for item in self.LastJson['comments']]

def get_comment(self):
    if len(self.comments):
        return random.choice(self.comments).strip()
    return "lol"
<commit_msg>fix no media bug<commit_after>"""
    All methods must return media_ids that can be
    passed into e.g. like() or comment() functions.
"""

import time
import random
from . import limits

# filters

def filter_not_liked(media_items, log=False):
    # TODO: convert these print function to logging

    not_liked_medias = []
    for m in media_items:
        if 'pk' in m.keys():
            if 'has_liked' in m.keys():
                if not m['has_liked']:
                    if m['like_count'] <= limits.MAX_LIKES_TO_LIKE:
                        not_liked_medias.append(m['pk'])
    if log:
        print ("  Recieved: %d. Already liked: %d." % (
                            len(media_items),
                            len(media_items) - len(not_liked_medias)
                            )
        )
    return not_liked_medias

def filter_media(media_items, log=False):
    # TODO: should return not liked and not commented medias
    # remove filter_not_liked when implemented
    return filter_not_liked(media_items, log=False)

def filter_users(user_items, log=False):
    # TODO: filter users from blacklist and already subscribed
    pass

# getters

def get_your_medias(self):
    self.getSelfUserFeed()
    return [item["pk"] for item in self.LastJson["items"]]

def get_timeline_medias(self):
    if not self.getTimelineFeed():
        self.logger.info("  Error while getting timeline feed")
        return False
    return filter_media(self.LastJson["items"])

def get_user_medias(self, user_id):
    self.getUserFeed(user_id)
    if self.LastJson["status"] == 'fail':
        self.logger.info("  This is a closed account")
        return False
    return filter_media(self.LastJson["items"])

def get_hashtag_medias(self, hashtag):
    if not self.getHashtagFeed(hashtag):
         self.logger.info("Error while getting hashtag feed")
         return False
    return filter_media(self.LastJson["items"])

def get_geotag_medias(self, geotag):
    # TODO: returns list of medias from geotag
    pass

def get_timeline_users(self):
    # TODO: returns list userids who just posted on your timeline feed
    pass

def get_hashtag_users(self, hashtag):
    users = []
    self.getHashtagFeed(hashtag)
    for i in self.LastJson['items']:
        users.append(i['user']['pk'])
    return users

def get_geotag_users(self, geotag):
    # TODO: returns list userids who just posted on this geotag
    pass

def get_userid_from_username(self, username):
    self.searchUsername(username)
    if "user" in self.LastJson:
        return str(self.LastJson["user"]["pk"])
    return None # Not found

def get_user_followers(self, user_id):
    followers = self.getTotalFollowers(user_id)
    return [item['pk'] for item in followers] if followers else False

def get_user_following(self, user_id):
    following = self.getTotalFollowings(user_id)
    return [item['pk'] for item in following] if following else False

def get_media_likers(self, media_id):
    self.getMediaLikers(media_id)
    if "users" not in self.LastJson:
        self.logger.info("Media with %s not found." % media_id)
        return False
    return [item['pk'] for item in self.LastJson["users"]]

def get_media_comments(self, media_id):
    # TODO:
    pass

def get_media_commenters(self, media_id):
    self.getMediaComments(media_id)
    if 'comments' not in self.LastJson:
        return []
    return [item["user"]["username"] for item in self.LastJson['comments']]

def get_comment(self):
    if len(self.comments):
        return random.choice(self.comments).strip()
    return "lol"
<|endoftext|>


<commit_before>#!/usr/bin/env python3
# Copyright 2017 The Kubernetes Authors.
#
# 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.

"""Receive push events for new builds and upload rows to BigQuery."""

import argparse
import json
import os
import pprint
import socket
import sys
import traceback
import time

import multiprocessing.pool

try:
    from google.cloud import bigquery
    from google.cloud import pubsub_v1
    import google.cloud.exceptions
except ImportError:
    print('WARNING: unable to load google cloud (test environment?)')
    traceback.print_exc()

import model
import make_db
import make_json

MAX_ROW_UPLOAD = 25

def process_changes(results):
    """Split GCS change events into trivial ack_ids and builds to further process."""
    ack_ids = []  # pubsub rec_message ids to acknowledge
    todo = []  # (id, job, build) of builds to grab
    # process results, find finished builds to process
    for rec_message in results:
        if rec_message.message.attributes['eventType'] != 'OBJECT_FINALIZE':
            ack_ids.append(rec_message.ack_id)
            continue
        obj = rec_message.message.attributes['objectId']
        if not obj.endswith('/finished.json'):
            ack_ids.append(rec_message.ack_id)
            continue
        job, build = obj[:-len('/finished.json')].rsplit('/', 1)
        job = 'gs://%s/%s' % (rec_message.message.attributes['bucketId'], job)
        todo.append((rec_message.ack_id, job, build))
    return ack_ids, todo

def get_started_finished(gcs_client, db, todo):
    """Download started/finished.json from build dirs in todo."""
    ack_ids = []
    build_dirs = []
    pool = multiprocessing.pool.ThreadPool(16)
    try:
        for ack_id, (build_dir, started, finished) in pool.imap_unordered(
                lambda ack_id_job_build: (ack_id_job_build[0], gcs_client.get_started_finished(
                    ack_id_job_build[1], ack_id_job_build[2])),
                todo):
            if finished:
                if not db.insert_build(build_dir, started, finished):
                    print('build dir already present in db: ', build_dir)
                start = time.localtime(started.get('timestamp', 0) if started else 0)
                print((build_dir, bool(started), bool(finished),
                       time.strftime('%F %T %Z', start),
                       finished and finished.get('result')))
                build_dirs.append(build_dir)
                ack_ids.append(ack_id)
            else:
                print('finished.json missing?', build_dir, started, finished)
    finally:
        pool.close()
    db.commit()
    return ack_ids, build_dirs

def retry(func, *args, **kwargs):
    """Run a function with arguments, retrying on server errors. """
    # pylint: disable=no-member
    for attempt in range(20):
        try:
            return func(*args, **kwargs)
        except (socket.error, google.cloud.exceptions.ServerError):
            # retry with exponential backoff
            traceback.print_exc()
            time.sleep(1.4 ** attempt)
    return func(*args, **kwargs)  # one last attempt

def insert_data(bq_client, table, rows_iter):
    """Upload rows from rows_iter into bigquery table table.

    rows_iter should return a series of (row_id, row dictionary) tuples.
    The row dictionary must match the table's schema.

    Args:
        bq_client: Client connection to BigQuery
        table: bigquery.Table object that points to a specific table
        rows_iter: row_id, dict representing a make_json.Build
    Returns the row_ids that were inserted.
    """
    def divide_chunks(l, bin_size=MAX_ROW_UPLOAD):
        # break up rows to not hit data limits
        for i in range(0, len(l), bin_size):
            yield l[i:i + bin_size]

    emitted, rows = [], []

    for row_id, build in rows_iter:
        emitted.append(row_id)
        rows.append(build)

    if not rows:  # nothing to do
        return []

    for chunk in divide_chunks(rows):
        # Insert rows with row_ids into table, retrying as necessary.
        errors = retry(bq_client.insert_rows, table, chunk, skip_invalid_rows=True)
        if not errors:
            print('Loaded {} builds into {}'.format(len(chunk), table.full_table_id))
        else:
            print('Errors:')
            pprint.pprint(errors)
            pprint.pprint(table.schema)

    return emitted

def main(
        db,
        subscriber,
        subscription_path,
        bq_client,
        tables,
        client_class=make_db.GCSClient,
        stop=None,
    ):
    # pylint: disable=too-many-locals
    gcs_client = client_class('', {})
    if stop is None:
        stop = lambda: False

    results = [0] * 1000  # don't sleep on first loop
    while not stop():
        print()
        if len(results) < 10 and client_class is make_db.GCSClient:
            time.sleep(5)  # slow down!

        print('====', time.strftime("%F %T %Z"), '=' * 40)

        results = retry(subscriber.pull, subscription=subscription_path, max_messages=1000)
        results = list(results.received_messages)
        start = time.time()
        while time.time() < start + 7:
            results_more = list(subscriber.pull(
                subscription=subscription_path,
                max_messages=1000,
                return_immediately=True).received_messages)
            if not results_more:
                break
            results.extend(results_more)

        print('PULLED', len(results))

        ack_ids, todo = process_changes(results)

        if ack_ids:
            print('ACK irrelevant', len(ack_ids))
            for n in range(0, len(ack_ids), 1000):
                retry(
                    subscriber.acknowledge,
                    subscription=subscription_path,
                    ack_ids=ack_ids[n: n + 1000])

        if todo:
            print('EXTEND-ACK ', len(todo))
            # give 3 minutes to grab build details
            retry(
                subscriber.modify_ack_deadline,
                subscription=subscription_path,
                ack_ids=[i for i, _j, _b in todo],
                ack_deadline_seconds=60*3)

        ack_ids, build_dirs = get_started_finished(gcs_client, db, todo)

        # notify pubsub queue that we've handled the finished.json messages
        if ack_ids:
            print('ACK "finished.json"', len(ack_ids))
            retry(subscriber.acknowledge, subscription=subscription_path, ack_ids=ack_ids)

        # grab junit files for new builds
        make_db.download_junit(db, 16, client_class)

        # stream new rows to tables
        if build_dirs and tables:
            for table, incremental_table in tables.values():
                builds = db.get_builds_from_paths(build_dirs, incremental_table)
                emitted = insert_data(bq_client, table, make_json.make_rows(db, builds))
                db.insert_emitted(emitted, incremental_table)

def load_sub(poll):
    """Return the PubSub subscription specified by the /-separated input.

    Args:
        poll: Follow GCS changes from project/topic/subscription
              Ex: kubernetes-jenkins/gcs-changes/kettle

    Return:
        Subscribed client
    """
    subscriber = pubsub_v1.SubscriberClient()
    project_id, _, sub = poll.split('/')
    subscription_path = f'projects/{project_id}/subscriptions/{sub}'
    return subscriber, subscription_path

def load_schema(schemafield):
    """Construct the expected BigQuery schema from files on disk.

    Only used for new tables."""
    basedir = os.path.dirname(__file__)
    schema_json = json.load(open(os.path.join(basedir, 'schema.json')))
    def make_field(spec):
        spec['field_type'] = spec.pop('type')
        if 'fields' in spec:
            spec['fields'] = [make_field(f) for f in spec['fields']]
        return schemafield(**spec)
    return [make_field(s) for s in schema_json]

def load_tables(dataset, tablespecs):
    """Construct a dictionary of BigQuery tables given the input tablespec.

    Args:
        dataset: bigquery.Dataset
        tablespecs: list of strings of "NAME:DAYS", e.g. ["day:1"]
    Returns:
        client, {name: (bigquery.Table, incremental table name)}
    """
    project, dataset_name = dataset.split(':')
    bq_client = bigquery.Client(project)

    tables = {}
    for spec in tablespecs:
        table_name, days = spec.split(':')
        table_ref = f'{project}.{dataset_name}.{table_name}'
        try:
            table = bq_client.get_table(table_ref) # pylint: disable=no-member
        except google.cloud.exceptions.NotFound:
            table = bq_client.create_table(table_ref) # pylint: disable=no-member
            table.schema = load_schema(bigquery.schema.SchemaField)
        tables[table_name] = (table, make_json.get_table(float(days)))
    return bq_client, tables

class StopWhen:
    """A simple object that returns True once when the given hour begins."""
    def __init__(self, target, clock=lambda: time.localtime().tm_hour):
        self.clock = clock
        self.last = self.clock()
        self.target = target

    def __call__(self):
        if os.path.exists('stop'):
            return True
        now = self.clock()
        last = self.last
        self.last = now
        return now != last and now == self.target

def get_options(argv):
    """Process command line arguments."""
    parser = argparse.ArgumentParser()
    parser.add_argument(
        '--poll',
        required=True,
        help='Follow GCS changes from project/topic/subscription',
    )
    parser.add_argument(
        '--dataset',
        help='BigQuery dataset (e.g. k8s-gubernator:build)'
    )
    parser.add_argument(
        '--tables',
        nargs='+',
        default=[],
        help='Upload rows to table:days [e.g. --tables day:1 week:7 all:0]',
    )
    parser.add_argument(
        '--stop_at',
        type=int,
        help='Terminate when this hour (0-23) rolls around (in local time).'
    )
    return parser.parse_args(argv)

if __name__ == '__main__':
    OPTIONS = get_options(sys.argv[1:])
    main(model.Database(),
         *load_sub(OPTIONS.poll),
         *load_tables(OPTIONS.dataset, OPTIONS.tables),
         stop=StopWhen(OPTIONS.stop_at))
<commit_msg>Decrease Max chunk size and catch Over Size error<commit_after>#!/usr/bin/env python3
# Copyright 2017 The Kubernetes Authors.
#
# 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.

"""Receive push events for new builds and upload rows to BigQuery."""

import argparse
import json
import os
import pprint
import socket
import sys
import traceback
import time

import multiprocessing.pool

try:
    from google.api_core import exceptions as api_exceptions
    from google.cloud import bigquery
    from google.cloud import pubsub_v1
    import google.cloud.exceptions
except ImportError:
    print('WARNING: unable to load google cloud (test environment?)')
    traceback.print_exc()

import model
import make_db
import make_json

MAX_ROW_UPLOAD = 10 # See https://github.com/googleapis/google-cloud-go/issues/2855

def process_changes(results):
    """Split GCS change events into trivial ack_ids and builds to further process."""
    ack_ids = []  # pubsub rec_message ids to acknowledge
    todo = []  # (id, job, build) of builds to grab
    # process results, find finished builds to process
    for rec_message in results:
        if rec_message.message.attributes['eventType'] != 'OBJECT_FINALIZE':
            ack_ids.append(rec_message.ack_id)
            continue
        obj = rec_message.message.attributes['objectId']
        if not obj.endswith('/finished.json'):
            ack_ids.append(rec_message.ack_id)
            continue
        job, build = obj[:-len('/finished.json')].rsplit('/', 1)
        job = 'gs://%s/%s' % (rec_message.message.attributes['bucketId'], job)
        todo.append((rec_message.ack_id, job, build))
    return ack_ids, todo

def get_started_finished(gcs_client, db, todo):
    """Download started/finished.json from build dirs in todo."""
    ack_ids = []
    build_dirs = []
    pool = multiprocessing.pool.ThreadPool(16)
    try:
        for ack_id, (build_dir, started, finished) in pool.imap_unordered(
                lambda ack_id_job_build: (ack_id_job_build[0], gcs_client.get_started_finished(
                    ack_id_job_build[1], ack_id_job_build[2])),
                todo):
            if finished:
                if not db.insert_build(build_dir, started, finished):
                    print('build dir already present in db: ', build_dir)
                start = time.localtime(started.get('timestamp', 0) if started else 0)
                print((build_dir, bool(started), bool(finished),
                       time.strftime('%F %T %Z', start),
                       finished and finished.get('result')))
                build_dirs.append(build_dir)
                ack_ids.append(ack_id)
            else:
                print('finished.json missing?', build_dir, started, finished)
    finally:
        pool.close()
    db.commit()
    return ack_ids, build_dirs

def retry(func, *args, **kwargs):
    """Run a function with arguments, retrying on server errors. """
    # pylint: disable=no-member
    for attempt in range(20):
        try:
            return func(*args, **kwargs)
        except (socket.error, google.cloud.exceptions.ServerError, api_exceptions.BadRequest):
            # retry with exponential backoff
            traceback.print_exc()
            time.sleep(1.4 ** attempt)
    return func(*args, **kwargs)  # one last attempt

def insert_data(bq_client, table, rows_iter):
    """Upload rows from rows_iter into bigquery table table.

    rows_iter should return a series of (row_id, row dictionary) tuples.
    The row dictionary must match the table's schema.

    Args:
        bq_client: Client connection to BigQuery
        table: bigquery.Table object that points to a specific table
        rows_iter: row_id, dict representing a make_json.Build
    Returns the row_ids that were inserted.
    """
    def divide_chunks(l, bin_size=MAX_ROW_UPLOAD):
        # break up rows to not hit data limits
        for i in range(0, len(l), bin_size):
            yield l[i:i + bin_size]

    emitted, rows = [], []

    for row_id, build in rows_iter:
        emitted.append(row_id)
        rows.append(build)

    if not rows:  # nothing to do
        return []

    for chunk in divide_chunks(rows):
        # Insert rows with row_ids into table, retrying as necessary.
        errors = retry(bq_client.insert_rows, table, chunk, skip_invalid_rows=True)
        if not errors:
            print(f'Loaded {len(chunk)} builds into {table.full_table_id}')
        else:
            print(f'Errors on Chunk: {chunk}')
            pprint.pprint(errors)
            pprint.pprint(table.schema)

    return emitted

def main(
        db,
        subscriber,
        subscription_path,
        bq_client,
        tables,
        client_class=make_db.GCSClient,
        stop=None,
    ):
    # pylint: disable=too-many-locals
    gcs_client = client_class('', {})
    if stop is None:
        stop = lambda: False

    results = [0] * 1000  # don't sleep on first loop
    while not stop():
        print()
        if len(results) < 10 and client_class is make_db.GCSClient:
            time.sleep(5)  # slow down!

        print('====', time.strftime("%F %T %Z"), '=' * 40)

        results = retry(subscriber.pull, subscription=subscription_path, max_messages=1000)
        results = list(results.received_messages)
        start = time.time()
        while time.time() < start + 7:
            results_more = list(subscriber.pull(
                subscription=subscription_path,
                max_messages=1000,
                return_immediately=True).received_messages)
            if not results_more:
                break
            results.extend(results_more)

        print('PULLED', len(results))

        ack_ids, todo = process_changes(results)

        if ack_ids:
            print('ACK irrelevant', len(ack_ids))
            for n in range(0, len(ack_ids), 1000):
                retry(
                    subscriber.acknowledge,
                    subscription=subscription_path,
                    ack_ids=ack_ids[n: n + 1000])

        if todo:
            print('EXTEND-ACK ', len(todo))
            # give 3 minutes to grab build details
            retry(
                subscriber.modify_ack_deadline,
                subscription=subscription_path,
                ack_ids=[i for i, _j, _b in todo],
                ack_deadline_seconds=60*3)

        ack_ids, build_dirs = get_started_finished(gcs_client, db, todo)

        # notify pubsub queue that we've handled the finished.json messages
        if ack_ids:
            print('ACK "finished.json"', len(ack_ids))
            retry(subscriber.acknowledge, subscription=subscription_path, ack_ids=ack_ids)

        # grab junit files for new builds
        make_db.download_junit(db, 16, client_class)

        # stream new rows to tables
        if build_dirs and tables:
            for table, incremental_table in tables.values():
                builds = db.get_builds_from_paths(build_dirs, incremental_table)
                emitted = insert_data(bq_client, table, make_json.make_rows(db, builds))
                db.insert_emitted(emitted, incremental_table)

def load_sub(poll):
    """Return the PubSub subscription specified by the /-separated input.

    Args:
        poll: Follow GCS changes from project/topic/subscription
              Ex: kubernetes-jenkins/gcs-changes/kettle

    Return:
        Subscribed client
    """
    subscriber = pubsub_v1.SubscriberClient()
    project_id, _, sub = poll.split('/')
    subscription_path = f'projects/{project_id}/subscriptions/{sub}'
    return subscriber, subscription_path

def load_schema(schemafield):
    """Construct the expected BigQuery schema from files on disk.

    Only used for new tables."""
    basedir = os.path.dirname(__file__)
    schema_json = json.load(open(os.path.join(basedir, 'schema.json')))
    def make_field(spec):
        spec['field_type'] = spec.pop('type')
        if 'fields' in spec:
            spec['fields'] = [make_field(f) for f in spec['fields']]
        return schemafield(**spec)
    return [make_field(s) for s in schema_json]

def load_tables(dataset, tablespecs):
    """Construct a dictionary of BigQuery tables given the input tablespec.

    Args:
        dataset: bigquery.Dataset
        tablespecs: list of strings of "NAME:DAYS", e.g. ["day:1"]
    Returns:
        client, {name: (bigquery.Table, incremental table name)}
    """
    project, dataset_name = dataset.split(':')
    bq_client = bigquery.Client(project)

    tables = {}
    for spec in tablespecs:
        table_name, days = spec.split(':')
        table_ref = f'{project}.{dataset_name}.{table_name}'
        try:
            table = bq_client.get_table(table_ref) # pylint: disable=no-member
        except google.cloud.exceptions.NotFound:
            table = bq_client.create_table(table_ref) # pylint: disable=no-member
            table.schema = load_schema(bigquery.schema.SchemaField)
        tables[table_name] = (table, make_json.get_table(float(days)))
    return bq_client, tables

class StopWhen:
    """A simple object that returns True once when the given hour begins."""
    def __init__(self, target, clock=lambda: time.localtime().tm_hour):
        self.clock = clock
        self.last = self.clock()
        self.target = target

    def __call__(self):
        if os.path.exists('stop'):
            return True
        now = self.clock()
        last = self.last
        self.last = now
        return now != last and now == self.target

def get_options(argv):
    """Process command line arguments."""
    parser = argparse.ArgumentParser()
    parser.add_argument(
        '--poll',
        required=True,
        help='Follow GCS changes from project/topic/subscription',
    )
    parser.add_argument(
        '--dataset',
        help='BigQuery dataset (e.g. k8s-gubernator:build)'
    )
    parser.add_argument(
        '--tables',
        nargs='+',
        default=[],
        help='Upload rows to table:days [e.g. --tables day:1 week:7 all:0]',
    )
    parser.add_argument(
        '--stop_at',
        type=int,
        help='Terminate when this hour (0-23) rolls around (in local time).'
    )
    return parser.parse_args(argv)

if __name__ == '__main__':
    OPTIONS = get_options(sys.argv[1:])
    main(model.Database(),
         *load_sub(OPTIONS.poll),
         *load_tables(OPTIONS.dataset, OPTIONS.tables),
         stop=StopWhen(OPTIONS.stop_at))
<|endoftext|>


<commit_before>import logging
import os

import pandas as pd
import synapseclient

from .example_filetype_format import FileTypeFormat
from . import process_functions

logger = logging.getLogger(__name__)

def validateSymbol(gene, bedDf, returnMappedDf=True):
    '''
    Validate gene symbol

    Args:
        gene: Gene name
        bedDf: Bed pandas dataframe
        returnMappedDf: Return a mapped gene. Defaults to True

    Returns:
        gene name or boolean for whether a gene is valid
    '''
    valid = False
    if sum(bedDf['Hugo_Symbol'] == gene) > 0:
        valid = True
    elif sum(bedDf['ID'] == gene) > 0:
        mismatch = bedDf[bedDf['ID'] == gene]
        mismatch.drop_duplicates(inplace=True)
        logger.info("{} will be remapped to {}".format(
            gene, mismatch['Hugo_Symbol'].values[0]))
        gene = mismatch['Hugo_Symbol'].values[0]
    else:
        logger.warning(
            "{} cannot be remapped and will not be released. The symbol "
            "must exist in your seq assay ids (bed files) and must be "
            "mappable to a gene.".format(gene))
        gene = pd.np.nan
    if returnMappedDf:
        return(gene)
    else:
        return(valid)

def makeCNARow(row, symbols):
    '''
    Make CNA Row (Deprecated function)

    CNA values are no longer stored in the database

    Args:
        row: one row in the CNA file
        symbols:  list of Gene symbols
    '''
    totalrow = "{symbols}\n{values}".format(
        symbols=",".join(symbols),
        values=",".join(row.astype(str)))
    totalrow = totalrow.replace(".0", "")
    return(totalrow)

def mergeCNAvalues(x):
    x.dropna(inplace=True)
    uniqueValues = set(x.unique())
    if len(uniqueValues) == 1:
        returnVal = x.tolist()[0]
    elif len(uniqueValues) <= 2:
        uniqueValues.discard(0)
        if len(uniqueValues) == 1:
            returnVal = list(uniqueValues)[0]
        else:
            returnVal = float('nan')
    else:
        returnVal = float('nan')
    return(returnVal)

def checkIfOneZero(x):
    assert len(set(x.tolist())) == 1, "Can only be one unique value"

class cna(FileTypeFormat):

    _fileType = "cna"

    _process_kwargs = ["newPath", 'databaseToSynIdMappingDf']

    _validation_kwargs = ['nosymbol_check']

    # VALIDATE FILENAME
    def _validateFilename(self, filePath):
        assert os.path.basename(filePath[0]) == \
            "data_CNA_{}.txt".format(self.center)

    def _process(self, cnaDf):
        cnaDf.rename(columns={
            cnaDf.columns[0]: cnaDf.columns[0].upper()}, inplace=True)
        cnaDf.rename(columns={
            "HUGO_SYMBOL": "Hugo_Symbol"}, inplace=True)

        index = [i for i, col in enumerate(cnaDf.columns)
                 if col.upper() == "ENTREZ_GENE_ID"]
        if len(index) > 0:
            del cnaDf[cnaDf.columns[index][0]]
        bedSynId = process_functions.getDatabaseSynId(
            self.syn, "bed", test=self.testing)
        bed = self.syn.tableQuery(
            "select Hugo_Symbol, ID from {} where CENTER = '{}'" .format(
                bedSynId, self.center))
        bedDf = bed.asDataFrame()
        cnaDf['Hugo_Symbol'] = \
            cnaDf['Hugo_Symbol'].apply(lambda x: validateSymbol(x, bedDf))
        order = cnaDf.columns
        cnaDf = cnaDf[~cnaDf['Hugo_Symbol'].isnull()]
        # cnaDf = cnaDf.applymap(str)
        duplicatedGenes = pd.DataFrame()
        duplicated_symbols = cnaDf['Hugo_Symbol'][
            cnaDf['Hugo_Symbol'].duplicated()].unique()
        for i in duplicated_symbols:
            dups = cnaDf[cnaDf['Hugo_Symbol'] == i]
            newVal = dups[dups.columns[dups.columns != "Hugo_Symbol"]].apply(
                mergeCNAvalues)
            temp = pd.DataFrame(newVal).transpose()
            temp['Hugo_Symbol'] = i
            duplicatedGenes = duplicatedGenes.append(temp, sort=False)
        cnaDf.drop_duplicates('Hugo_Symbol', keep=False, inplace=True)
        cnaDf = cnaDf.append(duplicatedGenes, sort=False)
        cnaDf = cnaDf[order]
        cnaDf.columns = [
            process_functions.checkGenieId(i, self.center)
            if i != "Hugo_Symbol" else i for i in cnaDf.columns]

        return(cnaDf)

    def process_steps(self, cnaDf, newPath, databaseToSynIdMappingDf):

        newCNA = self._process(cnaDf, test=self.testing)

        centerMafSynId = databaseToSynIdMappingDf.Id[
            databaseToSynIdMappingDf['Database'] == "centerMaf"][0]
        if not newCNA.empty:
            cnaText = process_functions.removePandasDfFloat(newCNA)
            # Replace blank with NA's
            cnaText = cnaText.replace(
                "\t\t", "\tNA\t").replace(
                "\t\t", "\tNA\t").replace(
                '\t\n', "\tNA\n")
            with open(newPath, "w") as cnaFile:
                cnaFile.write(cnaText)
            self.syn.store(synapseclient.File(newPath, parent=centerMafSynId))
        return(newPath)

    def _validate(self, cnvDF, nosymbol_check):
        total_error = ""
        warning = ""
        cnvDF.columns = [col.upper() for col in cnvDF.columns]

        if cnvDF.columns[0] != "HUGO_SYMBOL":
            total_error += "Your cnv file's first column must be Hugo_Symbol\n"
        haveColumn = process_functions.checkColExist(cnvDF, "HUGO_SYMBOL")
        if haveColumn:
            keepSymbols = cnvDF["HUGO_SYMBOL"]
            cnvDF.drop("HUGO_SYMBOL", axis=1, inplace=True)

        # if sum(cnvDF.apply(lambda x: sum(x.isnull()))) > 0:
        #   total_error += "Your cnv file must not have any empty values\n"

        if process_functions.checkColExist(cnvDF, "ENTREZ_GENE_ID"):
            del cnvDF['ENTREZ_GENE_ID']

        # cnvDF = cnvDF.fillna('')
        allowed_values = ['-2.0', '-2', '-1.5', '-1.0', '-1',
                          '0.0', '0', '0.5', '1.0', '1', '1.5',
                          '2', '2.0', 'nan']
        if not all(cnvDF.applymap(lambda x: str(x) in allowed_values).all()):
            total_error += (
                "All values must be NA/blank, -2, -1.5, -1, -0.5, "
                "0, 0.5, 1, 1.5, or 2.\n")
        else:
            cnvDF['HUGO_SYMBOL'] = keepSymbols
            if haveColumn and not nosymbol_check:
                bedSynId = process_functions.getDatabaseSynId(
                    self.syn, "bed", test=self.testing)
                bed = self.syn.tableQuery(
                    "select Hugo_Symbol, ID from {} where "
                    "CENTER = '{}'".format(bedSynId, self.center))
                bedDf = bed.asDataFrame()
                cnvDF['remapped'] = cnvDF['HUGO_SYMBOL'].apply(
                    lambda x: validateSymbol(x, bedDf))
                cnvDF = cnvDF[~cnvDF['remapped'].isnull()]

                # Do not allow any duplicated genes after symbols
                # have been remapped
                if sum(cnvDF['remapped'].duplicated()) > 0:
                    duplicated = cnvDF['remapped'].duplicated(keep=False)
                    total_error += (
                        "Your CNA file has duplicated Hugo_Symbols "
                        "(After remapping of genes): {} -> {}.\n".format(
                            ",".join(cnvDF['HUGO_SYMBOL'][duplicated]),
                            ",".join(cnvDF['remapped'][duplicated])))
        return(total_error, warning)
<commit_msg>Remove test parameter<commit_after>import logging
import os

import pandas as pd
import synapseclient

from .example_filetype_format import FileTypeFormat
from . import process_functions

logger = logging.getLogger(__name__)

def validateSymbol(gene, bedDf, returnMappedDf=True):
    '''
    Validate gene symbol

    Args:
        gene: Gene name
        bedDf: Bed pandas dataframe
        returnMappedDf: Return a mapped gene. Defaults to True

    Returns:
        gene name or boolean for whether a gene is valid
    '''
    valid = False
    if sum(bedDf['Hugo_Symbol'] == gene) > 0:
        valid = True
    elif sum(bedDf['ID'] == gene) > 0:
        mismatch = bedDf[bedDf['ID'] == gene]
        mismatch.drop_duplicates(inplace=True)
        logger.info("{} will be remapped to {}".format(
            gene, mismatch['Hugo_Symbol'].values[0]))
        gene = mismatch['Hugo_Symbol'].values[0]
    else:
        logger.warning(
            "{} cannot be remapped and will not be released. The symbol "
            "must exist in your seq assay ids (bed files) and must be "
            "mappable to a gene.".format(gene))
        gene = pd.np.nan
    if returnMappedDf:
        return(gene)
    else:
        return(valid)

def makeCNARow(row, symbols):
    '''
    Make CNA Row (Deprecated function)

    CNA values are no longer stored in the database

    Args:
        row: one row in the CNA file
        symbols:  list of Gene symbols
    '''
    totalrow = "{symbols}\n{values}".format(
        symbols=",".join(symbols),
        values=",".join(row.astype(str)))
    totalrow = totalrow.replace(".0", "")
    return(totalrow)

def mergeCNAvalues(x):
    x.dropna(inplace=True)
    uniqueValues = set(x.unique())
    if len(uniqueValues) == 1:
        returnVal = x.tolist()[0]
    elif len(uniqueValues) <= 2:
        uniqueValues.discard(0)
        if len(uniqueValues) == 1:
            returnVal = list(uniqueValues)[0]
        else:
            returnVal = float('nan')
    else:
        returnVal = float('nan')
    return(returnVal)

def checkIfOneZero(x):
    assert len(set(x.tolist())) == 1, "Can only be one unique value"

class cna(FileTypeFormat):

    _fileType = "cna"

    _process_kwargs = ["newPath", 'databaseToSynIdMappingDf']

    _validation_kwargs = ['nosymbol_check']

    # VALIDATE FILENAME
    def _validateFilename(self, filePath):
        assert os.path.basename(filePath[0]) == \
            "data_CNA_{}.txt".format(self.center)

    def _process(self, cnaDf):
        cnaDf.rename(columns={
            cnaDf.columns[0]: cnaDf.columns[0].upper()}, inplace=True)
        cnaDf.rename(columns={
            "HUGO_SYMBOL": "Hugo_Symbol"}, inplace=True)

        index = [i for i, col in enumerate(cnaDf.columns)
                 if col.upper() == "ENTREZ_GENE_ID"]
        if len(index) > 0:
            del cnaDf[cnaDf.columns[index][0]]
        bedSynId = process_functions.getDatabaseSynId(
            self.syn, "bed", test=self.testing)
        bed = self.syn.tableQuery(
            "select Hugo_Symbol, ID from {} where CENTER = '{}'" .format(
                bedSynId, self.center))
        bedDf = bed.asDataFrame()
        cnaDf['Hugo_Symbol'] = \
            cnaDf['Hugo_Symbol'].apply(lambda x: validateSymbol(x, bedDf))
        order = cnaDf.columns
        cnaDf = cnaDf[~cnaDf['Hugo_Symbol'].isnull()]
        # cnaDf = cnaDf.applymap(str)
        duplicatedGenes = pd.DataFrame()
        duplicated_symbols = cnaDf['Hugo_Symbol'][
            cnaDf['Hugo_Symbol'].duplicated()].unique()
        for i in duplicated_symbols:
            dups = cnaDf[cnaDf['Hugo_Symbol'] == i]
            newVal = dups[dups.columns[dups.columns != "Hugo_Symbol"]].apply(
                mergeCNAvalues)
            temp = pd.DataFrame(newVal).transpose()
            temp['Hugo_Symbol'] = i
            duplicatedGenes = duplicatedGenes.append(temp, sort=False)
        cnaDf.drop_duplicates('Hugo_Symbol', keep=False, inplace=True)
        cnaDf = cnaDf.append(duplicatedGenes, sort=False)
        cnaDf = cnaDf[order]
        cnaDf.columns = [
            process_functions.checkGenieId(i, self.center)
            if i != "Hugo_Symbol" else i for i in cnaDf.columns]

        return(cnaDf)

    def process_steps(self, cnaDf, newPath, databaseToSynIdMappingDf):

        newCNA = self._process(cnaDf)

        centerMafSynId = databaseToSynIdMappingDf.Id[
            databaseToSynIdMappingDf['Database'] == "centerMaf"][0]
        if not newCNA.empty:
            cnaText = process_functions.removePandasDfFloat(newCNA)
            # Replace blank with NA's
            cnaText = cnaText.replace(
                "\t\t", "\tNA\t").replace(
                "\t\t", "\tNA\t").replace(
                '\t\n', "\tNA\n")
            with open(newPath, "w") as cnaFile:
                cnaFile.write(cnaText)
            self.syn.store(synapseclient.File(newPath, parent=centerMafSynId))
        return(newPath)

    def _validate(self, cnvDF, nosymbol_check):
        total_error = ""
        warning = ""
        cnvDF.columns = [col.upper() for col in cnvDF.columns]

        if cnvDF.columns[0] != "HUGO_SYMBOL":
            total_error += "Your cnv file's first column must be Hugo_Symbol\n"
        haveColumn = process_functions.checkColExist(cnvDF, "HUGO_SYMBOL")
        if haveColumn:
            keepSymbols = cnvDF["HUGO_SYMBOL"]
            cnvDF.drop("HUGO_SYMBOL", axis=1, inplace=True)

        # if sum(cnvDF.apply(lambda x: sum(x.isnull()))) > 0:
        #   total_error += "Your cnv file must not have any empty values\n"

        if process_functions.checkColExist(cnvDF, "ENTREZ_GENE_ID"):
            del cnvDF['ENTREZ_GENE_ID']

        # cnvDF = cnvDF.fillna('')
        allowed_values = ['-2.0', '-2', '-1.5', '-1.0', '-1',
                          '0.0', '0', '0.5', '1.0', '1', '1.5',
                          '2', '2.0', 'nan']
        if not all(cnvDF.applymap(lambda x: str(x) in allowed_values).all()):
            total_error += (
                "All values must be NA/blank, -2, -1.5, -1, -0.5, "
                "0, 0.5, 1, 1.5, or 2.\n")
        else:
            cnvDF['HUGO_SYMBOL'] = keepSymbols
            if haveColumn and not nosymbol_check:
                bedSynId = process_functions.getDatabaseSynId(
                    self.syn, "bed", test=self.testing)
                bed = self.syn.tableQuery(
                    "select Hugo_Symbol, ID from {} where "
                    "CENTER = '{}'".format(bedSynId, self.center))
                bedDf = bed.asDataFrame()
                cnvDF['remapped'] = cnvDF['HUGO_SYMBOL'].apply(
                    lambda x: validateSymbol(x, bedDf))
                cnvDF = cnvDF[~cnvDF['remapped'].isnull()]

                # Do not allow any duplicated genes after symbols
                # have been remapped
                if sum(cnvDF['remapped'].duplicated()) > 0:
                    duplicated = cnvDF['remapped'].duplicated(keep=False)
                    total_error += (
                        "Your CNA file has duplicated Hugo_Symbols "
                        "(After remapping of genes): {} -> {}.\n".format(
                            ",".join(cnvDF['HUGO_SYMBOL'][duplicated]),
                            ",".join(cnvDF['remapped'][duplicated])))
        return(total_error, warning)
<|endoftext|>


<commit_before># -*- coding: utf-8 -*-
from django.db import models
from django.db.models.fields import FieldDoesNotExist
from django.db.models.related import RelatedObject
from django.utils import timezone
from django.utils.translation import ugettext as _
import logging
import re

from instagram.models import ApiModel
from m2m_history.fields import ManyToManyHistoryField

from . import fields
from .api import get_api
from .decorators import fetch_all

#from . import fields
#from .parser import get_replies
#api = get_api()

__all__ = ['User', 'Status', 'InstagramContentError', 'InstagramModel', 'InstagramManager', 'UserManager']

log = logging.getLogger('instagram_api')

class InstagramContentError(Exception):
    pass

class InstagramManager(models.Manager):

    '''
    Instagram Manager for RESTful CRUD operations
    '''

    def __init__(self, methods=None, remote_pk=None, *args, **kwargs):
        if methods and len(methods.items()) < 1:
            raise ValueError('Argument methods must contains at least 1 specified method')

        self.methods = methods or {}
        self.remote_pk = remote_pk or ('id',)
        if not isinstance(self.remote_pk, tuple):
            self.remote_pk = (self.remote_pk,)

        super(InstagramManager, self).__init__(*args, **kwargs)

    def get_by_url(self, url):
        '''
        Return object by url
        '''
        m = re.findall(r'(?:https?://)?(?:www\.)?instagram\.com/([^/]+)/?', url)
        if not len(m):
            raise ValueError("Url should be started with https://instagram.com/")

        return self.get_by_slug(m[0])

    def get_by_slug(self, slug):
        '''
        Return object by slug
        '''
        # TODO: change to self.get method
        return self.model.remote.fetch(slug)

    def get_or_create_from_instance(self, instance):

        remote_pk_dict = {}
        for field_name in self.remote_pk:
            remote_pk_dict[field_name] = getattr(instance, field_name)

        try:
            old_instance = self.model.objects.get(**remote_pk_dict)
            instance._substitute(old_instance)
            instance.save()
        except self.model.DoesNotExist:
            instance.save()
            log.debug('Fetch and create new object %s with remote pk %s' % (self.model, remote_pk_dict))

        return instance

#    def get_or_create_from_resource(self, resource):
#
#        instance = self.model()
#        instance.parse(dict(resource))
#
#        return self.get_or_create_from_instance(instance)

    def api_call(self, method, *args, **kwargs):
        api = get_api()

        if method in self.methods:
            method = self.methods[method]
        return getattr(api, method)(*args, **kwargs)

    def fetch(self, *args, **kwargs):
        '''
        Retrieve and save object to local DB
        '''
        result = self.get(*args, **kwargs)
        if isinstance(result, list):
            return [self.get_or_create_from_instance(instance) for instance in result]
        else:
            return self.get_or_create_from_instance(result)

    def get(self, *args, **kwargs):
        '''
        Retrieve objects from remote server
        '''
        extra_fields = kwargs.pop('extra_fields', {})
        extra_fields['fetched'] = timezone.now()
        response = self.api_call('get', *args, **kwargs)

        return self.parse_response(response, extra_fields)

    def parse_response(self, response, extra_fields=None):
        # if response is None:
        #     return []
        # el
        if isinstance(response, (list, tuple)):
            return self.parse_response_list(response, extra_fields)
        elif isinstance(response, ApiModel):
            return self.parse_response_object(response, extra_fields)
        else:
            raise InstagramContentError('Instagram response should be list or dict, not %s' % response)

    def parse_response_object(self, resource, extra_fields=None):

        instance = self.model()
        # important to do it before calling parse method
        if extra_fields:
            instance.__dict__.update(extra_fields)
        instance.set_tweepy(resource)
        instance.parse()

        return instance

    def parse_response_list(self, response_list, extra_fields=None):

        instances = []
        for response in response_list:

            if not isinstance(response, ApiModel):
                log.error("Resource %s is not dictionary" % response)
                continue

            instance = self.parse_response_object(response, extra_fields)
            instances += [instance]

        return instances

class UserManager(InstagramManager):

    def get_followers_ids_for_user(self, user, all=False, count=5000, **kwargs):
        # https://dev.instagram.com/docs/api/1.1/get/followers/ids
        if all:
            cursor = tweepy.Cursor(user.tweepy._api.followers_ids, id=user.pk, count=count)
            return list(cursor.items())
        else:
            raise NotImplementedError("This method implemented only with argument all=True")

    def fetch_followers_for_user(self, user, all=False, count=200, **kwargs):
        # https://dev.instagram.com/docs/api/1.1/get/followers/list
        # in docs default count is 20, but maximum is 200
        if all:
            # TODO: make optimization: break cursor iteration after getting already
            # existing user and switch to ids REST method
            user.followers.clear()
            cursor = tweepy.Cursor(user.tweepy._api.followers, id=user.pk, count=count)
            for instance in cursor.items():
                instance = self.parse_response_object(instance)
                instance = self.get_or_create_from_instance(instance)
                user.followers.add(instance)
        else:
            raise NotImplementedError("This method implemented only with argument all=True")
        return user.followers.all()

    def _______get_or_create_from_instance(self, instance):
        try:
            instance_old = self.model.objects.get(screen_name=instance.screen_name)
            if instance_old.pk == instance.pk:
                instance.save()
            else:
                # perhaps we already have old User with the same screen_name, but different id
                try:
                    self.fetch(instance_old.pk)
                except InstagramError, e:
                    if e.code == 34:
                        instance_old.delete()
                        instance.save()
                    else:
                        raise
            return instance
        except self.model.DoesNotExist:
            return super(UserManager, self).get_or_create_from_instance(instance)

class InstagramModel(models.Model):

    objects = models.Manager()

    _foreignkeys_pre_save = []

    class Meta:
        abstract = True

    def _substitute(self, old_instance):
        '''
        Substitute new user with old one while updating in method Manager.get_or_create_from_instance()
        Can be overrided in child models
        '''
        self.pk = old_instance.pk
        self.created_at = old_instance.created_at

    def save(self, *args, **kwargs):
        '''
        Save all related instances before or after current instance
        '''
        for field, instance in self._foreignkeys_pre_save:
            instance = instance.__class__.remote.get_or_create_from_instance(instance)
            setattr(self, field, instance)
        self._foreignkeys_pre_save = []

        super(InstagramModel, self).save(*args, **kwargs)

    def parse(self):
        '''
        Parse API response and define fields with values
        '''
        for key, value in self._response.items():
            if key == '_api':
                continue

            try:
                field, model, direct, m2m = self._meta.get_field_by_name(key)
            except FieldDoesNotExist:
                log.debug('Field with name "%s" doesn\'t exist in the model %s' % (key, type(self)))
                continue

            if isinstance(field, RelatedObject) and value:
                for item in value:
                    rel_instance = field.model.remote.parse_response_object(item)
                    self._external_links_post_save += [(field.field.name, rel_instance)]
            else:
                if isinstance(field, (models.BooleanField)):
                    value = bool(value)

                elif isinstance(field, (models.OneToOneField, models.ForeignKey)) and value:
                    rel_instance = field.rel.to.remote.parse_response_object(value)
                    value = rel_instance
                    if isinstance(field, models.ForeignKey):
                        self._foreignkeys_pre_save += [(key, rel_instance)]

                elif isinstance(field, (fields.CommaSeparatedCharField, models.CommaSeparatedIntegerField)) and isinstance(value, list):
                    value = ','.join([unicode(v) for v in value])

                elif isinstance(field, (models.CharField, models.TextField)) and value:
                    if isinstance(value, (str, unicode)):
                        value = value.strip()

                setattr(self, key, value)

class InstagramBaseModel(InstagramModel):

    _tweepy_model = None
    _response = None

    #id = models.BigIntegerField(primary_key=True)
    created_at = models.DateTimeField(auto_now_add=True)
    fetched = models.DateTimeField(u'Fetched', null=True, blank=True)
    #lang = models.CharField(max_length=10)
    #entities = fields.JSONField()

    class Meta:
        abstract = True

    def set_tweepy(self, model):
        self._tweepy_model = model
        self._response = dict(self._tweepy_model.__dict__)

    @property
    def tweepy(self):
        if not self._tweepy_model:
            # get fresh instance with the same ID, set tweepy object and refresh attributes
            instance = self.__class__.remote.get(self.pk)
            self.set_tweepy(instance.tweepy)
            self.parse()
        return self._tweepy_model

    def get_url(self):
        return 'https://instagram.com/%s' % self.slug

class User(InstagramBaseModel):

    username = models.CharField(max_length=50, unique=True)
    full_name = models.CharField(max_length=255)
    bio = models.CharField("BIO", max_length=255)

    profile_picture = models.URLField(max_length=300)
    website = models.URLField(max_length=300)

    followers_count = models.PositiveIntegerField(null=True)
    media_count = models.PositiveIntegerField(null=True)

    #followers = ManyToManyHistoryField('User', versions=True)

    objects = models.Manager()
    remote = UserManager(methods={
        'get': 'user',
    })

    def __unicode__(self):
        return self.username

    def parse(self):
        #self._response['favorites_count'] = self._response.pop('favourites_count', None)
        #self._response.pop('status', None)
        super(User, self).parse()

    def fetch_followers(self, **kwargs):
        return User.remote.fetch_followers_for_user(user=self, **kwargs)

    def get_followers_ids(self, **kwargs):
        return User.remote.get_followers_ids_for_user(user=self, **kwargs)

    def fetch_statuses(self, **kwargs):
        return Status.remote.fetch_for_user(user=self, **kwargs)

class MediaManager(InstagramManager):

    #def fetch(id):
    pass

class Media(InstagramBaseModel):

    id = models.CharField(max_length=100, primary_key=True)
    caption = models.CharField(max_length=255)
    link = models.URLField(max_length=300)

    #tags =
    created_time = models.DateTimeField(auto_now_add=True)

    comment_count = models.PositiveIntegerField(null=True)
    like_count = models.PositiveIntegerField(null=True)

    user = models.ForeignKey(User, related_name="media_feed")

    remote = MediaManager(methods={
        'get': 'media',
    })
<commit_msg>Comment model + fixes<commit_after># -*- coding: utf-8 -*-
from django.db import models
from django.db.models.fields import FieldDoesNotExist
from django.db.models.related import RelatedObject
from django.utils import timezone
from django.utils.translation import ugettext as _
import logging
import re

from instagram.models import ApiModel
from m2m_history.fields import ManyToManyHistoryField

from . import fields
from .api import get_api
from .decorators import fetch_all

#from . import fields
#from .parser import get_replies
#api = get_api()

__all__ = ['User', 'Status', 'InstagramContentError', 'InstagramModel', 'InstagramManager', 'UserManager']

log = logging.getLogger('instagram_api')

class InstagramContentError(Exception):
    pass

class InstagramManager(models.Manager):

    '''
    Instagram Manager for RESTful CRUD operations
    '''

    def __init__(self, methods=None, remote_pk=None, *args, **kwargs):
        self.api = get_api()

        if methods and len(methods.items()) < 1:
            raise ValueError('Argument methods must contains at least 1 specified method')

        self.methods = methods or {}
        self.remote_pk = remote_pk or ('id',)
        if not isinstance(self.remote_pk, tuple):
            self.remote_pk = (self.remote_pk,)

        super(InstagramManager, self).__init__(*args, **kwargs)

    def get_by_url(self, url):
        '''
        Return object by url
        '''
        m = re.findall(r'(?:https?://)?(?:www\.)?instagram\.com/([^/]+)/?', url)
        if not len(m):
            raise ValueError("Url should be started with https://instagram.com/")

        return self.get_by_slug(m[0])

    def get_by_slug(self, slug):
        '''
        Return object by slug
        '''
        # TODO: change to self.get method
        return self.model.remote.fetch(slug)

    def get_or_create_from_instance(self, instance):

        remote_pk_dict = {}
        for field_name in self.remote_pk:
            remote_pk_dict[field_name] = getattr(instance, field_name)

        try:
            old_instance = self.model.objects.get(**remote_pk_dict)
            instance._substitute(old_instance)
            instance.save()
        except self.model.DoesNotExist:
            instance.save()
            log.debug('Fetch and create new object %s with remote pk %s' % (self.model, remote_pk_dict))

        return instance

#    def get_or_create_from_resource(self, resource):
#
#        instance = self.model()
#        instance.parse(dict(resource))
#
#        return self.get_or_create_from_instance(instance)

    def api_call(self, method, *args, **kwargs):
        if method in self.methods:
            method = self.methods[method]
        return getattr(self.api, method)(*args, **kwargs)

    def fetch(self, *args, **kwargs):
        '''
        Retrieve and save object to local DB
        '''
        result = self.get(*args, **kwargs)
        if isinstance(result, list):
            return [self.get_or_create_from_instance(instance) for instance in result]
        else:
            return self.get_or_create_from_instance(result)

    def get(self, *args, **kwargs):
        '''
        Retrieve objects from remote server
        '''
        extra_fields = kwargs.pop('extra_fields', {})
        extra_fields['fetched'] = timezone.now()
        response = self.api_call('get', *args, **kwargs)

        return self.parse_response(response, extra_fields)

    def parse_response(self, response, extra_fields=None):
        # if response is None:
        #     return []
        # el
        if isinstance(response, (list, tuple)):
            return self.parse_response_list(response, extra_fields)
        elif isinstance(response, ApiModel):
            return self.parse_response_object(response, extra_fields)
        else:
            raise InstagramContentError('Instagram response should be list or dict, not %s' % response)

    def parse_response_object(self, resource, extra_fields=None):

        instance = self.model()
        # important to do it before calling parse method
        if extra_fields:
            instance.__dict__.update(extra_fields)
        instance.set_tweepy(resource)
        instance.parse()

        return instance

    def parse_response_list(self, response_list, extra_fields=None):

        instances = []
        for response in response_list:

            if not isinstance(response, ApiModel):
                log.error("Resource %s is not dictionary" % response)
                continue

            instance = self.parse_response_object(response, extra_fields)
            instances += [instance]

        return instances

class UserManager(InstagramManager):

    def get_followers_ids_for_user(self, user, all=False, count=5000, **kwargs):
        # https://dev.instagram.com/docs/api/1.1/get/followers/ids
        if all:
            cursor = tweepy.Cursor(user.tweepy._api.followers_ids, id=user.pk, count=count)
            return list(cursor.items())
        else:
            raise NotImplementedError("This method implemented only with argument all=True")

    def fetch_followers_for_user(self, user, all=False, count=200, **kwargs):
        # https://dev.instagram.com/docs/api/1.1/get/followers/list
        # in docs default count is 20, but maximum is 200
        if all:
            # TODO: make optimization: break cursor iteration after getting already
            # existing user and switch to ids REST method
            user.followers.clear()
            cursor = tweepy.Cursor(user.tweepy._api.followers, id=user.pk, count=count)
            for instance in cursor.items():
                instance = self.parse_response_object(instance)
                instance = self.get_or_create_from_instance(instance)
                user.followers.add(instance)
        else:
            raise NotImplementedError("This method implemented only with argument all=True")
        return user.followers.all()

    def _______get_or_create_from_instance(self, instance):
        try:
            instance_old = self.model.objects.get(screen_name=instance.screen_name)
            if instance_old.pk == instance.pk:
                instance.save()
            else:
                # perhaps we already have old User with the same screen_name, but different id
                try:
                    self.fetch(instance_old.pk)
                except InstagramError, e:
                    if e.code == 34:
                        instance_old.delete()
                        instance.save()
                    else:
                        raise
            return instance
        except self.model.DoesNotExist:
            return super(UserManager, self).get_or_create_from_instance(instance)

class InstagramModel(models.Model):

    objects = models.Manager()

    class Meta:
        abstract = True

    def __init__(self, *args, **kwargs):
        super(InstagramModel, self).__init__(*args, **kwargs)

        # different lists for saving related objects
        self._external_links_post_save = []
        self._foreignkeys_pre_save = []
        self._external_links_to_add = []

    def _substitute(self, old_instance):
        '''
        Substitute new user with old one while updating in method Manager.get_or_create_from_instance()
        Can be overrided in child models
        '''
        self.pk = old_instance.pk
        self.created_at = old_instance.created_at

    def save(self, *args, **kwargs):
        '''
        Save all related instances before or after current instance
        '''
        for field, instance in self._foreignkeys_pre_save:
            instance = instance.__class__.remote.get_or_create_from_instance(instance)
            setattr(self, field, instance)
        self._foreignkeys_pre_save = []

        super(InstagramModel, self).save(*args, **kwargs)

    def parse(self):
        '''
        Parse API response and define fields with values
        '''
        for key, value in self._response.items():
            if key == '_api':
                continue

            try:
                field, model, direct, m2m = self._meta.get_field_by_name(key)
            except FieldDoesNotExist:
                log.debug('Field with name "%s" doesn\'t exist in the model %s' % (key, type(self)))
                continue

            if isinstance(field, RelatedObject) and value:
                for item in value:
                    rel_instance = field.model.remote.parse_response_object(item)
                    self._external_links_post_save += [(field.field.name, rel_instance)]
            else:
                if isinstance(field, (models.BooleanField)):
                    value = bool(value)

                elif isinstance(field, (models.OneToOneField, models.ForeignKey)) and value:
                    rel_instance = field.rel.to.remote.parse_response_object(value)
                    value = rel_instance
                    if isinstance(field, models.ForeignKey):
                        self._foreignkeys_pre_save += [(key, rel_instance)]

                elif isinstance(field, (fields.CommaSeparatedCharField, models.CommaSeparatedIntegerField)) and isinstance(value, list):
                    value = ','.join([unicode(v) for v in value])

                elif isinstance(field, (models.CharField, models.TextField)) and value:
                    if isinstance(value, (str, unicode)):
                        value = value.strip()

                setattr(self, key, value)

class InstagramBaseModel(InstagramModel):

    _tweepy_model = None
    _response = None

    #id = models.BigIntegerField(primary_key=True)
    created_at = models.DateTimeField(auto_now_add=True)
    fetched = models.DateTimeField(u'Fetched', null=True, blank=True)
    #lang = models.CharField(max_length=10)
    #entities = fields.JSONField()

    class Meta:
        abstract = True

    def set_tweepy(self, model):
        self._tweepy_model = model
        self._response = dict(self._tweepy_model.__dict__)

    @property
    def tweepy(self):
        if not self._tweepy_model:
            # get fresh instance with the same ID, set tweepy object and refresh attributes
            instance = self.__class__.remote.get(self.pk)
            self.set_tweepy(instance.tweepy)
            self.parse()
        return self._tweepy_model

    def get_url(self):
        return 'https://instagram.com/%s' % self.slug

class User(InstagramBaseModel):

    username = models.CharField(max_length=50, unique=True)
    full_name = models.CharField(max_length=255)
    bio = models.CharField("BIO", max_length=255)

    profile_picture = models.URLField(max_length=300)
    website = models.URLField(max_length=300)

    followers_count = models.PositiveIntegerField(null=True)
    media_count = models.PositiveIntegerField(null=True)

    #followers = ManyToManyHistoryField('User', versions=True)

    objects = models.Manager()
    remote = UserManager(methods={
        'get': 'user',
    })

    def __unicode__(self):
        return self.username

    def parse(self):
        #self._response['favorites_count'] = self._response.pop('favourites_count', None)
        #self._response.pop('status', None)
        super(User, self).parse()

    def fetch_followers(self, **kwargs):
        return User.remote.fetch_followers_for_user(user=self, **kwargs)

    def get_followers_ids(self, **kwargs):
        return User.remote.get_followers_ids_for_user(user=self, **kwargs)

    def fetch_statuses(self, **kwargs):
        return Status.remote.fetch_for_user(user=self, **kwargs)

class MediaManager(InstagramManager):
    pass

class Media(InstagramBaseModel):

    id = models.CharField(max_length=100, primary_key=True)
    caption = models.CharField(max_length=255)
    link = models.URLField(max_length=300)

    #tags =
    created_time = models.DateTimeField()

    comment_count = models.PositiveIntegerField(null=True)
    like_count = models.PositiveIntegerField(null=True)

    user = models.ForeignKey(User, related_name="media_feed")

    remote = MediaManager(methods={
        'get': 'media',
    })

    def fetch_comments(self):
        return Comment.remote.fetch_media_comments(self)

class CommentManager(InstagramManager):
    def fetch_media_comments(self, media):

        extra_fields = {}
        extra_fields['fetched'] = timezone.now()
        extra_fields['media_id'] = media.id

        #comments
        response = self.api.media_comments(media.id)
        result = self.parse_response(response, extra_fields)
        return [self.get_or_create_from_instance(instance) for instance in result]

class Comment(InstagramBaseModel):

    user = models.ForeignKey(User)
    media = models.ForeignKey(Media, related_name="comments")

    id = models.BigIntegerField(primary_key=True)
    text = models.TextField()
    #created_at = models.DateTimeField()

    remote = CommentManager()

<|endoftext|>


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

import os, sys
sys.path.append(os.path.join(os.path.dirname(__file__), '..', 'lib'))

import numpy as np
import glob, logging, math, queue, socket, subprocess, threading
import tensorflow as tf

from support import Config
from task_usage import DistributedDatabase
import support

class Learn:
    def __init__(self, config):
        graph = tf.Graph()
        with graph.as_default():
            model = Model(config)
            with tf.variable_scope('optimization'):
                state = tf.Variable(
                    [0, 0], name='state', dtype=tf.int64, trainable=False)
                state_update = tf.placeholder(
                    tf.int64, shape=(2), name='state_update')
                update_state = state.assign(state_update)
                parameters = tf.trainable_variables()
                gradient = tf.gradients(model.loss, parameters)
                gradient, _ = tf.clip_by_global_norm(
                    gradient, config.gradient_clip)
                optimizer = tf.train.AdamOptimizer(config.learning_rate)
                train = optimizer.apply_gradients(zip(gradient, parameters))
            with tf.variable_scope('summary'):
                tf.summary.scalar(
                    'log_loss', tf.log(tf.reduce_sum(model.loss)))
            logger = tf.summary.FileWriter(config.log_path, graph)
            summary = tf.summary.merge_all()
            initialize = tf.variables_initializer(
                tf.global_variables(), name='initialize')
            saver = Saver(config)

        self.graph = graph
        self.model = model
        self.state = state
        self.state_update = state_update
        self.update_state = update_state
        self.parameters = parameters
        self.train = train
        self.logger = logger
        self.summary = summary
        self.initialize = initialize
        self.saver = saver

    def count_parameters(self):
        return np.sum([int(np.prod(p.get_shape())) for p in self.parameters])

    def run(self, target, monitor, config):
        support.log(self, 'Parameters: {}', self.count_parameters())
        support.log(self, 'Samples: {}', target.sample_count)
        session = tf.Session(graph=self.graph)
        session.run(self.initialize)
        self.saver.restore(session)
        state = State.deserialize(session.run(self.state))
        for _ in range(config.epoch_count - state.epoch % config.epoch_count):
            self._run_epoch(target, monitor, config, session, state)
            state.increment_epoch()
            session.run(self.update_state, {
                self.state_update: state.serialize(),
            })
            self.saver.save(session)

    def _run_epoch(self, target, monitor, config, session, state):
        for _ in range(target.sample_count):
            if monitor.should_train(state.time):
                self._run_train(target, monitor, config, session, state)
            if monitor.should_predict(state.time):
                self._run_predict(target, monitor, config, session, state)
            state.increment_time()

    def _run_train(self, target, monitor, config, session, state):
        sample = target.compute(state.sample)
        feed = {
            self.model.start: self._zero_start(),
            self.model.x: np.reshape(sample, [1, -1, target.dimension_count]),
            self.model.y: np.reshape(
                support.shift(sample, -1), [1, -1, target.dimension_count]),
        }
        fetch = {
            'train': self.train,
            'loss': self.model.loss,
            'summary': self.summary,
        }
        result = session.run(fetch, feed)
        loss = result['loss'].flatten()
        assert(np.all([not math.isnan(loss) for loss in loss]))
        monitor.train(loss, state)
        self.logger.add_summary(result['summary'], state.time)

    def _run_predict(self, target, monitor, config, session, state):
        sample = target.compute((state.sample + 1) % target.sample_count)
        step_count = sample.shape[0]
        feed = {
            self.model.start: self._zero_start(),
        }
        fetch = {
            'y_hat': self.model.y_hat,
            'finish': self.model.finish,
        }
        for i in range(step_count):
            feed[self.model.x] = np.reshape(
                sample[:(i + 1), :], [1, i + 1, -1])
            y_hat = np.zeros([step_count, target.dimension_count])
            for j in range(step_count - i - 1):
                result = session.run(fetch, feed)
                feed[self.model.start] = result['finish']
                y_hat[j, :] = result['y_hat'][-1, :]
                feed[self.model.x] = np.reshape(y_hat[j, :], [1, 1, -1])
            if not monitor.predict(support.shift(sample, -i - 1), y_hat):
                break

    def _zero_start(self):
        return np.zeros(self.model.start.get_shape(), np.float32)

class Model:
    def __init__(self, config):
        x = tf.placeholder(
            tf.float32, [1, None, config.dimension_count], name='x')
        y = tf.placeholder(
            tf.float32, [1, None, config.dimension_count], name='y')
        with tf.variable_scope('network') as scope:
            cell = tf.nn.rnn_cell.LSTMCell(
                config.unit_count, state_is_tuple=True,
                cell_clip=config.cell_clip, forget_bias=config.forget_bias,
                use_peepholes=config.use_peepholes,
                initializer=config.network_initializer)
            cell = tf.nn.rnn_cell.MultiRNNCell(
                [cell] * config.layer_count, state_is_tuple=True)
            start, state = Model._initialize(config)
            h, state = tf.nn.dynamic_rnn(
                cell, x, initial_state=state, parallel_iterations=1)
            finish = Model._finalize(state, config)
        y_hat, loss = Model._regress(h, y, config)

        self.x = x
        self.y = y
        self.y_hat = y_hat
        self.loss = loss
        self.start = start
        self.finish = finish

    def _finalize(state, config):
        parts = []
        for i in range(config.layer_count):
            parts.append(state[i].c)
            parts.append(state[i].h)
        return tf.pack(parts, name='finish')

    def _initialize(config):
        start = tf.placeholder(
            tf.float32, [2 * config.layer_count, 1, config.unit_count],
            name='start')
        parts = tf.unpack(start)
        state = []
        for i in range(config.layer_count):
            c, h = parts[2 * i], parts[2*i + 1]
            state.append(tf.nn.rnn_cell.LSTMStateTuple(c, h))
        return start, tuple(state)

    def _regress(x, y, config):
        with tf.variable_scope('regression') as scope:
            unroll_count = tf.shape(x)[1]
            x = tf.squeeze(x, squeeze_dims=[0])
            y = tf.squeeze(y, squeeze_dims=[0])
            w = tf.get_variable(
                'w', [config.unit_count, config.dimension_count],
                initializer=config.regression_initializer)
            b = tf.get_variable('b', [1, config.dimension_count])
            y_hat = tf.matmul(x, w) + tf.tile(b, [unroll_count, 1])
            loss = tf.reduce_mean(tf.squared_difference(y_hat, y))
        return y_hat, loss

class Monitor:
    def __init__(self, config):
        assert(np.all(np.diff(config.work_schedule) >= 0))
        self.bind_address = config.bind_address
        self.work_schedule = config.work_schedule
        self.channels = {}
        self.lock = threading.Lock()
        threading.Thread(target=self._predict_server, daemon=True).start()

    def should_train(self, time):
        return True

    def should_predict(self, time):
        if len(self.channels) == 0:
            return False
        time = time % self.work_schedule[-1]
        phase = np.nonzero(self.work_schedule >= time)[0][0]
        if phase % 2 != 1:
            return False
        return True

    def train(self, loss, state):
        sys.stdout.write(
            '%10d %4d %10d' % (state.time, state.epoch, state.sample))
        [sys.stdout.write(' %12.4e' % loss) for loss in loss]
        sys.stdout.write('\n')

    def predict(self, y, y_hat):
        self.lock.acquire()
        try:
            for channel in self.channels:
                channel.put((y, y_hat))
        finally:
            self.lock.release()
        return len(self.channels) > 0

    def _predict_client(self, connection, address):
        support.log(self, 'Start serving {}.', address)
        channel = queue.Queue()
        self.lock.acquire()
        try:
            self.channels[channel] = True
        finally:
            self.lock.release()
        try:
            client = connection.makefile(mode="w")
            while True:
                y, y_hat = channel.get()
                values = [str(value) for value in y.flatten()]
                client.write(','.join(values) + ',')
                values = [str(value) for value in y_hat.flatten()]
                client.write(','.join(values) + '\n')
        except Exception as e:
            support.log(self, 'Stop serving {} ({}).', address, e)
        self.lock.acquire()
        try:
            del self.channels[channel]
        finally:
            self.lock.release()

    def _predict_server(self):
        server = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
        server.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
        server.bind(self.bind_address)
        server.listen(1)
        support.log(self, 'Listening to {}...', self.bind_address)
        while True:
            try:
                connection, address = server.accept()
                worker = threading.Thread(daemon=True,
                                          target=self._predict_client,
                                          args=(connection, address))
                worker.start()
            except Exception as e:
                support.log(self, 'Encountered a problem ({}).', e)

class Saver:
    def __init__(self, config):
        self.backend = tf.train.Saver()
        self.path = config.save_path

    def save(self, session):
        path = self.backend.save(session, self.path)
        support.log(self, 'Saved in "{}".', path)

    def restore(self, session):
        if len(glob.glob("{}*".format(self.path))) > 0:
            if input('Restore from "{}"? '.format(self.path)) != 'no':
                self.backend.restore(session, self.path)
                support.log(self, 'Restored. Continue learning...')

class State:
    def deserialize(state):
        return State(state[0], state[1])

    def __init__(self, time, epoch):
        self.time = time
        self.epoch = epoch
        self.sample = 0

    def increment_epoch(self):
        self.epoch += 1
        self.sample = 0

    def increment_time(self):
        self.time += 1
        self.sample += 1

    def serialize(self):
        return [self.time, self.epoch]

class Target:
    def __init__(self, config):
        self.database = DistributedDatabase(config.data_path)
        self.dimension_count = 1
        self.sample_count = self.database.count()

    def compute(self, k):
        raise Exception('not implemented yet')

class TestTarget:
    def __init__(self, config):
        self.dimension_count = 1
        self.sample_count = 100000

    def compute(self, k):
        return np.reshape(np.sin(4 * np.pi / 40 * np.arange(0, 40)), [-1, 1])

def main(config):
    learn = Learn(config)
    target = Target(config)
    monitor = Monitor(config)
    learn.run(target, monitor, config)

if __name__ == '__main__':
    logging.basicConfig(format='%(asctime)s %(message)s',
                        datefmt='%Y-%m-%d %H:%M:%S',
                        level=logging.INFO)
    assert(len(sys.argv) == 2)
    config = Config({
        'dimension_count': 1,
        'data_path': sys.argv[1],
        'layer_count': 1,
        'unit_count': 200,
        'cell_clip': 1.0,
        'forget_bias': 1.0,
        'use_peepholes': True,
        'network_initializer': tf.random_uniform_initializer(-0.01, 0.01),
        'regression_initializer': tf.random_normal_initializer(stddev=0.01),
        'learning_rate': 1e-3,
        'gradient_clip': 1.0,
        'epoch_count': 100,
        'log_path': os.path.join('output', 'log'),
        'save_path': os.path.join('output', 'model'),
        'bind_address': ('0.0.0.0', 4242),
        'work_schedule': np.cumsum([1000 - 10, 10]),
    })
    main(config)
<commit_msg>Make a cosmetic adjustment<commit_after>#!/usr/bin/env python3

import os, sys
sys.path.append(os.path.join(os.path.dirname(__file__), '..', 'lib'))

import numpy as np
import glob, logging, math, queue, socket, subprocess, threading
import tensorflow as tf

from support import Config
from task_usage import DistributedDatabase
import support

class Learn:
    def __init__(self, config):
        graph = tf.Graph()
        with graph.as_default():
            model = Model(config)
            with tf.variable_scope('optimization'):
                state = tf.Variable(
                    [0, 0], name='state', dtype=tf.int64, trainable=False)
                state_update = tf.placeholder(
                    tf.int64, shape=(2), name='state_update')
                update_state = state.assign(state_update)
                parameters = tf.trainable_variables()
                gradient = tf.gradients(model.loss, parameters)
                gradient, _ = tf.clip_by_global_norm(
                    gradient, config.gradient_clip)
                optimizer = tf.train.AdamOptimizer(config.learning_rate)
                train = optimizer.apply_gradients(zip(gradient, parameters))
            with tf.variable_scope('summary'):
                tf.summary.scalar(
                    'log_loss', tf.log(tf.reduce_sum(model.loss)))
            logger = tf.summary.FileWriter(config.log_path, graph)
            summary = tf.summary.merge_all()
            initialize = tf.variables_initializer(
                tf.global_variables(), name='initialize')
            saver = Saver(config)

        self.graph = graph
        self.model = model
        self.state = state
        self.state_update = state_update
        self.update_state = update_state
        self.parameters = parameters
        self.train = train
        self.logger = logger
        self.summary = summary
        self.initialize = initialize
        self.saver = saver

    def count_parameters(self):
        return np.sum([int(np.prod(p.get_shape())) for p in self.parameters])

    def run(self, target, monitor, config):
        support.log(self, 'Parameters: {}', self.count_parameters())
        support.log(self, 'Samples: {}', target.sample_count)
        session = tf.Session(graph=self.graph)
        session.run(self.initialize)
        self.saver.restore(session)
        state = State.deserialize(session.run(self.state))
        for _ in range(config.epoch_count - state.epoch % config.epoch_count):
            self._run_epoch(target, monitor, config, session, state)
            state.increment_epoch()
            session.run(self.update_state, {
                self.state_update: state.serialize(),
            })
            self.saver.save(session)

    def _run_epoch(self, target, monitor, config, session, state):
        for _ in range(target.sample_count):
            if monitor.should_train(state.time):
                self._run_train(target, monitor, config, session, state)
            if monitor.should_predict(state.time):
                self._run_predict(target, monitor, config, session, state)
            state.increment_time()

    def _run_train(self, target, monitor, config, session, state):
        sample = target.compute(state.sample)
        feed = {
            self.model.start: self._zero_start(),
            self.model.x: np.reshape(sample, [1, -1, target.dimension_count]),
            self.model.y: np.reshape(
                support.shift(sample, -1), [1, -1, target.dimension_count]),
        }
        fetch = {
            'train': self.train,
            'loss': self.model.loss,
            'summary': self.summary,
        }
        result = session.run(fetch, feed)
        loss = result['loss'].flatten()
        assert(np.all([not math.isnan(loss) for loss in loss]))
        monitor.train(loss, state)
        self.logger.add_summary(result['summary'], state.time)

    def _run_predict(self, target, monitor, config, session, state):
        sample = target.compute((state.sample + 1) % target.sample_count)
        step_count = sample.shape[0]
        feed = {
            self.model.start: self._zero_start(),
        }
        fetch = {
            'y_hat': self.model.y_hat,
            'finish': self.model.finish,
        }
        for i in range(step_count):
            feed[self.model.x] = np.reshape(
                sample[:(i + 1), :], [1, i + 1, -1])
            y_hat = np.zeros([step_count, target.dimension_count])
            for j in range(step_count - i - 1):
                result = session.run(fetch, feed)
                feed[self.model.start] = result['finish']
                y_hat[j, :] = result['y_hat'][-1, :]
                feed[self.model.x] = np.reshape(y_hat[j, :], [1, 1, -1])
            if not monitor.predict(support.shift(sample, -i - 1), y_hat):
                break

    def _zero_start(self):
        return np.zeros(self.model.start.get_shape(), np.float32)

class Model:
    def __init__(self, config):
        x = tf.placeholder(
            tf.float32, [1, None, config.dimension_count], name='x')
        y = tf.placeholder(
            tf.float32, [1, None, config.dimension_count], name='y')
        with tf.variable_scope('network') as scope:
            cell = tf.nn.rnn_cell.LSTMCell(
                config.unit_count, state_is_tuple=True,
                cell_clip=config.cell_clip, forget_bias=config.forget_bias,
                use_peepholes=config.use_peepholes,
                initializer=config.network_initializer)
            cell = tf.nn.rnn_cell.MultiRNNCell(
                [cell] * config.layer_count, state_is_tuple=True)
            start, state = Model._initialize(config)
            h, state = tf.nn.dynamic_rnn(
                cell, x, initial_state=state, parallel_iterations=1)
            finish = Model._finalize(state, config)
        y_hat, loss = Model._regress(h, y, config)

        self.x = x
        self.y = y
        self.y_hat = y_hat
        self.loss = loss
        self.start = start
        self.finish = finish

    def _finalize(state, config):
        parts = []
        for i in range(config.layer_count):
            parts.append(state[i].c)
            parts.append(state[i].h)
        return tf.pack(parts, name='finish')

    def _initialize(config):
        start = tf.placeholder(
            tf.float32, [2 * config.layer_count, 1, config.unit_count],
            name='start')
        parts = tf.unpack(start)
        state = []
        for i in range(config.layer_count):
            c, h = parts[2 * i], parts[2*i + 1]
            state.append(tf.nn.rnn_cell.LSTMStateTuple(c, h))
        return start, tuple(state)

    def _regress(x, y, config):
        with tf.variable_scope('regression') as scope:
            unroll_count = tf.shape(x)[1]
            x = tf.squeeze(x, squeeze_dims=[0])
            y = tf.squeeze(y, squeeze_dims=[0])
            w = tf.get_variable(
                'w', [config.unit_count, config.dimension_count],
                initializer=config.regression_initializer)
            b = tf.get_variable('b', [1, config.dimension_count])
            y_hat = tf.matmul(x, w) + tf.tile(b, [unroll_count, 1])
            loss = tf.reduce_mean(tf.squared_difference(y_hat, y))
        return y_hat, loss

class Monitor:
    def __init__(self, config):
        assert(np.all(np.diff(config.work_schedule) >= 0))
        self.bind_address = config.bind_address
        self.work_schedule = config.work_schedule
        self.channels = {}
        self.lock = threading.Lock()
        worker = threading.Thread(daemon=True, target=self._predict_server)
        worker.start()

    def should_train(self, time):
        return True

    def should_predict(self, time):
        if len(self.channels) == 0:
            return False
        time = time % self.work_schedule[-1]
        phase = np.nonzero(self.work_schedule >= time)[0][0]
        if phase % 2 != 1:
            return False
        return True

    def train(self, loss, state):
        sys.stdout.write(
            '%10d %4d %10d' % (state.time, state.epoch, state.sample))
        [sys.stdout.write(' %12.4e' % loss) for loss in loss]
        sys.stdout.write('\n')

    def predict(self, y, y_hat):
        self.lock.acquire()
        try:
            for channel in self.channels:
                channel.put((y, y_hat))
        finally:
            self.lock.release()
        return len(self.channels) > 0

    def _predict_client(self, connection, address):
        support.log(self, 'Start serving {}.', address)
        channel = queue.Queue()
        self.lock.acquire()
        try:
            self.channels[channel] = True
        finally:
            self.lock.release()
        try:
            client = connection.makefile(mode="w")
            while True:
                y, y_hat = channel.get()
                values = [str(value) for value in y.flatten()]
                client.write(','.join(values) + ',')
                values = [str(value) for value in y_hat.flatten()]
                client.write(','.join(values) + '\n')
        except Exception as e:
            support.log(self, 'Stop serving {} ({}).', address, e)
        self.lock.acquire()
        try:
            del self.channels[channel]
        finally:
            self.lock.release()

    def _predict_server(self):
        server = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
        server.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
        server.bind(self.bind_address)
        server.listen(1)
        support.log(self, 'Listening to {}...', self.bind_address)
        while True:
            try:
                connection, address = server.accept()
                worker = threading.Thread(daemon=True,
                                          target=self._predict_client,
                                          args=(connection, address))
                worker.start()
            except Exception as e:
                support.log(self, 'Encountered a problem ({}).', e)

class Saver:
    def __init__(self, config):
        self.backend = tf.train.Saver()
        self.path = config.save_path

    def save(self, session):
        path = self.backend.save(session, self.path)
        support.log(self, 'Saved in "{}".', path)

    def restore(self, session):
        if len(glob.glob("{}*".format(self.path))) > 0:
            if input('Restore from "{}"? '.format(self.path)) != 'no':
                self.backend.restore(session, self.path)
                support.log(self, 'Restored. Continue learning...')

class State:
    def deserialize(state):
        return State(state[0], state[1])

    def __init__(self, time, epoch):
        self.time = time
        self.epoch = epoch
        self.sample = 0

    def increment_epoch(self):
        self.epoch += 1
        self.sample = 0

    def increment_time(self):
        self.time += 1
        self.sample += 1

    def serialize(self):
        return [self.time, self.epoch]

class Target:
    def __init__(self, config):
        self.database = DistributedDatabase(config.data_path)
        self.dimension_count = 1
        self.sample_count = self.database.count()

    def compute(self, k):
        raise Exception('not implemented yet')

class TestTarget:
    def __init__(self, config):
        self.dimension_count = 1
        self.sample_count = 100000

    def compute(self, k):
        return np.reshape(np.sin(4 * np.pi / 40 * np.arange(0, 40)), [-1, 1])

def main(config):
    learn = Learn(config)
    target = Target(config)
    monitor = Monitor(config)
    learn.run(target, monitor, config)

if __name__ == '__main__':
    logging.basicConfig(format='%(asctime)s %(message)s',
                        datefmt='%Y-%m-%d %H:%M:%S',
                        level=logging.INFO)
    assert(len(sys.argv) == 2)
    config = Config({
        'dimension_count': 1,
        'data_path': sys.argv[1],
        'layer_count': 1,
        'unit_count': 200,
        'cell_clip': 1.0,
        'forget_bias': 1.0,
        'use_peepholes': True,
        'network_initializer': tf.random_uniform_initializer(-0.01, 0.01),
        'regression_initializer': tf.random_normal_initializer(stddev=0.01),
        'learning_rate': 1e-3,
        'gradient_clip': 1.0,
        'epoch_count': 100,
        'log_path': os.path.join('output', 'log'),
        'save_path': os.path.join('output', 'model'),
        'bind_address': ('0.0.0.0', 4242),
        'work_schedule': np.cumsum([1000 - 10, 10]),
    })
    main(config)
<|endoftext|>