<commit_before>import os
import glob
import pickle
import datetime

import numpy as np

from keras.layers import (Conv3D, AveragePooling3D, MaxPooling3D, Activation, UpSampling3D, merge, Input)
from keras import backend as K
from keras.models import Model,load_model
from keras.optimizers import Adam

import SimpleITK as sitk

pool_size = (2, 2, 2)
image_shape = (144, 240, 240)
n_channels = 3
input_shape = tuple([n_channels] + list(image_shape))
n_labels = 5
batch_size = 1
n_test_subjects = 40
z_crop = 155 - image_shape[0]

def pickle_dump(item, out_file):
    with open(out_file, "wb") as opened_file:
        pickle.dump(item, opened_file)

def pickle_load(in_file):
    with open(in_file, "rb") as opened_file:
        return pickle.load(opened_file)

K.set_image_dim_ordering('th')
smooth = 1.

def dice_coef(y_true, y_pred):
    y_true_f = K.flatten(y_true)
    y_pred_f = K.flatten(y_pred)
    intersection = K.sum(y_true_f * y_pred_f)
    return (2. * intersection + smooth) / (K.sum(y_true_f) + K.sum(y_pred_f) + smooth)

def dice_coef_loss(y_true, y_pred):
    return -dice_coef(y_true, y_pred)

def unet_model():
    inputs = Input(input_shape)
    conv1 = Conv3D(32, 3, 3, 3, activation='relu', border_mode='same')(inputs)
    conv1 = Conv3D(32, 3, 3, 3, activation='relu', border_mode='same')(conv1)
    pool1 = MaxPooling3D(pool_size=pool_size)(conv1)

    conv2 = Conv3D(64, 3, 3, 3, activation='relu', border_mode='same')(pool1)
    conv2 = Conv3D(64, 3, 3, 3, activation='relu', border_mode='same')(conv2)
    pool2 = MaxPooling3D(pool_size=pool_size)(conv2)

    conv3 = Conv3D(128, 3, 3, 3, activation='relu', border_mode='same')(pool2)
    conv3 = Conv3D(128, 3, 3, 3, activation='relu', border_mode='same')(conv3)
    pool3 = MaxPooling3D(pool_size=pool_size)(conv3)

    conv4 = Conv3D(256, 3, 3, 3, activation='relu', border_mode='same')(pool3)
    conv4 = Conv3D(256, 3, 3, 3, activation='relu', border_mode='same')(conv4)
    pool4 = MaxPooling3D(pool_size=pool_size)(conv4)

    conv5 = Conv3D(512, 3, 3, 3, activation='relu', border_mode='same')(pool4)
    conv5 = Conv3D(512, 3, 3, 3, activation='relu', border_mode='same')(conv5)

    up6 = merge([UpSampling3D(size=pool_size)(conv5), conv4], mode='concat', concat_axis=1)
    conv6 = Conv3D(256, 3, 3, 3, activation='relu', border_mode='same')(up6)
    conv6 = Conv3D(256, 3, 3, 3, activation='relu', border_mode='same')(conv6)

    up7 = merge([UpSampling3D(size=pool_size)(conv6), conv3], mode='concat', concat_axis=1)
    conv7 = Conv3D(128, 3, 3, 3, activation='relu', border_mode='same')(up7)
    conv7 = Conv3D(128, 3, 3, 3, activation='relu', border_mode='same')(conv7)

    up8 = merge([UpSampling3D(size=pool_size)(conv7), conv2], mode='concat', concat_axis=1)
    conv8 = Conv3D(64, 3, 3, 3, activation='relu', border_mode='same')(up8)
    conv8 = Conv3D(64, 3, 3, 3, activation='relu', border_mode='same')(conv8)

    up9 = merge([UpSampling3D(size=pool_size)(conv8), conv1], mode='concat', concat_axis=1)
    conv9 = Conv3D(32, 3, 3, 3, activation='relu', border_mode='same')(up9)
    conv9 = Conv3D(32, 3, 3, 3, activation='relu', border_mode='same')(conv9)

    conv10 = Conv3D(n_labels, 1, 1, 1, activation='sigmoid')(conv9)

    model = Model(input=inputs, output=conv10)

    model.compile(optimizer=Adam(lr=1e-5), loss=dice_coef_loss, metrics=[dice_coef])

    return model

def train_batch(batch, model):
    x_train = batch[:,:3]
    y_train = get_truth(batch)
    del(batch)
    model.train_on_batch(x_train, y_train)
    del(x_train)
    del(y_train)

def read_subject_folder(folder):
    flair_image = sitk.ReadImage(os.path.join(folder, "Flair.nii.gz"))
    t1_image = sitk.ReadImage(os.path.join(folder, "T1.nii.gz"))
    t1c_image = sitk.ReadImage(os.path.join(folder, "T1c.nii.gz"))
    truth_image = sitk.ReadImage(os.path.join(folder, "truth.nii.gz"))
    background_image = sitk.ReadImage(os.path.join(folder, "background.nii.gz"))
    return np.array([sitk.GetArrayFromImage(t1_image),
                     sitk.GetArrayFromImage(t1c_image),
                     sitk.GetArrayFromImage(flair_image),
                     sitk.GetArrayFromImage(truth_image),
                     sitk.GetArrayFromImage(background_image)])

def crop_data(data, background_channel=4):
    if np.all(data[background_channel, :z_crop] == 1):
        return data[:, z_crop:]
    elif np.all(data[background_channel, data.shape[1] - z_crop:] == 1):
        return data[:, :data.shape[1] - z_crop]
    else:
        upper = z_crop/2
        lower = z_crop - upper
        return data[:, lower:data.shape[1] - upper]

def get_truth(batch, truth_channel=3):
    truth = np.array(batch)[:, truth_channel]
    batch_list = []
    for sample_number in range(truth.shape[0]):
        sample_list = []
        for label in range(1, n_labels+1):
            array = np.zeros_like(truth[sample_number])
            array[truth[sample_number] == label] = 1
            sample_list.append(array)
        batch_list.append(sample_list)
    return np.array(batch_list)

def get_subject_id(subject_dir):
    return subject_dir.split("_")[-2]

def main(overwrite=False):
    model_file = os.path.abspath("3d_unet_model.h5")
    if not overwrite and os.path.exists(model_file):
        model = load_model(model_file)
    else:
        model = unet_model()

    processed_list_file = os.path.abspath("processed_subjects.pkl")
    if overwrite or not os.path.exists(processed_list_file):
        processed_list = []
    else:
        processed_list = pickle_load(processed_list_file)

    subject_dirs = glob.glob("../data/*/*")

    if os.path.exists("testing_ids.pkl") and not overwrite:
        testing_ids = pickle.load("testing_ids.pkl")
    else:
        # reomove duplicate sessions
        subjects = dict()
        for dirname in subject_dirs:
            subjects[dirname.split('_')[-2]] = dirname

        subject_ids = subjects.keys()
        np.random.shuffle(subject_ids)
        testing_ids = subject_ids[n_test_subjects:]
        pickle_dump(testing_ids, "testing_ids.pkl")

    batch = []
    for subject_dir in subject_dirs:

        subject_id = get_subject_id(subject_dir)
        if subject_id in testing_ids or subject_id in processed_list:
            continue

        processed_list.append(subject_id)

        batch.append(crop_data(read_subject_folder(subject_dir)))
        if len(batch) >= batch_size:
            train_batch(np.array(batch), model)
            del(batch)
            batch = []
            print("Saving: " + model_file)
            pickle_dump(processed_list, processed_list_file)
            model.save(model_file)

if __name__ == "__main__":
    main(overwrite=True)
<commit_msg>FIX: Allows for the model to load and skip the training on that has already been completed<commit_after>import os
import glob
import pickle
import datetime

import numpy as np

from keras.layers import (Conv3D, AveragePooling3D, MaxPooling3D, Activation, UpSampling3D, merge, Input)
from keras import backend as K
from keras.models import Model,load_model
from keras.optimizers import Adam

import SimpleITK as sitk

pool_size = (2, 2, 2)
image_shape = (144, 240, 240)
n_channels = 3
input_shape = tuple([n_channels] + list(image_shape))
n_labels = 5
batch_size = 1
n_test_subjects = 40
z_crop = 155 - image_shape[0]

def pickle_dump(item, out_file):
    with open(out_file, "wb") as opened_file:
        pickle.dump(item, opened_file)

def pickle_load(in_file):
    with open(in_file, "rb") as opened_file:
        return pickle.load(opened_file)

K.set_image_dim_ordering('th')
smooth = 1.

def dice_coef(y_true, y_pred):
    y_true_f = K.flatten(y_true)
    y_pred_f = K.flatten(y_pred)
    intersection = K.sum(y_true_f * y_pred_f)
    return (2. * intersection + smooth) / (K.sum(y_true_f) + K.sum(y_pred_f) + smooth)

def dice_coef_loss(y_true, y_pred):
    return -dice_coef(y_true, y_pred)

def unet_model():
    inputs = Input(input_shape)
    conv1 = Conv3D(32, 3, 3, 3, activation='relu', border_mode='same')(inputs)
    conv1 = Conv3D(32, 3, 3, 3, activation='relu', border_mode='same')(conv1)
    pool1 = MaxPooling3D(pool_size=pool_size)(conv1)

    conv2 = Conv3D(64, 3, 3, 3, activation='relu', border_mode='same')(pool1)
    conv2 = Conv3D(64, 3, 3, 3, activation='relu', border_mode='same')(conv2)
    pool2 = MaxPooling3D(pool_size=pool_size)(conv2)

    conv3 = Conv3D(128, 3, 3, 3, activation='relu', border_mode='same')(pool2)
    conv3 = Conv3D(128, 3, 3, 3, activation='relu', border_mode='same')(conv3)
    pool3 = MaxPooling3D(pool_size=pool_size)(conv3)

    conv4 = Conv3D(256, 3, 3, 3, activation='relu', border_mode='same')(pool3)
    conv4 = Conv3D(256, 3, 3, 3, activation='relu', border_mode='same')(conv4)
    pool4 = MaxPooling3D(pool_size=pool_size)(conv4)

    conv5 = Conv3D(512, 3, 3, 3, activation='relu', border_mode='same')(pool4)
    conv5 = Conv3D(512, 3, 3, 3, activation='relu', border_mode='same')(conv5)

    up6 = merge([UpSampling3D(size=pool_size)(conv5), conv4], mode='concat', concat_axis=1)
    conv6 = Conv3D(256, 3, 3, 3, activation='relu', border_mode='same')(up6)
    conv6 = Conv3D(256, 3, 3, 3, activation='relu', border_mode='same')(conv6)

    up7 = merge([UpSampling3D(size=pool_size)(conv6), conv3], mode='concat', concat_axis=1)
    conv7 = Conv3D(128, 3, 3, 3, activation='relu', border_mode='same')(up7)
    conv7 = Conv3D(128, 3, 3, 3, activation='relu', border_mode='same')(conv7)

    up8 = merge([UpSampling3D(size=pool_size)(conv7), conv2], mode='concat', concat_axis=1)
    conv8 = Conv3D(64, 3, 3, 3, activation='relu', border_mode='same')(up8)
    conv8 = Conv3D(64, 3, 3, 3, activation='relu', border_mode='same')(conv8)

    up9 = merge([UpSampling3D(size=pool_size)(conv8), conv1], mode='concat', concat_axis=1)
    conv9 = Conv3D(32, 3, 3, 3, activation='relu', border_mode='same')(up9)
    conv9 = Conv3D(32, 3, 3, 3, activation='relu', border_mode='same')(conv9)

    conv10 = Conv3D(n_labels, 1, 1, 1, activation='sigmoid')(conv9)

    model = Model(input=inputs, output=conv10)

    model.compile(optimizer=Adam(lr=1e-5), loss=dice_coef_loss, metrics=[dice_coef])

    return model

def train_batch(batch, model):
    x_train = batch[:,:3]
    y_train = get_truth(batch)
    del(batch)
    model.train_on_batch(x_train, y_train)
    del(x_train)
    del(y_train)

def read_subject_folder(folder):
    flair_image = sitk.ReadImage(os.path.join(folder, "Flair.nii.gz"))
    t1_image = sitk.ReadImage(os.path.join(folder, "T1.nii.gz"))
    t1c_image = sitk.ReadImage(os.path.join(folder, "T1c.nii.gz"))
    truth_image = sitk.ReadImage(os.path.join(folder, "truth.nii.gz"))
    background_image = sitk.ReadImage(os.path.join(folder, "background.nii.gz"))
    return np.array([sitk.GetArrayFromImage(t1_image),
                     sitk.GetArrayFromImage(t1c_image),
                     sitk.GetArrayFromImage(flair_image),
                     sitk.GetArrayFromImage(truth_image),
                     sitk.GetArrayFromImage(background_image)])

def crop_data(data, background_channel=4):
    if np.all(data[background_channel, :z_crop] == 1):
        return data[:, z_crop:]
    elif np.all(data[background_channel, data.shape[1] - z_crop:] == 1):
        return data[:, :data.shape[1] - z_crop]
    else:
        upper = z_crop/2
        lower = z_crop - upper
        return data[:, lower:data.shape[1] - upper]

def get_truth(batch, truth_channel=3):
    truth = np.array(batch)[:, truth_channel]
    batch_list = []
    for sample_number in range(truth.shape[0]):
        sample_list = []
        for label in range(1, n_labels+1):
            array = np.zeros_like(truth[sample_number])
            array[truth[sample_number] == label] = 1
            sample_list.append(array)
        batch_list.append(sample_list)
    return np.array(batch_list)

def get_subject_id(subject_dir):
    return subject_dir.split("_")[-2]

def main(overwrite=False):
    model_file = os.path.abspath("3d_unet_model.h5")
    if not overwrite and os.path.exists(model_file):
        model = load_model(model_file, custom_objects={'dice_coef_loss': dice_coef_loss, 'dice_coef': dice_coef})
    else:
        model = unet_model()

    processed_list_file = os.path.abspath("processed_subjects.pkl")
    if overwrite or not os.path.exists(processed_list_file):
        processed_list = []
    else:
        processed_list = pickle_load(processed_list_file)

    subject_dirs = glob.glob("../data/*/*")

    if os.path.exists("testing_ids.pkl") and not overwrite:
        testing_ids = pickle_load("testing_ids.pkl")
    else:
        # reomove duplicate sessions
        subjects = dict()
        for dirname in subject_dirs:
            subjects[dirname.split('_')[-2]] = dirname

        subject_ids = subjects.keys()
        np.random.shuffle(subject_ids)
        testing_ids = subject_ids[n_test_subjects:]
        pickle_dump(testing_ids, "testing_ids.pkl")

    batch = []
    for subject_dir in subject_dirs:

        subject_id = get_subject_id(subject_dir)
        if subject_id in testing_ids or subject_id in processed_list:
            continue

        processed_list.append(subject_id)

        batch.append(crop_data(read_subject_folder(subject_dir)))
        if len(batch) >= batch_size:
            train_batch(np.array(batch), model)
            del(batch)
            batch = []
            print("Saving: " + model_file)
            pickle_dump(processed_list, processed_list_file)
            model.save(model_file)

if __name__ == "__main__":
    main(overwrite=False)
<|endoftext|>


<commit_before>#! /usr/bin/python2
# vim: set fileencoding=utf-8
"""Try to describe venue by various features."""
from matplotlib import pyplot as pp
from collections import Counter, defaultdict, OrderedDict
import CommonMongo as cm
import FSCategories as fsc
import explore as xp
import numpy as np
import pandas as pd
import utils as u
import random as r
import scipy.cluster.vq as cluster
DB = None
CLIENT = None

def venues_info(vids, visits=None, visitors=None, depth=10):
    """Return various info about from the venue ids `vids`."""
    tags = defaultdict(int)
    city = DB.venue.find_one({'_id': vids[0]})['city']
    if not visits:
        visits = xp.get_visits(CLIENT, xp.Entity.venue, city)
    if not visitors:
        visitors = xp.get_visitors(CLIENT, city)
    venues = list(DB.venue.find({'_id': {'$in': vids}},
                                {'cat': 1, 'name': 1, 'loc': 1,
                                 'price': 1, 'rating': 1, 'tags': 1,
                                 'likes': 1, 'usersCount': 1,
                                 'checkinsCount': 1}))

    res = pd.DataFrame(index=[_['_id'] for _ in venues])

    def add_col(field):
        res[field.replace('Count', '')] = [_[field] for _ in venues]
    for field in ['name', 'price', 'rating', 'likes',
                  'usersCount', 'checkinsCount']:
        add_col(field)
    res['tags'] = [[normalized_tag(t) for t in _['tags']] for _ in venues]
    res['loc'] = [_['loc']['coordinates'] for _ in venues]
    res['cat'] = [parenting_cat(_['cat'], depth) for _ in venues]
    res['vis'] = [len(visits[id_]) for id_ in res.index]
    res['H'] = [venue_entropy(visitors[id_]) for id_ in res.index]
    for venue in venues:
        for tag in venue['tags']:
            tags[normalized_tag(tag)] += 1
    return res, OrderedDict(sorted(tags.iteritems(), key=lambda x: x[1],
                                   reverse=True))

def venue_entropy(visitors):
    """Compute the entropy of venue given the list of its `visitors`."""
    # pylint: disable=E1101
    c = np.array(Counter(visitors).values(), dtype=float)
    N = np.sum(c)
    return np.log(N) - np.sum(c*np.log(c))/N

def normalized_tag(tag):
    """normalize `tag` by removing punctuation and space character."""
    return tag.replace(' ', '')

def count_tags(tags):
    """Count occurence of a list of list of tags."""
    return Counter([normalized_tag(t) for oneset in tags for t in oneset])

def parenting_cat(cat, depth=1):
    """Return the name of category id `cat` (or name), stopping at level
    `depth`."""
    if not cat:
        return None
    _, path = fsc.search_categories(cat)
    cat_is_name = fsc.choose_type(cat)
    answer = path[depth] if len(path) > depth else path[-1]
    if cat_is_name:
        return answer
    return fsc.CAT_TO_ID[:answer]

def get_loc(vid):
    """Return coordinated of the venue `vid` (or None if it's not in DB)."""
    res = DB.venue.find_one({'_id': vid}, {'loc': 1})
    if res:
        return u.get_nested(res, ['loc', 'coordinates'])
    return None

def get_venue(vid, depth=1):
    """Return a textual description of venue `vid` or None."""
    venue = DB.venue.find_one({'_id': vid}, {'cat': 1, 'name': 1})
    if not venue:
        return None
    cat = parenting_cat(venue.get('cat'), depth)
    venue['cat'] = cat or '???'
    return (venue['cat'], venue['name'], vid)

def photos_around(id_, centroid, offset, daily, radius=200):
    """Gather photos timestamp in a `radius` around `id_` and return its time
    pattern (`daily` or not), and its distance to every `centroid`."""
    center = get_loc(id_)
    photos = xp.get_visits(CLIENT, xp.Entity.photo, ball=(center, radius))
    kind = xp.to_frequency(xp.aggregate_visits(photos.values(), offset)[daily])
    nb_class = centroid.shape[0]
    # pylint: disable=E1101
    classes = np.linalg.norm(np.tile(kind, (nb_class, 1)) - centroid, axis=1)
    return len(photos), kind, classes, np.argmin(classes)

def named_ticks(kind, offset=0):
    """Return ticks label for kind in ('day', 'week', 'mix')."""
    if kind is 'day':
        period = lambda i: '{}--{}'.format(i % 24, (i+3) % 24)
        return [period(i) for i in range(0+offset, 24+offset, 3)]
    days = 'mon tue wed thu fri sat sun'.split()
    if kind is 'week':
        return days
    if kind is 'mix':
        period = '1 2 3'.split()
        return [d+''+p for d in days for p in period]
    raise ValueError('`kind` arguments is not valid')

def draw_classes(centroid, offset):
    """Plot each time patterns in `centroid`."""
    size = centroid.shape[0]
    for i, marker in zip(range(size), legend[:size]):
        pp.plot(centroid[i, :], marker+'-', ms=11)
    if centroid.shape[1] == 8:
        pp.xticks(range(8), named_ticks('day', offset))
    else:
        pp.xticks(range(7*3), named_ticks('mix'))

if __name__ == '__main__':
    #pylint: disable=C0103
    import arguments
    args = arguments.city_parser().parse_args()
    city = args.city
    DB, CLIENT = cm.connect_to_db('foursquare', args.host, args.port)

    legend = 'v^<>s*xo|8d+'

    # pylint: disable=E1101
    def get_distorsion(ak, kl, sval):
        """Compute the sum of euclidean distance from `sval` to its
        centroid"""
        return np.sum(np.linalg.norm(ak[kl, :] - sval, axis=1))

    do_cluster = lambda val, k: cluster.kmeans2(val, k, 20, minit='points')

    def getclass(c, kl, visits):
        """Return {id: number of visits} of the venues in classs `c` of
        `kl`."""
        return {v[0]: v[1] for v, k in zip(visits.iteritems(), kl) if k == c}

    def peek_at_class(c, kl, visits, k=15):
        """Return a table of `k` randomly chosen venues in class `c` of
        `kl`."""
        sample = r.sample([get_venue(i)
                           for i in getclass(c, kl, visits).keys()], k)
        return pd.DataFrame({'cat': [_[0] for _ in sample],
                             'name': [_[1] for _ in sample],
                             'id': [_[2] for _ in sample]})
<commit_msg>normalize tag properly<commit_after>#! /usr/bin/python2
# vim: set fileencoding=utf-8
"""Try to describe venue by various features."""
from matplotlib import pyplot as pp
from collections import Counter, defaultdict, OrderedDict
import CommonMongo as cm
import FSCategories as fsc
import explore as xp
import numpy as np
import pandas as pd
import utils as u
import random as r
import scipy.cluster.vq as cluster
import re
import string
NOISE = re.compile(r'[\s'+string.punctuation+r']')
DB = None
CLIENT = None

def venues_info(vids, visits=None, visitors=None, depth=10):
    """Return various info about from the venue ids `vids`."""
    tags = defaultdict(int)
    city = DB.venue.find_one({'_id': vids[0]})['city']
    if not visits:
        visits = xp.get_visits(CLIENT, xp.Entity.venue, city)
    if not visitors:
        visitors = xp.get_visitors(CLIENT, city)
    venues = list(DB.venue.find({'_id': {'$in': vids}},
                                {'cat': 1, 'name': 1, 'loc': 1,
                                 'price': 1, 'rating': 1, 'tags': 1,
                                 'likes': 1, 'usersCount': 1,
                                 'checkinsCount': 1}))

    res = pd.DataFrame(index=[_['_id'] for _ in venues])

    def add_col(field):
        res[field.replace('Count', '')] = [_[field] for _ in venues]
    for field in ['name', 'price', 'rating', 'likes',
                  'usersCount', 'checkinsCount']:
        add_col(field)
    res['tags'] = [[normalized_tag(t) for t in _['tags']] for _ in venues]
    res['loc'] = [_['loc']['coordinates'] for _ in venues]
    res['cat'] = [parenting_cat(_['cat'], depth) for _ in venues]
    res['vis'] = [len(visits[id_]) for id_ in res.index]
    res['H'] = [venue_entropy(visitors[id_]) for id_ in res.index]
    for venue in venues:
        for tag in venue['tags']:
            tags[normalized_tag(tag)] += 1
    return res, OrderedDict(sorted(tags.iteritems(), key=lambda x: x[1],
                                   reverse=True))

def venue_entropy(visitors):
    """Compute the entropy of venue given the list of its `visitors`."""
    # pylint: disable=E1101
    c = np.array(Counter(visitors).values(), dtype=float)
    N = np.sum(c)
    return np.log(N) - np.sum(c*np.log(c))/N

def normalized_tag(tag):
    """normalize `tag` by removing punctuation and space character."""
    return NOISE.sub('', tag).lower()

def count_tags(tags):
    """Count occurence of a list of list of tags."""
    return Counter([normalized_tag(t) for oneset in tags for t in oneset])

def parenting_cat(cat, depth=1):
    """Return the name of category id `cat` (or name), stopping at level
    `depth`."""
    if not cat:
        return None
    _, path = fsc.search_categories(cat)
    cat_is_name = fsc.choose_type(cat)
    answer = path[depth] if len(path) > depth else path[-1]
    if cat_is_name:
        return answer
    return fsc.CAT_TO_ID[:answer]

def get_loc(vid):
    """Return coordinated of the venue `vid` (or None if it's not in DB)."""
    res = DB.venue.find_one({'_id': vid}, {'loc': 1})
    if res:
        return u.get_nested(res, ['loc', 'coordinates'])
    return None

def get_venue(vid, depth=1):
    """Return a textual description of venue `vid` or None."""
    venue = DB.venue.find_one({'_id': vid}, {'cat': 1, 'name': 1})
    if not venue:
        return None
    cat = parenting_cat(venue.get('cat'), depth)
    venue['cat'] = cat or '???'
    return (venue['cat'], venue['name'], vid)

def photos_around(id_, centroid, offset, daily, radius=200):
    """Gather photos timestamp in a `radius` around `id_` and return its time
    pattern (`daily` or not), and its distance to every `centroid`."""
    center = get_loc(id_)
    photos = xp.get_visits(CLIENT, xp.Entity.photo, ball=(center, radius))
    kind = xp.to_frequency(xp.aggregate_visits(photos.values(), offset)[daily])
    nb_class = centroid.shape[0]
    # pylint: disable=E1101
    classes = np.linalg.norm(np.tile(kind, (nb_class, 1)) - centroid, axis=1)
    return len(photos), kind, classes, np.argmin(classes)

def named_ticks(kind, offset=0):
    """Return ticks label for kind in ('day', 'week', 'mix')."""
    if kind is 'day':
        period = lambda i: '{}--{}'.format(i % 24, (i+3) % 24)
        return [period(i) for i in range(0+offset, 24+offset, 3)]
    days = 'mon tue wed thu fri sat sun'.split()
    if kind is 'week':
        return days
    if kind is 'mix':
        period = '1 2 3'.split()
        return [d+''+p for d in days for p in period]
    raise ValueError('`kind` arguments is not valid')

def draw_classes(centroid, offset):
    """Plot each time patterns in `centroid`."""
    size = centroid.shape[0]
    for i, marker in zip(range(size), legend[:size]):
        pp.plot(centroid[i, :], marker+'-', ms=11)
    if centroid.shape[1] == 8:
        pp.xticks(range(8), named_ticks('day', offset))
    else:
        pp.xticks(range(7*3), named_ticks('mix'))

if __name__ == '__main__':
    #pylint: disable=C0103
    import arguments
    args = arguments.city_parser().parse_args()
    city = args.city
    DB, CLIENT = cm.connect_to_db('foursquare', args.host, args.port)

    legend = 'v^<>s*xo|8d+'

    # pylint: disable=E1101
    def get_distorsion(ak, kl, sval):
        """Compute the sum of euclidean distance from `sval` to its
        centroid"""
        return np.sum(np.linalg.norm(ak[kl, :] - sval, axis=1))

    do_cluster = lambda val, k: cluster.kmeans2(val, k, 20, minit='points')

    def getclass(c, kl, visits):
        """Return {id: time pattern} of the venues in classs `c` of
        `kl`."""
        return {v[0]: v[1] for v, k in zip(visits.iteritems(), kl) if k == c}

    def peek_at_class(c, kl, visits, k=15):
        """Return a table of `k` randomly chosen venues in class `c` of
        `kl`."""
        sample = r.sample([get_venue(i)
                           for i in getclass(c, kl, visits).keys()], k)
        return pd.DataFrame({'cat': [_[0] for _ in sample],
                             'name': [_[1] for _ in sample],
                             'id': [_[2] for _ in sample]})
<|endoftext|>


<commit_before>#!/usr/bin/env python
# -*- coding: utf-8 -*-

"""
:mod:`utils` --- Utilities
==========================
"""

__docformat__ = "restructuredtext en"

import os
import sys
import re
import time

def camelcase_to_underscore(string):
    """
    Converts CamelCased string to lower-cased string with underscores.

    :param string: a CamelCased string to convert
    :type string: `unicode`

    :return: converted underscored string
    :rtype: `unicode`

    >>> camelcase_to_underscore("ThisIsACamelCasedString")
    'this_is_a_camel_cased_string'
    """

    pattern = "(((?<=[a-z])[A-Z])|([A-Z](?![A-Z]|$)))"
    return re.sub(pattern, "_\\1", string).lower().strip("_")

def retry_on_exception(exception, tries=10, delay=3, backoff=1):
    """
    Decorator.

    Retries a function until it stops raising specified exception(s) or reaches
    a certain number of tries.

    :param exception: exception(s) to be risen by decorated function
    :type exception: Exception class or a tuple of Exceptions classes

    :param tries: number of tries
    :type tries: `int`

    :param delay: delay between tries in seconds
    :type delay: `int`

    :param backoff: factor by which the delay should lengthen after each
        failure
    :type backoff: `int`

    Usage:

    >>> # This useless function will be retried 2 times.
    >>> @retry_on_exception((Exception, IOError), tries=2, delay=1, backoff=0)
    ... def useless_function():
    ...     print "derp"
    ...     raise Exception
    ...
    >>> useless_function()
    derp
    derp
    """

    def wrapper(f):
        def f_retry(*args, **kwargs):
            mtries, mdelay = tries, delay
            while mtries > 0:
                try:
                    return f(*args, **kwargs)
                except exception:
                    time.sleep(mdelay)
                    mtries -= 1
                    mdelay *= backoff
            return
        return f_retry
    return wrapper

def get_current_file_path():
    """
    py2exe / py2app workaround. Returns current file path based on script
    frozen state, platform and current system encoding.

    :return: current file path
    :rtype: `unicode`
    """

    encoding = sys.getfilesystemencoding()
    frozen = hasattr(sys, "frozen")
    if not frozen:
        return unicode(__file__, encoding)
    if frozen in ("dll", "console_exe", "windows_exe"):
        return unicode(sys.executable, encoding)
    if frozen in ("macosx_app"):
        return unicode(os.environ["RESOURCEPATH"], encoding)

if __name__ == "__main__":
    import doctest
    doctest.testmod()
<commit_msg>Typo fixes<commit_after>#!/usr/bin/env python
# -*- coding: utf-8 -*-

"""
:mod:`utils` --- Utilities
==========================
"""

__docformat__ = "restructuredtext en"

import os
import sys
import re
import time

def camelcase_to_underscore(string):
    """
    Converts CamelCased string to lower-cased string with underscores.

    :param string: a CamelCased string to convert
    :type string: `unicode`

    :return: converted underscored string
    :rtype: `unicode`

    >>> camelcase_to_underscore("ThisIsACamelCasedString")
    'this_is_a_camel_cased_string'
    """

    pattern = "(((?<=[a-z])[A-Z])|([A-Z](?![A-Z]|$)))"
    return re.sub(pattern, "_\\1", string).lower().strip("_")

def retry_on_exception(exception, tries=10, delay=3, backoff=1):
    """
    Decorator.

    Retries a function until it stops raising specified exception(s) or reaches
    a certain number of tries.

    :param exception: exception(s) to be risen by decorated function
    :type exception: Exception class or a tuple of Exceptions classes

    :param tries: number of tries
    :type tries: `int`

    :param delay: delay between tries in seconds
    :type delay: `int`

    :param backoff: factor by which the delay should lengthen after each
        failure
    :type backoff: `int`

    Usage:

    >>> # The following useless function will be consecutively called 2 times.
    >>> @retry_on_exception((Exception, IOError), tries=2, delay=0, backoff=0)
    ... def useless_function():
    ...     print "derp"
    ...     raise Exception
    ...
    >>> useless_function()
    derp
    derp
    """

    def wrapper(f):
        def f_retry(*args, **kwargs):
            mtries, mdelay = tries, delay
            while mtries > 0:
                try:
                    return f(*args, **kwargs)
                except exception:
                    time.sleep(mdelay)
                    mtries -= 1
                    mdelay *= backoff
            return
        return f_retry
    return wrapper

def get_current_file_path():
    """
    py2exe / py2app workaround. Returns current file path based on script's
    frozen state, current system platform and encoding.

    :return: current file path
    :rtype: `unicode`
    """

    encoding = sys.getfilesystemencoding()
    frozen = hasattr(sys, "frozen")
    if not frozen:
        return unicode(__file__, encoding)
    if frozen in ("dll", "console_exe", "windows_exe"):
        return unicode(sys.executable, encoding)
    if frozen in ("macosx_app"):
        return unicode(os.environ["RESOURCEPATH"], encoding)

if __name__ == "__main__":
    import doctest
    doctest.testmod()
<|endoftext|>


<commit_before>from twisted.python import log
from twisted.python.log import logging
from twisted.internet.defer import inlineCallbacks, returnValue
from twisted.internet.task import LoopingCall
from twisted.internet import reactor

from vumi.service import Worker, Consumer, Publisher
from vumi.message import Message
from vumi.workers.smpp.client import EsmeTransceiverFactory, EsmeTransceiver

import json
import re

#import os
#os.environ['DJANGO_SETTINGS_MODULE'] = 'vumi.webapp.settings'
from vumi.webapp.api import models
from vumi.webapp.api import forms
from vumi.webapp.api import utils
from vumi.utils import *
import vumi.options

import urllib
import urllib2

from datetime import datetime, timedelta
import time
import redis

class SmppConsumer(Consumer):
    """
    This consumer creates the generic outbound SMPP transport.
    Anything published to the `vumi.smpp` exchange with
    routing key smpp.* (* == single word match, # == zero or more words)
    """
    exchange_name = "vumi"
    exchange_type = "direct"
    durable = True
    auto_delete = False
    queue_name = "sms.outbound.fallback"
    routing_key = "sms.outbound.fallback"

    def __init__(self, r_server, r_prefix, send_callback):
        self.r_server = r_server
        self.r_prefix = r_prefix
        self.send = send_callback
        log.msg("Consuming on %s -> %s" % (self.routing_key, self.queue_name))

    def consume_message(self, message):
        log.msg("Consumed JSON", message)
        sequence_number = self.send(**message.payload)
        formdict = {
                "sent_sms":message.payload.get("id"),
                "sequence_number": sequence_number,
                }
        log.msg("SMPPLinkForm", repr(formdict))
        form = forms.SMPPLinkForm(formdict)
        form.save()
        self.r_server.set("%s_%s#last_sequence_number" % (self.r_prefix, vumi.options.get_all()['config']['smpp_increment']),
                sequence_number)
        self.r_server.set("%s#%s" % (self.r_prefix, sequence_number),
                message.payload.get("id"))
        return True

    def consume(self, message):
        if self.consume_message(Message.from_json(message.content.body)):
            self.ack(message)

def dynamically_create_smpp_consumer(name, **kwargs):
    return type("%s_SmppConsumer" % name, (SmppConsumer,), kwargs)

class SmppPublisher(Publisher):
    """
    This publisher publishes all incoming SMPP messages to the
    `vumi.smpp` exchange, its default routing key is `smpp.fallback`
    """
    exchange_name = "vumi"
    exchange_type = "topic"             # -> route based on pattern matching
    routing_key = 'smpp.fallback'       # -> overriden in publish method
    durable = False                     # -> not created at boot
    auto_delete = False                 # -> auto delete if no consumers bound
    delivery_mode = 2                   # -> save to disk

    def publish_message(self, message, **kwargs):
        log.msg("Publishing Message %s with extra args: %s" % (message, kwargs))
        super(SmppPublisher, self).publish_message(message, **kwargs)

class SmppTransport(Worker):
    """
    The SmppTransport
    """

    def startWorker(self):
        # Connect to Redis
        self.r_server = redis.Redis("localhost", db=vumi.options.get_deploy_int(self.vhost))
        log.msg("Connected to Redis")
        config = vumi.options.get_all()['config']
        self.r_prefix = "%s@%s:%s" % (config['system_id'], config['host'], config['port'])
        log.msg("r_prefix = %s" % self.r_prefix)

        log.msg("Starting the SmppTransport")
        # start the Smpp transport
        factory = EsmeTransceiverFactory(
                int(self.config['smpp_increment']),
                int(self.config['smpp_offset']))
        factory.loadDefaults(self.config)
        self.sequence_key = "%s_%s#last_sequence_number" % (self.r_prefix, config['smpp_increment'])
        log.msg("sequence_key = %s", self.sequence_key)
        last_sequence_number = int(self.r_server.get(self.sequence_key) or 1551)
        factory.setLatestSequenceNumber(last_sequence_number)
        factory.setConnectCallback(self.esme_connected)
        factory.setDisconnectCallback(self.esme_disconnected)
        factory.setSubmitSMRespCallback(self.submit_sm_resp)
        factory.setDeliveryReportCallback(self.delivery_report)
        factory.setDeliverSMCallback(self.deliver_sm)
        factory.setLoopingQuerySMCallback(self.query_sm_group)
        log.msg(factory.defaults)
        reactor.connectTCP(
                factory.defaults['host'],
                factory.defaults['port'],
                factory)

    #def getLatestSequenceNumber(self):
        #sequence_number = 0
        #try:
            #sequence_number = models.SMPPLink.objects.latest().sequence_number
        #except Exception, e:
            #log.msg("No SMPPLink entries yet")
        #return sequence_number

    @inlineCallbacks
    def esme_connected(self, client):
        log.msg("ESME Connected, adding handlers")
        self.esme_client = client
        self.esme_client.set_handler(self)

        # Start the publisher
        self.publisher = yield self.start_publisher(SmppPublisher)
        # Start the consumer, pass along the send_smpp callback for sending
        # back consumed AMQP messages over SMPP.
        #self.consumer = yield self.start_consumer(SmppConsumer, self.send_smpp)
        transport = self.config.get('TRANSPORT_NAME', 'fallback').lower()
        yield self.start_consumer(dynamically_create_smpp_consumer(transport,
                    routing_key='sms.outbound.%s' % transport,
                    queue_name='sms.outbound.%s' % transport
                ), self.r_server, self.r_prefix, self.send_smpp)

    @inlineCallbacks
    def esme_disconnected(self):
        log.msg("ESME Disconnected, stopping consumer")
        stop = yield self.consumer.stop()

    @inlineCallbacks
    def submit_sm_resp(self, *args, **kwargs):
        smpplink = models.SMPPLink.objects \
                .filter(sequence_number=kwargs['sequence_number']) \
                .latest('created_at')
        redis_key = "%s#%s" % (self.r_prefix, kwargs['sequence_number'])
        log.msg("redis_key = %s" % redis_key)
        sent_sms_id = self.r_server.get(redis_key)
        self.r_server.delete(redis_key)
        kwargs.update({'sent_sms':sent_sms_id})
        log.msg("SMPPRespForm <- %s" % repr(kwargs))
        form = forms.SMPPRespForm(kwargs)
        form.save()
        yield log.msg("SUBMIT SM RESP %s" % repr(kwargs))

    @inlineCallbacks
    def query_sm_group(self, *args, **kwargs):
        try:
            self.second_counter += 1
            if self.second_counter >= 60:
                self.second_counter = 0
        except:
            self.second_counter = 0
        fromdate = datetime.now() - timedelta(days=1)
        smppRespList = models.SMPPResp.objects \
                .filter(created_at__gte=fromdate) \
                .extra(where=['ROUND(EXTRACT(SECOND FROM created_at)) = %d' % (self.second_counter)]) \
                .order_by('-created_at')
        for r in smppRespList:
            route = get_operator_number(
                    r.sent_sms.to_msisdn,
                    self.config['COUNTRY_CODE'],
                    self.config.get('OPERATOR_PREFIX',{}),
                    self.config.get('OPERATOR_NUMBER',{}))
            sequence_number = self.esme_client.query_sm(
                    message_id = r.message_id,
                    source_addr = route
                    )
        yield log.msg("LOOPING QUERY SM" % repr(kwargs))

    @inlineCallbacks
    def delivery_report(self, *args, **kwargs):
        transport_name = self.config.get('TRANSPORT_NAME', 'fallback').lower()
        log.msg("DELIVERY REPORT", kwargs)
        dictionary = {
            'transport_name': transport_name,
            'transport_msg_id': kwargs['delivery_report']['id'],
            'transport_status': kwargs['delivery_report']['stat'],
            'transport_delivered_at': datetime.strptime(
                kwargs['delivery_report']['done_date'],
                "%y%m%d%H%M%S")
        }
        yield self.publisher.publish_message(Message(**dictionary),
            routing_key='sms.receipt.%s' % (transport_name,))

    @inlineCallbacks
    def deliver_sm(self, *args, **kwargs):
        yield self.publisher.publish_message(Message(**kwargs),
            routing_key='sms.inbound.%s.%s' % (
                self.config.get('TRANSPORT_NAME', 'fallback').lower(),
                kwargs.get('destination_addr')))

    def send_smpp(self, id, to_msisdn, message, *args, **kwargs):
        log.msg("Sending SMPP, to: %s, message: %s" % (to_msisdn, repr(message)))
        # first do a lookup in our YAML to see if we've got a source_addr
        # defined for the given MT number, if not, trust the from_msisdn
        # in the message
        route = get_operator_number(to_msisdn,
                self.config.get('COUNTRY_CODE',''),
                self.config.get('OPERATOR_PREFIX',{}),
                self.config.get('OPERATOR_NUMBER',{})) or kwargs.get('from_msisdn', '')
        sequence_number = self.esme_client.submit_sm(
                short_message = message.encode('utf-8'),
                destination_addr = str(to_msisdn),
                source_addr = route,
                )
        #self.deliver_sm(
                #short_message=str(message),
                #destination_addr=str(to_msisdn),
                #source_addr=route)
        return sequence_number

    def stopWorker(self):
        log.msg("Stopping the SMPPTransport")

<commit_msg>finally - seeds smpp sequence numbers off redis<commit_after>from twisted.python import log
from twisted.python.log import logging
from twisted.internet.defer import inlineCallbacks, returnValue
from twisted.internet.task import LoopingCall
from twisted.internet import reactor

from vumi.service import Worker, Consumer, Publisher
from vumi.message import Message
from vumi.workers.smpp.client import EsmeTransceiverFactory, EsmeTransceiver

import json
import re

#import os
#os.environ['DJANGO_SETTINGS_MODULE'] = 'vumi.webapp.settings'
from vumi.webapp.api import models
from vumi.webapp.api import forms
from vumi.webapp.api import utils
from vumi.utils import *
import vumi.options

import urllib
import urllib2

from datetime import datetime, timedelta
import time
import redis

class SmppConsumer(Consumer):
    """
    This consumer creates the generic outbound SMPP transport.
    Anything published to the `vumi.smpp` exchange with
    routing key smpp.* (* == single word match, # == zero or more words)
    """
    exchange_name = "vumi"
    exchange_type = "direct"
    durable = True
    auto_delete = False
    queue_name = "sms.outbound.fallback"
    routing_key = "sms.outbound.fallback"

    def __init__(self, r_server, r_prefix, send_callback):
        self.r_server = r_server
        self.r_prefix = r_prefix
        self.send = send_callback
        log.msg("Consuming on %s -> %s" % (self.routing_key, self.queue_name))

    def consume_message(self, message):
        log.msg("Consumed JSON", message)
        sequence_number = self.send(**message.payload)
        formdict = {
                "sent_sms":message.payload.get("id"),
                "sequence_number": sequence_number,
                }
        log.msg("SMPPLinkForm", repr(formdict))
        form = forms.SMPPLinkForm(formdict)
        form.save()
        self.r_server.set("%s_%s#last_sequence_number" % (self.r_prefix, vumi.options.get_all()['config']['smpp_offset']),
                sequence_number)
        self.r_server.set("%s#%s" % (self.r_prefix, sequence_number),
                message.payload.get("id"))
        return True

    def consume(self, message):
        if self.consume_message(Message.from_json(message.content.body)):
            self.ack(message)

def dynamically_create_smpp_consumer(name, **kwargs):
    return type("%s_SmppConsumer" % name, (SmppConsumer,), kwargs)

class SmppPublisher(Publisher):
    """
    This publisher publishes all incoming SMPP messages to the
    `vumi.smpp` exchange, its default routing key is `smpp.fallback`
    """
    exchange_name = "vumi"
    exchange_type = "topic"             # -> route based on pattern matching
    routing_key = 'smpp.fallback'       # -> overriden in publish method
    durable = False                     # -> not created at boot
    auto_delete = False                 # -> auto delete if no consumers bound
    delivery_mode = 2                   # -> save to disk

    def publish_message(self, message, **kwargs):
        log.msg("Publishing Message %s with extra args: %s" % (message, kwargs))
        super(SmppPublisher, self).publish_message(message, **kwargs)

class SmppTransport(Worker):
    """
    The SmppTransport
    """

    def startWorker(self):
        # Connect to Redis
        self.r_server = redis.Redis("localhost", db=vumi.options.get_deploy_int(self.vhost))
        log.msg("Connected to Redis")
        config = vumi.options.get_all()['config']
        self.r_prefix = "%s@%s:%s" % (config['system_id'], config['host'], config['port'])
        log.msg("r_prefix = %s" % self.r_prefix)

        log.msg("Starting the SmppTransport")
        # start the Smpp transport
        factory = EsmeTransceiverFactory(
                int(self.config['smpp_increment']),
                int(self.config['smpp_offset']))
        factory.loadDefaults(self.config)
        self.sequence_key = "%s_%s#last_sequence_number" % (self.r_prefix, config['smpp_offset'])
        log.msg("sequence_key = %s" % (self.sequence_key))
        last_sequence_number = int(self.r_server.get(self.sequence_key) or 0)
        factory.setLatestSequenceNumber(last_sequence_number)
        factory.setConnectCallback(self.esme_connected)
        factory.setDisconnectCallback(self.esme_disconnected)
        factory.setSubmitSMRespCallback(self.submit_sm_resp)
        factory.setDeliveryReportCallback(self.delivery_report)
        factory.setDeliverSMCallback(self.deliver_sm)
        factory.setLoopingQuerySMCallback(self.query_sm_group)
        log.msg(factory.defaults)
        reactor.connectTCP(
                factory.defaults['host'],
                factory.defaults['port'],
                factory)

    #def getLatestSequenceNumber(self):
        #sequence_number = 0
        #try:
            #sequence_number = models.SMPPLink.objects.latest().sequence_number
        #except Exception, e:
            #log.msg("No SMPPLink entries yet")
        #return sequence_number

    @inlineCallbacks
    def esme_connected(self, client):
        log.msg("ESME Connected, adding handlers")
        self.esme_client = client
        self.esme_client.set_handler(self)

        # Start the publisher
        self.publisher = yield self.start_publisher(SmppPublisher)
        # Start the consumer, pass along the send_smpp callback for sending
        # back consumed AMQP messages over SMPP.
        #self.consumer = yield self.start_consumer(SmppConsumer, self.send_smpp)
        transport = self.config.get('TRANSPORT_NAME', 'fallback').lower()
        yield self.start_consumer(dynamically_create_smpp_consumer(transport,
                    routing_key='sms.outbound.%s' % transport,
                    queue_name='sms.outbound.%s' % transport
                ), self.r_server, self.r_prefix, self.send_smpp)

    @inlineCallbacks
    def esme_disconnected(self):
        log.msg("ESME Disconnected, stopping consumer")
        stop = yield self.consumer.stop()

    @inlineCallbacks
    def submit_sm_resp(self, *args, **kwargs):
        smpplink = models.SMPPLink.objects \
                .filter(sequence_number=kwargs['sequence_number']) \
                .latest('created_at')
        redis_key = "%s#%s" % (self.r_prefix, kwargs['sequence_number'])
        log.msg("redis_key = %s" % redis_key)
        sent_sms_id = self.r_server.get(redis_key)
        self.r_server.delete(redis_key)
        kwargs.update({'sent_sms':sent_sms_id})
        log.msg("SMPPRespForm <- %s" % repr(kwargs))
        form = forms.SMPPRespForm(kwargs)
        form.save()
        yield log.msg("SUBMIT SM RESP %s" % repr(kwargs))

    @inlineCallbacks
    def query_sm_group(self, *args, **kwargs):
        try:
            self.second_counter += 1
            if self.second_counter >= 60:
                self.second_counter = 0
        except:
            self.second_counter = 0
        fromdate = datetime.now() - timedelta(days=1)
        smppRespList = models.SMPPResp.objects \
                .filter(created_at__gte=fromdate) \
                .extra(where=['ROUND(EXTRACT(SECOND FROM created_at)) = %d' % (self.second_counter)]) \
                .order_by('-created_at')
        for r in smppRespList:
            route = get_operator_number(
                    r.sent_sms.to_msisdn,
                    self.config['COUNTRY_CODE'],
                    self.config.get('OPERATOR_PREFIX',{}),
                    self.config.get('OPERATOR_NUMBER',{}))
            sequence_number = self.esme_client.query_sm(
                    message_id = r.message_id,
                    source_addr = route
                    )
        yield log.msg("LOOPING QUERY SM" % repr(kwargs))

    @inlineCallbacks
    def delivery_report(self, *args, **kwargs):
        transport_name = self.config.get('TRANSPORT_NAME', 'fallback').lower()
        log.msg("DELIVERY REPORT", kwargs)
        dictionary = {
            'transport_name': transport_name,
            'transport_msg_id': kwargs['delivery_report']['id'],
            'transport_status': kwargs['delivery_report']['stat'],
            'transport_delivered_at': datetime.strptime(
                kwargs['delivery_report']['done_date'],
                "%y%m%d%H%M%S")
        }
        yield self.publisher.publish_message(Message(**dictionary),
            routing_key='sms.receipt.%s' % (transport_name,))

    @inlineCallbacks
    def deliver_sm(self, *args, **kwargs):
        yield self.publisher.publish_message(Message(**kwargs),
            routing_key='sms.inbound.%s.%s' % (
                self.config.get('TRANSPORT_NAME', 'fallback').lower(),
                kwargs.get('destination_addr')))

    def send_smpp(self, id, to_msisdn, message, *args, **kwargs):
        log.msg("Sending SMPP, to: %s, message: %s" % (to_msisdn, repr(message)))
        # first do a lookup in our YAML to see if we've got a source_addr
        # defined for the given MT number, if not, trust the from_msisdn
        # in the message
        route = get_operator_number(to_msisdn,
                self.config.get('COUNTRY_CODE',''),
                self.config.get('OPERATOR_PREFIX',{}),
                self.config.get('OPERATOR_NUMBER',{})) or kwargs.get('from_msisdn', '')
        sequence_number = self.esme_client.submit_sm(
                short_message = message.encode('utf-8'),
                destination_addr = str(to_msisdn),
                source_addr = route,
                )
        #self.deliver_sm(
                #short_message=str(message),
                #destination_addr=str(to_msisdn),
                #source_addr=route)
        return sequence_number

    def stopWorker(self):
        log.msg("Stopping the SMPPTransport")

<|endoftext|>


<commit_before>version_info = (2, 13, 0)
version = '.'.join(map(str, version_info))
<commit_msg>up version<commit_after>version_info = (2, 13, 1)
version = '.'.join(map(str, version_info))
<|endoftext|>


<commit_before><commit_msg>Added missing methods to cairowrapper.py<commit_after><|endoftext|>


<commit_before>import csv
import datetime

from subprocess import PIPE

from wal_e.piper import popen_nonblock
from wal_e.exception import UserException

PSQL_BIN = 'psql'

class UTC(datetime.tzinfo):
    """
    UTC timezone

    Adapted from a Python example

    """

    ZERO = datetime.timedelta(0)
    HOUR = datetime.timedelta(hours=1)

    def utcoffset(self, dt):
        return self.ZERO

    def tzname(self, dt):
        return "UTC"

    def dst(self, dt):
        return self.ZERO

def psql_csv_run(sql_command, error_handler=None):
    """
    Runs psql and returns a CSVReader object from the query

    This CSVReader includes header names as the first record in all
    situations.  The output is fully buffered into Python.

    """
    # Explicitly disable the query timeout.
    csv_query = ['SET statement_timeout = 0']
    csv_query.append('COPY ({}) TO STDOUT WITH CSV HEADER'.format(sql_command))

    psql_proc = popen_nonblock([PSQL_BIN, '-d', 'postgres', '--no-password',
                                '-c', '; '.join(csv_query)],
                               stdout=PIPE)
    stdout = psql_proc.communicate()[0]

    if psql_proc.returncode != 0:
        if error_handler is not None:
            error_handler(psql_proc)
        else:
            assert error_handler is None
            raise UserException(
                'could not csv-execute a query successfully via psql',
                'Query was "{query}".'.format(sql_command),
                'You may have to set some libpq environment '
                'variables if you are sure the server is running.')

    # Previous code must raise any desired exceptions for non-zero
    # exit codes
    assert psql_proc.returncode == 0

    # Fake enough iterator interface to get a CSV Reader object
    # that works.
    return csv.reader(iter(stdout.strip().split('\n')))

class PgBackupStatements(object):
    """
    Contains operators to start and stop a backup on a Postgres server

    Relies on PsqlHelp for underlying mechanism.

    """

    @staticmethod
    def _dict_transform(csv_reader):
        rows = list(csv_reader)
        assert len(rows) == 2, 'Expect header row and data row'
        return dict(zip(*rows))

    @classmethod
    def run_start_backup(cls):
        """
        Connects to a server and attempts to start a hot backup

        Yields the WAL information in a dictionary for bookkeeping and
        recording.

        """
        def handler(popen):
            assert popen.returncode != 0
            raise UserException('Could not start hot backup')

        # The difficulty of getting a timezone-stamped, UTC,
        # ISO-formatted datetime is downright embarrassing.
        #
        # See http://bugs.python.org/issue5094
        label = 'freeze_start_' + (datetime.datetime.utcnow()
                                   .replace(tzinfo=UTC()).isoformat())

        return cls._dict_transform(psql_csv_run(
                "SELECT file_name, "
                "  lpad(file_offset::text, 8, '0') AS file_offset "
                "FROM pg_xlogfile_name_offset("
                "  pg_start_backup('{0}'))".format(label),
                error_handler=handler))

    @classmethod
    def run_stop_backup(cls):
        """
        Stop a hot backup, if it was running, or error

        Return the last WAL file name and position that is required to
        gain consistency on the captured heap.

        """
        def handler(popen):
            assert popen.returncode != 0
            raise UserException('Could not stop hot backup')

        return cls._dict_transform(psql_csv_run(
                "SELECT file_name, "
                "  lpad(file_offset::text, 8, '0') AS file_offset "
                "FROM pg_xlogfile_name_offset("
                "  pg_stop_backup())", error_handler=handler))

    @classmethod
    def pg_version(cls):
        """
        Get a very informative version string from Postgres

        Includes minor version, major version, and architecture, among
        other details.

        """
        return cls._dict_transform(psql_csv_run('SELECT * FROM version()'))
<commit_msg>Revert "Disable query timeout for backup push."<commit_after>import csv
import datetime

from subprocess import PIPE

from wal_e.piper import popen_nonblock
from wal_e.exception import UserException

PSQL_BIN = 'psql'

class UTC(datetime.tzinfo):
    """
    UTC timezone

    Adapted from a Python example

    """

    ZERO = datetime.timedelta(0)
    HOUR = datetime.timedelta(hours=1)

    def utcoffset(self, dt):
        return self.ZERO

    def tzname(self, dt):
        return "UTC"

    def dst(self, dt):
        return self.ZERO

def psql_csv_run(sql_command, error_handler=None):
    """
    Runs psql and returns a CSVReader object from the query

    This CSVReader includes header names as the first record in all
    situations.  The output is fully buffered into Python.

    """
    csv_query = ('COPY ({query}) TO STDOUT WITH CSV HEADER;'
                 .format(query=sql_command))

    psql_proc = popen_nonblock([PSQL_BIN, '-d', 'postgres', '--no-password',
                                '-c', csv_query],
                               stdout=PIPE)
    stdout = psql_proc.communicate()[0]

    if psql_proc.returncode != 0:
        if error_handler is not None:
            error_handler(psql_proc)
        else:
            assert error_handler is None
            raise UserException(
                'could not csv-execute a query successfully via psql',
                'Query was "{query}".'.format(sql_command),
                'You may have to set some libpq environment '
                'variables if you are sure the server is running.')

    # Previous code must raise any desired exceptions for non-zero
    # exit codes
    assert psql_proc.returncode == 0

    # Fake enough iterator interface to get a CSV Reader object
    # that works.
    return csv.reader(iter(stdout.strip().split('\n')))

class PgBackupStatements(object):
    """
    Contains operators to start and stop a backup on a Postgres server

    Relies on PsqlHelp for underlying mechanism.

    """

    @staticmethod
    def _dict_transform(csv_reader):
        rows = list(csv_reader)
        assert len(rows) == 2, 'Expect header row and data row'
        return dict(zip(*rows))

    @classmethod
    def run_start_backup(cls):
        """
        Connects to a server and attempts to start a hot backup

        Yields the WAL information in a dictionary for bookkeeping and
        recording.

        """
        def handler(popen):
            assert popen.returncode != 0
            raise UserException('Could not start hot backup')

        # The difficulty of getting a timezone-stamped, UTC,
        # ISO-formatted datetime is downright embarrassing.
        #
        # See http://bugs.python.org/issue5094
        label = 'freeze_start_' + (datetime.datetime.utcnow()
                                   .replace(tzinfo=UTC()).isoformat())

        return cls._dict_transform(psql_csv_run(
                "SELECT file_name, "
                "  lpad(file_offset::text, 8, '0') AS file_offset "
                "FROM pg_xlogfile_name_offset("
                "  pg_start_backup('{0}'))".format(label),
                error_handler=handler))

    @classmethod
    def run_stop_backup(cls):
        """
        Stop a hot backup, if it was running, or error

        Return the last WAL file name and position that is required to
        gain consistency on the captured heap.

        """
        def handler(popen):
            assert popen.returncode != 0
            raise UserException('Could not stop hot backup')

        return cls._dict_transform(psql_csv_run(
                "SELECT file_name, "
                "  lpad(file_offset::text, 8, '0') AS file_offset "
                "FROM pg_xlogfile_name_offset("
                "  pg_stop_backup())", error_handler=handler))

    @classmethod
    def pg_version(cls):
        """
        Get a very informative version string from Postgres

        Includes minor version, major version, and architecture, among
        other details.

        """
        return cls._dict_transform(psql_csv_run('SELECT * FROM version()'))
<|endoftext|>


<commit_before>#!/usr/bin/python

import re
import time,logging
import argparse, string,datetime,sys
import afs
from afs.util.AfsConfig import parseDefaultConfig
from afs.util.afsutil import parseHumanWriteableSize
from afs.service.OSDVolService import OSDVolService
from afs.service.OSDCellService import OSDCellService
from afs.service.OSDFsService import OSDFsService
from afs.service.ProjectService import ProjectService
from afs.model.Volume import Volume
from afs.dao.VolumeDAO import VolumeDAO
from afs.dao.VLDbDAO import VLDbDAO

global FS,PS,VS,VD,VlD

myParser=argparse.ArgumentParser(parents=[afs.argParser], add_help=False)
myParser.add_argument("--ssrv", required=True, help="server to empty")
myParser.add_argument("--spart", help="partition on server to empty. defaults to all")
myParser.add_argument("--dsrv" , help="force server to be filled. Otherwise chosen automatically by project.")
myParser.add_argument("--dpart" , help="force partition on server to be filled. Otherwise chosen automatically by project.")
group = myParser.add_mutually_exclusive_group()
group.add_argument("--ignorerx", dest="ignoreRX", action="append",  help="regEx for volumenames to ignore. All volumes not matching this will be moved.")
group.add_argument("--onlyrx", dest="onlyRX", action="append", help="regEx for volumenames to include all volumes not matching this will be ignored.")
group.add_argument("--ignoreproject", dest="ignoreProjects", action="append", help="ignore volumes of given project.")
group.add_argument("--onlyproject", dest="onlyProjects", action="append", help="only move volumes of given project.")
myParser.add_argument("--dryrun",action="store_true", help="Just print out what would be done, but don't do it.")
myParser.add_argument("--maxnum", default = 0, type=int, help="max number of Volumes to move.")
myParser.add_argument("--untilfree", default = "0", help="move until # is free on spart.")
myParser.add_argument("--rwvols", dest="moveRWVols", default=False, action="store_true", help="move rwvols with their accompanying ROs.")
myParser.add_argument("--solitaryrovols", dest="moveSolitaryROVols", default=False, action="store_true", help="move solitary rovols.")
myParser.add_argument("--minsize", dest="minVolumeUsage", default="0", help="only move volumes with minimalsize of")
myParser.add_argument("--osdvolumes", dest="moveOSDVOlumes", default=False, action="store_true", help="also move OSD-Volumes")

parseDefaultConfig(myParser)
FS=OSDFsService()
PS=ProjectService()
VS=OSDVolService()
VD=VolumeDAO()
VlD=VLDbDAO()

if not afs.defaultConfig.moveRWVols and not afs.defaultConfig.moveSolitaryROVols :
    sys.stderr.write("If you want to nmake me do anything, specify --rwvols and/or --solitaryrovols\n")
    sys.exit(1)

if afs.defaultConfig.ignoreRX != None :
    ignoreRX=[]
    for rx in afs.defaultConfig.ignoreRX :
        try :
            ignoreRX.append(re.compile(rx))
        except :
            sys.stderr.write("Cannot compile regular expression: '%s'\n" % rx)
            sys.exit(1)
elif afs.defaultConfig.onlyRX != None :
    onlyRX=[]
    for rx in afs.defaultConfig.onlyRX :
        try :
            onlyRX.append(re.compile(rx))
        except :
            sys.stderr.write("Cannot compile regular expression: '%s'\n" % rx)
            sys.exit(1)
srcFS=FS.getFileServer(afs.defaultConfig.ssrv)

if srcFS == None :
    sys.stderr.write("src server %s does not exist." % srcFS)
    sys.exit(2)
if afs.defaultConfig.spart != None :
    if not afs.defaultConfig.spart in srcFS.parts :
        sys.stderr.write("Partition %s does not exist on server %s\n" % (afs.defaultConfig.spart, afs.defaultConfig.ssrv))
        sys.exit(2)
    # we cycle thorugh this list later
    srcParts=[afs.defaultConfig.spart]
else :
    srcParts=srcFS.parts

if afs.defaultConfig.dsrv != None :
    dstFS=FS.getFileServer(afs.defaultConfig.dsrv)
    if dstFS == None :
        sys.stderr.write("destination server %s does not exist.\n" % srcFS)
        sys.exit(2)

    if afs.defaultConfig.dpart != None :
        if not afs.defaultConfig.dpart in dstFS.parts :
            sys.stderr.write("Partition %s does not exist on server %s\n" % (afs.defaultConfig.dpart, afs.defaultConfig.dsrv))
            sys.exit(2)
        else :
            reqDstPart=afs.defaultConfig.dpart
    else :
        reqDstPart=None
else :
    dstFS=None
    if afs.defaultConfig.dpart != None :
        sys.stderr.write("Warning: ignoring given dpart=%s, because no dsrv has been specified.\n" % afs.defaultConfig.dpart)

# XXX we should handle everything internally with bytes
untilFree=parseHumanWriteableSize(afs.defaultConfig.untilfree)/1024
minUsage=parseHumanWriteableSize(afs.defaultConfig.minVolumeUsage)/1024
VolObj = Volume()
movedVolcount=0
for srcP in srcParts :
    print "Processing Partition %s...." % srcP
    # check if partition is freed enough
    parts=FS.getPartitions(afs.defaultConfig.ssrv)
    if parts[srcP]["free"] > untilFree and untilFree > 0 :
        print "already %s Bytes free on spart %s." % (afs.util.afsutil.humanReadableSize(parts[srcP]["free"]/1024),srcP )
        continue
    # get list of volumes to move
    srcVolList=FS.getVolList(srcFS.servernames[0],srcP,cached=False)
    # get RW Volumes :
    RWVols=[]
    solitaryROVols=[]
    for v in srcVolList :
        if "type" in v.keys() :
            if v["type"] == "RW" : RWVols.append(v)
    for v in srcVolList :
        if "type" in v.keys() :
            isSolitary = True
            if v["type"] == "RO" :
                for rw in RWVols :
                    if rw["name"] == v["name"][:-9] :
                        isSolitary = False
                if isSolitary : solitaryROVols.append(v)
    if afs.defaultConfig.moveRWVols :
        for v in RWVols :
            # check for moving osd-volumes
            if v.get("osdPolicy",0) != 0 :
                if not afs.defaultConfig.moveOSDVOlumes : continue
            # check for minSize
            if minUsage != -1 :
                if int(v.get("diskused")) < minUsage : continue

            # check for name with given regex, these checks are mutually exclusive.
            skip_it=False
            if afs.defaultConfig.ignoreRX != None :
                skip_it=False
                for rx in ignoreRX :
                    if rx.match(v["name"]) : skip_it = True
            elif afs.defaultConfig.onlyRX != None :
                skip_it=True
                for rx in onlyRX :
                    if rx.match(v["name"]) : skip_it = False
            elif afs.defaultConfig.onlyProjects != None :
                skip_it=True
                volProjects=PS.getProjectsByVolumeName(v["name"])
                for prj in volProjects :
                    if prj.name in afs.defaultConfig.onlyProjects : skip_it = False
            elif afs.defaultConfig.ignoreProjects != None :
                skip_it=False
                volProjects=PS.getProjectsByVolumeName(v["name"])
                for prj in volProjects :
                    if prj.name in afs.defaultConfig.ignoreProjects : skip_it = True

            if skip_it  : continue
            # remove osd-attributes from dict, create Obj
            v.pop("filequota")
            v.pop("osdPolicy")
            v["serv_uuid"]=afs.LookupUtil[afs.defaultConfig.CELL_NAME].getFSUUID(v["servername"])
            VolObj.setByDict(v)
            if dstFS == None :
                minNestingLevel=-1 # NestingLevel should be "spezifizitaet"
                for p in PS.getProjectsByVolumeName(v["name"]) :
                    if p.NestingLevel <  minNestingLevel or  minNestingLevel == -1 :
                        minNestingLevel =  p.NestingLevel
                        PrjObj=p
                dstSrv,dstP = PS.getNewVolumeLocation(PrjObj.name,VolObj)
                if dstSrv == None :
                    sys.stderr.write("found no appropriate location for %s, part of project %s. Skipping\n" % (VolObj.name,PrjObj.name))
                    continue
            else :
                dstSrv = dstFS.servernames[0]
                parts=FS.getPartitions(dstSrv)
                if reqDstPart == None :
                    maxFree=0
                    for p in parts :
                        if parts[p]["free"] > maxFree :
                            dstP=p
                            maxFree=parts[p]["free"]
                else :
                    dstP=reqDstPart
            print "moving volume %s from %s %s to %s %s" % (v["name"],srcFS.servernames[0],srcP,dstSrv,dstP)
            if not afs.defaultConfig.dryrun :
                VD.move(v["name"],srcFS.servernames[0],srcP,dstSrv,dstP)
                try :
                    # add RO to dstSrv if there is none yet.
                    hasRO=False
                    for ov in VS.getVolume("%s.readonly" % VolObj.name, cached=False) :
                        if ov.servername == dstSrv and ov.part == dstP :
                            hasRO = True
                    if hasRO :  # find a second server to move the RO to then
                        dstSrv,dstP = PS.getNewVolumeLocation(PrjObj.name,ov)
                        if dstSrv == None :
                            sys.stderr.write("found no appropriate location for %s. Skipping\n" % VolObj.name)
                            continue
                except : # there is no RO, so just skip this.
                    continue
            print "Moving accompanying RO to  %s %s" % (dstSrv,dstP)
            if not afs.defaultConfig.dryrun  :
                VlD.addsite(v["name"],dstSrv,dstP)
                VD.release(v["name"])
                # only remove accompanying RO from srcSRV if we are sure there is one!
                for ov in VS.getVolume("%s.readonly" % VolObj.name, cached=False) :
                    if ov.servername == srcFS.servernames[0] and ov.part == srcP :
                        VD.remove("%s.readonly" % v["name"],srcFS.servernames[0],srcP)
                        break
            movedVolcount += 1
            if movedVolcount > afs.defaultConfig.maxnum and afs.defaultConfig.maxnum > 0 :
                print "moved %d volumes. Terminating." % afs.defaultConfig.maxnum
                sys.exit(0)
            # check if partition is freed enough
            parts=FS.getPartitions(afs.defaultConfig.ssrv)
            if parts[srcP]["free"] > untilFree and untilFree > 0 :
                print "%s bytes free on spart %s." % (afs.util.afsutil.humanReadableSize(parts[srcP]["free"]/1024),srcP )
                continue
    if afs.defaultConfig.moveSolitaryROVols :
        for v in solitaryROVols :
            # check for moving osd-volumes
            if v.get("osdPolicy",0) != 0 :
                if not afs.defaultConfig.moveOSDVOlumes : continue
            # check for minSize
            if minUsage != -1 :
                if int(v.get("diskused")) < minUsage : continue

            # get RWVolName
            RWVolName=v['name'][:-len(".readonly")]
            if afs.defaultConfig.ignoreRX != None :
                skip_it=False
                for rx in ignoreRX :
                    if rx.match(RWVolName) : skip_it = True
            elif afs.defaultConfig.onlyRX != None :
                skip_it=True
                for rx in onlyRX :
                    if rx.match(RWVolName) : skip_it = False
            elif afs.defaultConfig.onlyProjects != None :
                skip_it=True
                volProjects=PS.getProjectsByVolumeName(RWVolName)
                for prj in volProjects :
                    if prj.name in afs.defaultConfig.onlyProjects : skip_it = False
            elif afs.defaultConfig.ignoreProjects != None :
                skip_it=False
                volProjects=PS.getProjectsByVolumeName(RWVolName)
                for prj in volProjects :
                    if prj.name in afs.defaultConfig.ignoreProjects : skip_it = True
            else :
                skip_it = False
            if skip_it  : continue

            # remove osd-attributes from dict, create object
            v.pop("filequota")
            v.pop("osdPolicy")
            v["serv_uuid"]=afs.LookupUtil[afs.defaultConfig.CELL_NAME].getFSUUID(v["servername"])
            VolObj.setByDict(v)

            if dstFS == None :
                # get actual stuff from live
                #Vol=VS.getVolume(v["name"],cached=False)
                minNestingLevel=-1 # NestingLevel should be "spezifizitaet"
                for p in PS.getProjectsByVolumeName(RWVolName) :
                    if p.NestingLevel <  minNestingLevel or  minNestingLevel == -1 :
                        minNestingLevel =  p.NestingLevel
                        PrjObj=p
                dstSrv,dstP = PS.getNewVolumeLocation(PrjObj.name,VolObj)
                if dstSrv == None :
                    sys.stderr.write("found no appropriate location for %s. Skipping\n" % VolObj.name)
                    continue
            else :
                dstSrv = dstFS.servernames[0]
                if reqDstPart == None :
                    dstP=None
                    maxFree=0
                    parts=FS.getPartitions(dstSrv)
                    for p in parts :
                        if parts[p]["free"] > maxFree :
                            dstP=p
                            maxFree=parts[p]["free"]

                    if dstP == None :
                        sys.stderr.write("found no appropriate partition on server %s for %s. Skipping\n" % (dstSrv,VolObj.name))
                        continue
                else :
                    # check if requested destination is ok
                    dstP=reqDstPart
                    skip_it= False
                    for rov in VS.getVolGroup(VolObj.name,cached=False)["RO"] :
                        if rov.serv_uuid == dstFS.uuid :
                            print "Found one existing RO-copy of %s on server %s. Skipping." % (VolObj.name,dstSrv)
                            skip_it = True
                    if skip_it : continue

            # XXX only move RO, if we don't have enough of them -> auto-healing
            # get RO-Sites of the Volume :
            # the most specific project for this volums counts
            print "Moving solitary RO %s to %s %s" % (v['name'],dstSrv,dstP)
            if not afs.defaultConfig.dryrun :
                VlD.addsite(RWVolName,dstSrv,dstP)
                VD.release(RWVolName)
                VD.remove(v['name'],srcFS.servernames[0],srcP)
            movedVolcount += 1
            if movedVolcount > afs.defaultConfig.maxnum and afs.defaultConfig.maxnum > 0 :
                print "moved %d volumes. Terminating." % afs.defaultConfig.maxnum
                sys.exit(0)

            # check if partition is freed enough
            parts=FS.getPartitions(afs.defaultConfig.ssrv)
            if parts[srcP]["free"] > untilFree and untilFree > 0 :
                print "%s bytes free on spart %s." % (afs.util.afsutil.humanReadableSize(parts[srcP]["free"]/1024),srcP )
                continue

<commit_msg>examples/emptyServer.py: some bugfixes<commit_after>#!/usr/bin/python

import re
import time,logging
import argparse, string,datetime,sys
import afs
from afs.util.AfsConfig import parseDefaultConfig
from afs.util.afsutil import parseHumanWriteableSize
from afs.service.OSDVolService import OSDVolService
from afs.service.OSDCellService import OSDCellService
from afs.service.OSDFsService import OSDFsService
from afs.service.ProjectService import ProjectService
from afs.model.Volume import Volume
from afs.dao.VolumeDAO import VolumeDAO
from afs.dao.VLDbDAO import VLDbDAO

global FS,PS,VS,VD,VlD

myParser=argparse.ArgumentParser(parents=[afs.argParser], add_help=False)
myParser.add_argument("--ssrv", required=True, help="server to empty")
myParser.add_argument("--spart", help="partition on server to empty. defaults to all")
myParser.add_argument("--dsrv" , help="force server to be filled. Otherwise chosen automatically by project.")
myParser.add_argument("--dpart" , help="force partition on server to be filled. Otherwise chosen automatically by project.")
group = myParser.add_mutually_exclusive_group()
group.add_argument("--ignorerx", dest="ignoreRX", action="append",  help="regEx for volumenames to ignore. All volumes not matching this will be moved.")
group.add_argument("--onlyrx", dest="onlyRX", action="append", help="regEx for volumenames to include all volumes not matching this will be ignored.")
group.add_argument("--ignoreproject", dest="ignoreProjects", action="append", help="ignore volumes of given project.")
group.add_argument("--onlyproject", dest="onlyProjects", action="append", help="only move volumes of given project.")
myParser.add_argument("--dryrun",action="store_true", help="Just print out what would be done, but don't do it.")
myParser.add_argument("--maxnum", default = 0, type=int, help="max number of Volumes to move.")
myParser.add_argument("--untilfree", default = "0", help="move until # is free on spart.")
myParser.add_argument("--rwvols", dest="moveRWVols", default=False, action="store_true", help="move rwvols with their accompanying ROs.")
myParser.add_argument("--solitaryrovols", dest="moveSolitaryROVols", default=False, action="store_true", help="move solitary rovols.")
myParser.add_argument("--minsize", dest="minVolumeUsage", default="0", help="only move volumes with minimalsize of")
myParser.add_argument("--osdvolumes", dest="moveOSDVOlumes", default=False, action="store_true", help="also move OSD-Volumes")

parseDefaultConfig(myParser)
FS=OSDFsService()
PS=ProjectService()
VS=OSDVolService()
VD=VolumeDAO()
VlD=VLDbDAO()

if not afs.defaultConfig.moveRWVols and not afs.defaultConfig.moveSolitaryROVols :
    sys.stderr.write("If you want to nmake me do anything, specify --rwvols and/or --solitaryrovols\n")
    sys.exit(1)

if afs.defaultConfig.ignoreRX != None :
    ignoreRX=[]
    for rx in afs.defaultConfig.ignoreRX :
        try :
            ignoreRX.append(re.compile(rx))
        except :
            sys.stderr.write("Cannot compile regular expression: '%s'\n" % rx)
            sys.exit(1)
elif afs.defaultConfig.onlyRX != None :
    onlyRX=[]
    for rx in afs.defaultConfig.onlyRX :
        try :
            onlyRX.append(re.compile(rx))
        except :
            sys.stderr.write("Cannot compile regular expression: '%s'\n" % rx)
            sys.exit(1)
srcFS=FS.getFileServer(afs.defaultConfig.ssrv)

if srcFS == None :
    sys.stderr.write("src server %s does not exist." % srcFS)
    sys.exit(2)
if afs.defaultConfig.spart != None :
    if not afs.defaultConfig.spart in srcFS.parts :
        sys.stderr.write("Partition %s does not exist on server %s\n" % (afs.defaultConfig.spart, afs.defaultConfig.ssrv))
        sys.exit(2)
    # we cycle thorugh this list later
    srcParts=[afs.defaultConfig.spart]
else :
    srcParts=srcFS.parts

if afs.defaultConfig.dsrv != None :
    dstFS=FS.getFileServer(afs.defaultConfig.dsrv)
    if dstFS == None :
        sys.stderr.write("destination server %s does not exist.\n" % srcFS)
        sys.exit(2)

    if afs.defaultConfig.dpart != None :
        if not afs.defaultConfig.dpart in dstFS.parts :
            sys.stderr.write("Partition %s does not exist on server %s\n" % (afs.defaultConfig.dpart, afs.defaultConfig.dsrv))
            sys.exit(2)
        else :
            reqDstPart=afs.defaultConfig.dpart
    else :
        reqDstPart=None
else :
    dstFS=None
    if afs.defaultConfig.dpart != None :
        sys.stderr.write("Warning: ignoring given dpart=%s, because no dsrv has been specified.\n" % afs.defaultConfig.dpart)

# XXX we should handle everything internally with bytes
untilFree=parseHumanWriteableSize(afs.defaultConfig.untilfree)/1024
minUsage=parseHumanWriteableSize(afs.defaultConfig.minVolumeUsage)/1024
VolObj = Volume()
movedVolcount=0
for srcP in srcParts :
    print "Processing Partition %s...." % srcP
    # check if partition is freed enough
    parts=FS.getPartitions(afs.defaultConfig.ssrv)
    if parts[srcP]["free"] > untilFree and untilFree > 0 :
        print "already %s Bytes free on spart %s." % (afs.util.afsutil.humanReadableSize(parts[srcP]["free"]),srcP )
        continue
    # get list of volumes to move
    srcVolList=FS.getVolList(srcFS.servernames[0],srcP,cached=False)
    # get RW Volumes :
    RWVols=[]
    solitaryROVols=[]
    for v in srcVolList :
        if "type" in v.keys() :
            if v["type"] == "RW" : RWVols.append(v)
    for v in srcVolList :
        if "type" in v.keys() :
            isSolitary = True
            if v["type"] == "RO" :
                for rw in RWVols :
                    if rw["name"] == v["name"][:-9] :
                        isSolitary = False
                if isSolitary : solitaryROVols.append(v)
    if afs.defaultConfig.moveRWVols :
        for v in RWVols :
            # check for moving osd-volumes
            if v.get("osdPolicy",0) != 0 :
                if not afs.defaultConfig.moveOSDVOlumes :
                    print "Skipping %s, because it is an OSD-Volume" % v["name"]
                    continue
            # check for minSize
            if minUsage != 0 :
                if int(v.get("diskused")) < minUsage :
                    print "Skipping %s, because it is smaller than %s" % minUsage
                    continue

            # check for name with given regex, these checks are mutually exclusive.
            skip_it=False
            if afs.defaultConfig.ignoreRX != None :
                skip_it=False
                for rx in ignoreRX :
                    if rx.match(v["name"]) : skip_it = True
            elif afs.defaultConfig.onlyRX != None :
                skip_it=True
                for rx in onlyRX :
                    if rx.match(v["name"]) : skip_it = False
            elif afs.defaultConfig.onlyProjects != None :
                skip_it=True
                volProjects=PS.getProjectsByVolumeName(v["name"])
                for prj in volProjects :
                    if prj.name in afs.defaultConfig.onlyProjects : skip_it = False
            elif afs.defaultConfig.ignoreProjects != None :
                skip_it=False
                volProjects=PS.getProjectsByVolumeName(v["name"])
                for prj in volProjects :
                    if prj.name in afs.defaultConfig.ignoreProjects : skip_it = True

            if skip_it  : continue
            # remove osd-attributes from dict, create Obj
            v.pop("filequota")
            v.pop("osdPolicy")
            v["serv_uuid"]=afs.LookupUtil[afs.defaultConfig.CELL_NAME].getFSUUID(v["servername"])
            VolObj.setByDict(v)
            if dstFS == None :
                minNestingLevel=-1 # NestingLevel should be "spezifizitaet"
                for p in PS.getProjectsByVolumeName(v["name"]) :
                    if p.NestingLevel <  minNestingLevel or  minNestingLevel == -1 :
                        minNestingLevel =  p.NestingLevel
                        PrjObj=p
                dstSrv,dstP = PS.getNewVolumeLocation(PrjObj.name,VolObj)
                if dstSrv == None :
                    sys.stderr.write("found no appropriate location for %s, part of project %s. Skipping\n" % (VolObj.name,PrjObj.name))
                    continue
            else :
                dstSrv = dstFS.servernames[0]
                parts=FS.getPartitions(dstSrv)
                if reqDstPart == None :
                    maxFree=0
                    for p in parts :
                        if parts[p]["free"] > maxFree :
                            dstP=p
                            maxFree=parts[p]["free"]
                else :
                    dstP=reqDstPart
            print "moving volume %s from %s %s to %s %s" % (v["name"],srcFS.servernames[0],srcP,dstSrv,dstP)
            if not afs.defaultConfig.dryrun :
                VD.move(v["name"],srcFS.servernames[0],srcP,dstSrv,dstP)
                try :
                    # add RO to dstSrv if there is none yet.
                    hasRO=False
                    for ov in VS.getVolume("%s.readonly" % VolObj.name, cached=False) :
                        if ov.servername == dstSrv and ov.part == dstP :
                            hasRO = True
                    if hasRO :  # find a second server to move the RO to then
                        dstSrv,dstP = PS.getNewVolumeLocation(PrjObj.name,ov)
                        if dstSrv == None :
                            sys.stderr.write("found no appropriate location for %s. Skipping\n" % VolObj.name)
                            continue
                except : # there is no RO, so just skip this.
                    continue
            print "Moving accompanying RO to  %s %s" % (dstSrv,dstP)
            if not afs.defaultConfig.dryrun  :
                VlD.addsite(v["name"],dstSrv,dstP)
                VD.release(v["name"])
                # only remove accompanying RO from srcSRV if we are sure there is one!
                for ov in VS.getVolume("%s.readonly" % VolObj.name, cached=False) :
                    if ov.servername == srcFS.servernames[0] and ov.part == srcP :
                        VD.remove("%s.readonly" % v["name"],srcFS.servernames[0],srcP)
                        break
            movedVolcount += 1
            if movedVolcount > afs.defaultConfig.maxnum and afs.defaultConfig.maxnum > 0 :
                print "moved %d volumes. Terminating." % afs.defaultConfig.maxnum
                sys.exit(0)
            # check if partition is freed enough
            parts=FS.getPartitions(afs.defaultConfig.ssrv)
            if parts[srcP]["free"] > untilFree and untilFree > 0 :
                print "%s bytes free on spart %s." % (afs.util.afsutil.humanReadableSize(parts[srcP]["free"]),srcP )
                break
    if afs.defaultConfig.moveSolitaryROVols :
        for v in solitaryROVols :
            # get RWVolName
            RWVolName=v['name'][:-len(".readonly")]
            # check for moving osd-volumes
            if v.get("osdPolicy",0) != 0 :
                if not afs.defaultConfig.moveOSDVOlumes :
                    print "Skipping %s, because it is an OSD-Volume" % v["name"]
                    continue
            # check for minSize
            if minUsage != 0 :
                if int(v.get("diskused")) < minUsage :
                    print "Skipping %s, because it is smaller than %s" % minUsage
                    continue

            if afs.defaultConfig.ignoreRX != None :
                skip_it=False
                for rx in ignoreRX :
                    if rx.match(RWVolName) : skip_it = True
            elif afs.defaultConfig.onlyRX != None :
                skip_it=True
                for rx in onlyRX :
                    if rx.match(RWVolName) : skip_it = False
            elif afs.defaultConfig.onlyProjects != None :
                skip_it=True
                volProjects=PS.getProjectsByVolumeName(RWVolName)
                for prj in volProjects :
                    if prj.name in afs.defaultConfig.onlyProjects : skip_it = False
            elif afs.defaultConfig.ignoreProjects != None :
                skip_it=False
                volProjects=PS.getProjectsByVolumeName(RWVolName)
                for prj in volProjects :
                    if prj.name in afs.defaultConfig.ignoreProjects : skip_it = True
            else :
                skip_it = False
            if skip_it  : continue

            # remove osd-attributes from dict, create object
            v.pop("filequota")
            v.pop("osdPolicy")
            v["serv_uuid"]=afs.LookupUtil[afs.defaultConfig.CELL_NAME].getFSUUID(v["servername"])
            VolObj.setByDict(v)

            if dstFS == None :
                # get actual stuff from live
                #Vol=VS.getVolume(v["name"],cached=False)
                minNestingLevel=-1 # NestingLevel should be "spezifizitaet"
                for p in PS.getProjectsByVolumeName(RWVolName) :
                    if p.NestingLevel <  minNestingLevel or  minNestingLevel == -1 :
                        minNestingLevel =  p.NestingLevel
                        PrjObj=p
                dstSrv,dstP = PS.getNewVolumeLocation(PrjObj.name,VolObj)
                if dstSrv == None :
                    sys.stderr.write("found no appropriate location for %s. Skipping\n" % VolObj.name)
                    continue
            else :
                dstSrv = dstFS.servernames[0]
                if reqDstPart == None :
                    dstP=None
                    maxFree=0
                    parts=FS.getPartitions(dstSrv)
                    for p in parts :
                        if parts[p]["free"] > maxFree :
                            dstP=p
                            maxFree=parts[p]["free"]

                    if dstP == None :
                        sys.stderr.write("found no appropriate partition on server %s for %s. Skipping\n" % (dstSrv,VolObj.name))
                        continue
                else :
                    # check if requested destination is ok
                    dstP=reqDstPart
                    skip_it= False
                    for rov in VS.getVolGroup(VolObj.name,cached=False)["RO"] :
                        if rov.serv_uuid == dstFS.uuid :
                            print "Found one existing RO-copy of %s on server %s. Skipping." % (VolObj.name,dstSrv)
                            skip_it = True
                    if skip_it : continue

            # XXX only move RO, if we don't have enough of them -> auto-healing
            # get RO-Sites of the Volume :
            # the most specific project for this volums counts
            print "Moving solitary RO %s to %s %s" % (v['name'],dstSrv,dstP)
            if not afs.defaultConfig.dryrun :
                VlD.addsite(RWVolName,dstSrv,dstP)
                VD.release(RWVolName)
                VD.remove(v['name'],srcFS.servernames[0],srcP)
            movedVolcount += 1
            if movedVolcount > afs.defaultConfig.maxnum and afs.defaultConfig.maxnum > 0 :
                print "moved %d volumes. Terminating." % afs.defaultConfig.maxnum
                sys.exit(0)

            # check if partition is freed enough
            parts=FS.getPartitions(afs.defaultConfig.ssrv)
            if parts[srcP]["free"] > untilFree and untilFree > 0 :
                print "%s bytes free on spart %s." % (afs.util.afsutil.humanReadableSize(parts[srcP]["free"]),srcP )
                break

<|endoftext|>


<commit_before>from django.conf import settings
from django.db import models
import logging

log = logging.getLogger('genoome.twentythree.models')

class Token23(models.Model):
    user = models.ForeignKey(settings.AUTH_USER_MODEL, unique=True)
    access_token = models.TextField()
    refresh_token = models.TextField()
    scope = models.TextField()
    created_at = models.DateTimeField(auto_now_add=True)
    modified_at = models.DateTimeField(auto_now=True)

    def _api_get(self, url):
        headers = {
            'Authorization': 'Bearer %s' % self.access_token,
        }
        response = requests.get(url, headers=headers)
        if response.status_code != 200:
            # https://api.23andme.com/docs/errors/
            log.warning('23andme error response: %s\nurl:%s', response.text, url)
            raise self.ClientError
        return response.json()

    @classmethod
    def get_by_code(klass, user, code):
        # https://api.23andme.com/docs/authentication/
        # curl https://api.23andme.com/token/
        #       -d client_id=xxx \
        #       -d client_secret=yyy \
        #       -d grant_type=authorization_code \
        #       -d code=zzz \
        #       -d "redirect_uri=https://localhost:5000/receive_code/"
        #       -d "scope=basic%20rs3094315"
        post_data = {
            'client_id': settings.CLIENT_ID23,
            'client_secret': settings.CLIENT_SECRET23,
            'scope': 'basic genomes',
            'code': code,
            'grant_type': 'authorization_code',
            'redirect_uri': settings.COMEBACK_URL23,
        }
        response = requests.post('/token/', data=post_data, timeout=30.00, verify=True)
        if response.status_code != 200:
            log.error('Problem fetching token %', response.status_code, response.text)
            raise self.ClientError
        data = response.json()
        return klass.objects.create(user=user, access_token=data['access_token'], refresh_token=data['refresh_token'], scope=data['scope'])

    def refresh(self):
        post_data = {
            'client_id': settings.CLIENT_ID23,
            'client_secret': settings.CLIENT_SECRET23,
            'scope': self.scope,
            'refresh_token': self.refresh_token,
            'grant_type': 'refresh_token',
            'redirect_uri': settings.COMEBACK_URL23,
        }
        response = requests.post('/token/', data=post_data, timeout=30.00, verify=True)
        if response.status_code != 200:
            log.error('Problem refreshing token %', response.status_code, response.text)
            raise self.ClientError
        data = response.json()
        self.access_token = data['access_token']
        self.refresh_token = data['refresh_token']
        self.scope = data['scope']
        self.save()

    def get_genome(self, profile_id23):
        # GET /1/genomes/profile_id/?unfiltered=...
        # curl https://api.23andme.com/1/genomes/c44.../ -H "..."
        # https://api.23andme.com/res/txt/snps.b4e00fe1db50.data
        # scope required: genomes
        data = self._api_get('/1/genomes/%s/' % (profile_id23,))
        return data['genome']

    def get_profiles(self):
        # GET /1/user/
        # # JSON response:
        #{
        #    "id": "a42e94634e3f7683",
        #    "profiles": [
        #        {
        #            "genotyped": true,
        #            "id": "c4480ba411939067"
        #        }, ...
        #    ]
        #}
        # scope required: basic

        data = self._api_get('/1/user/')
        return data['profiles']

class CeleryTask23(models.Model):
    STATUS_CHOICES = (
        ('new', 'new'),
        ('fetching', 'fetching genome'),
        ('parsing', 'parsing genome'),
        ('error', 'error'),
    )
    user = models.ForeignKey(settings.AUTH_USER_MODEL, unique=True)
    chosen_profile = models.TextField()
    fetch_task_id = models.TextField()
    analyze_order = models.ForeignKey('disease.models.AnalyzeDataOrder', null=True)
    process_task_id = models.TextField(null=True)
    status = models.TextField(choices=STATUS_CHOICES, default='new')
<commit_msg>23andme wip<commit_after>from django.conf import settings
from django.db import models
import logging

log = logging.getLogger('genoome.twentythree.models')

class Token23(models.Model):
    user = models.ForeignKey(settings.AUTH_USER_MODEL, unique=True)
    access_token = models.TextField()
    refresh_token = models.TextField()
    scope = models.TextField()
    created_at = models.DateTimeField(auto_now_add=True)
    modified_at = models.DateTimeField(auto_now=True)

    class ClientError(Exception):
        pass

    def _api_get(self, url):
        headers = {
            'Authorization': 'Bearer %s' % self.access_token,
        }
        response = requests.get(url, headers=headers)
        if response.status_code != 200:
            # https://api.23andme.com/docs/errors/
            log.warning('23andme error response: %s\nurl:%s', response.text, url)
            raise self.ClientError
        return response.json()

    @classmethod
    def get_by_code(klass, user, code):
        # https://api.23andme.com/docs/authentication/
        # curl https://api.23andme.com/token/
        #       -d client_id=xxx \
        #       -d client_secret=yyy \
        #       -d grant_type=authorization_code \
        #       -d code=zzz \
        #       -d "redirect_uri=https://localhost:5000/receive_code/"
        #       -d "scope=basic%20rs3094315"
        post_data = {
            'client_id': settings.CLIENT_ID23,
            'client_secret': settings.CLIENT_SECRET23,
            'scope': 'basic genomes',
            'code': code,
            'grant_type': 'authorization_code',
            'redirect_uri': settings.COMEBACK_URL23,
        }
        response = requests.post('/token/', data=post_data, timeout=30.00, verify=True)
        if response.status_code != 200:
            log.error('Problem fetching token %', response.status_code, response.text)
            raise self.ClientError
        data = response.json()
        return klass.objects.create(user=user, access_token=data['access_token'], refresh_token=data['refresh_token'], scope=data['scope'])

    def refresh(self):
        post_data = {
            'client_id': settings.CLIENT_ID23,
            'client_secret': settings.CLIENT_SECRET23,
            'scope': self.scope,
            'refresh_token': self.refresh_token,
            'grant_type': 'refresh_token',
            'redirect_uri': settings.COMEBACK_URL23,
        }
        response = requests.post('/token/', data=post_data, timeout=30.00, verify=True)
        if response.status_code != 200:
            log.error('Problem refreshing token %', response.status_code, response.text)
            raise self.ClientError
        data = response.json()
        self.access_token = data['access_token']
        self.refresh_token = data['refresh_token']
        self.scope = data['scope']
        self.save()

    def get_genome(self, profile_id23):
        # GET /1/genomes/profile_id/?unfiltered=...
        # curl https://api.23andme.com/1/genomes/c44.../ -H "..."
        # https://api.23andme.com/res/txt/snps.b4e00fe1db50.data
        # scope required: genomes
        data = self._api_get('/1/genomes/%s/' % (profile_id23,))
        return data['genome']

    def get_profiles(self):
        # GET /1/user/
        # # JSON response:
        #{
        #    "id": "a42e94634e3f7683",
        #    "profiles": [
        #        {
        #            "genotyped": true,
        #            "id": "c4480ba411939067"
        #        }, ...
        #    ]
        #}
        # scope required: basic

        data = self._api_get('/1/user/')
        return data['profiles']

class CeleryTask23(models.Model):
    STATUS_CHOICES = (
        ('new', 'new'),
        ('fetching', 'fetching genome'),
        ('parsing', 'parsing genome'),
        ('error', 'error'),
    )
    user = models.ForeignKey(settings.AUTH_USER_MODEL, unique=True)
    chosen_profile = models.TextField()
    fetch_task_id = models.TextField()
    analyze_order = models.ForeignKey('disease.models.AnalyzeDataOrder', null=True)
    process_task_id = models.TextField(null=True)
    status = models.TextField(choices=STATUS_CHOICES, default='new')
<|endoftext|>