<commit_before>
def plotter(title,
            df,
            x_label = False,
            y_label = False,
            style = 'ggplot',
            figsize = (13, 6),
            save = False,
            legend = 'best',
            num_to_plot = 7,
            tex = 'try',
            colours = 'Paired',
            cumulative = False,
            **kwargs):
    """plot interrogator() or editor() output.

    **kwargs are for pandas first, which can then send them through to matplotlib.plot():

    http://pandas.pydata.org/pandas-docs/dev/generated/pandas.DataFrame.plot.html
    http://matplotlib.org/api/pyplot_api.html#matplotlib.pyplot.plot

    """

    import os
    import matplotlib.pyplot as plt
    from matplotlib import rc
    import pandas
    import pandas as pd
    from pandas import DataFrame

    import numpy
    from time import localtime, strftime
    from corpkit.tests import check_pytex

    have_python_tex = check_pytex()

    # are we doing subplots?
    sbplt = False
    if 'subplots' in kwargs:
        if kwargs['subplots'] is True:
            sbplt = True

    if colours is True:
        colours = 'Paired'

    styles = ['dark_background', 'bmh', 'grayscale', 'ggplot', 'fivethirtyeight', 'matplotlib']
    if style not in styles:
        raise ValueError('Style %s not found. Use %s' % (style, ', '.join(styles)))

    # try to use tex
    using_tex = False
    if tex == 'try' or tex is True:
        try:
            rc('text', usetex=True)
            rc('font', **{'family': 'serif', 'serif': ['Computer Modern']})
            using_tex = True
        except:
            rc('text', usetex=False)
    else:
        rc('text', usetex=False)

  # pie chart requires subplots
    piemode = False
    if 'kind' in kwargs:
        if kwargs['kind'] == 'pie':
            piemode = True
            sbplt = True

    kwargs['subplots'] = sbplt
    # copy data, make series into df
    dataframe = df.copy()
    was_series = False
    if type(dataframe) == pandas.core.series.Series:
        was_series = True
        dataframe = DataFrame(dataframe)
    # attempt to convert x axis to ints:
    try:
        dataframe.index = [int(i) for i in list(dataframe.index)]
    except:
        pass

    # remove totals if there ... maybe redundant
    try:
        dataframe = dataframe.drop('Total', axis = 0)
    except:
        pass

    if cumulative:
        dataframe = DataFrame(dataframe.cumsum())

    if num_to_plot == 'all':
        if was_series:
            num_to_plot = len(dataframe)
        else:
            num_to_plot = len(list(dataframe.columns))

    #  use colormap if need be:
    colormap = False
    if num_to_plot > 7:
        if 'kind' in kwargs:
            if kwargs['kind'] in ['pie', 'line', 'area']:
                if colours:
                    kwargs['colormap'] = colours
        else:
            if colours:
                kwargs['colormap'] = colours
    if 'kind' in kwargs:
        if colours:
            if kwargs['kind'].startswith('bar'):
                if len(list(dataframe.columns)) == 1:
                    import numpy as np
                    the_range = np.linspace(0, 1, num_to_plot)
                    cmap = plt.get_cmap(colours)
                    kwargs['colors'] = [cmap(n) for n in the_range]

    # no legend for bar chart if just one label
    if 'kind' in kwargs:
        if kwargs['kind'] in ['bar', 'barh', 'area', 'line']:
            if was_series:
                legend = False

    # cut dataframe if just_totals
    try:
        tst = dataframe['Combined total']
        dataframe = dataframe.head(num_to_plot)
    except:
        pass

    # rotate automatically
    #if len(max(list(dataframe.columns), key=len)) > 6:
        #kwargs['rot'] = 45

    # no title for subplots because ugly
    if sbplt and not piemode:
        title_to_show = ''
        figsize = (figsize[0], figsize[1] * 2.5)
    else:
        title_to_show = title
    if piemode:
        title_to_show = ''
        figsize = (figsize[0], figsize[0])
        legend = False

    # this gets tid of the y_label thing showing up for pie mode...
    if piemode:
        if len(dataframe.columns) == 1:
            dataframe.columns = ['']

    # try to use styles
    if style != 'matplotlib' and style is not False:
        with plt.style.context((style)):
            if type(legend) == bool:
                a_plot = DataFrame(dataframe[list(dataframe.columns)[:num_to_plot]]).plot(title = title_to_show, figsize = figsize, legend = legend, **kwargs)
            else:
                a_plot = DataFrame(dataframe[list(dataframe.columns)[:num_to_plot]]).plot(title = title_to_show, figsize = figsize, **kwargs)
    # plot without style
    else:
        if type(legend) == bool:
            a_plot = DataFrame(dataframe[list(dataframe.columns)[:num_to_plot]]).plot(title = title_to_show, figsize = figsize, legend = legend, **kwargs)
        else:
            a_plot = DataFrame(dataframe[list(dataframe.columns)[:num_to_plot]]).plot(title = title_to_show, figsize = figsize, **kwargs)

    # bigger ticks?
    #if using_tex:
    #    size = (figsize[0] / 2) + 12
    #else:
    #    size = (figsize[0] / 2) + 7
    #if type(a_plot) != numpy.ndarray:
    #    a_plot.tick_params(axis='x', which='major', labelsize=size)
    # make and set x label:
    if x_label is False:
        check_x_axis = list(dataframe.index)[0] # get first entry# get second entry of first entry (year, count)
        try:
            check_x_axis = int(check_x_axis)
            if 1500 < check_x_axis < 2050:
                x_label = 'Year'
            else:
                x_label = 'Group'
        except:
            x_label = 'Group'

    if x_label is not None:
        if not sbplt:
            plt.xlabel(x_label)

    # make and set y label
    if y_label is False:
        try:
            if type(dataframe[list(dataframe.columns)[0]][list(dataframe.index)[0]]) == numpy.float64:
                y_label = 'Percentage'
            else:
                y_label = 'Absolute frequency'
        except:
            if type(dataframe['Total'][list(dataframe.index)[0]]) == numpy.float64:
                y_label = 'Percentage'
            else:
                y_label = 'Absolute frequency'

    if y_label is not None:
        if not sbplt:
            plt.ylabel(y_label)

    # hacky: turn legend into subplot titles :)
    if sbplt:
        for index, f in enumerate(a_plot):
            titletext = list(dataframe.columns)[index]
            if not piemode:
                f.legend_.remove()
                f.set_title(titletext)

    # legend values
    possible = {'best': 0, 'upper right': 1, 'upper left': 2, 'lower left': 3, 'lower right': 4,
                'right': 5, 'center left': 6, 'center right': 7, 'lower center': 8,
                'upper center': 9, 'center': 10}

    if legend:
        if not legend.startswith('o'):
            if legend == 'outside right':
                legend = 'outside upper right'
            if type(legend) == int:
                the_loc = legend
            elif type(legend) == str:
                try:
                    the_loc = possible[legend]
                except KeyError:
                    raise KeyError('legend value must be one of:\n%s\n or an int between 0-10.' %', '.join(possible.keys()))
            else:
                raise KeyError('legend value must be one of:\n%s' %', '.join(possible.keys()))
            if not sbplt:
                plt.legend(loc=the_loc)
        elif legend.startswith('o'):
            if legend.startswith('outside r') or legend.startswith('o r'):
                legend = 'outside upper right'
            if not sbplt:
                os, plc = legend.split(' ', 1)
                try:
                    if plc == 'upper right':
                    #the_loc = possible[plc]
                        plt.legend(bbox_to_anchor=(1.02, 1), loc=2, borderaxespad=1)
                    if plc == 'center right':
                        plt.legend(bbox_to_anchor=(1.02, 0.5), loc='center left', borderaxespad=1)
                    if plc == 'lower right':
                        plt.legend(bbox_to_anchor=(1.02, 0), loc='lower left', borderaxespad=1)
                #if plc == 'upper left':
                    #plt.legend(bbox_to_anchor=(1, 0), loc=1, borderaxespad=1)
                except KeyError:
                    raise KeyError('legend value must be one of: %s\n or an int between 0-10.' %', '.join(possible.keys()))
        else:
            raise KeyError('legend value must be one of: %s' %', '.join(possible.keys()))

    # make figure
    plt.subplots_adjust(bottom=0.20)
    if not sbplt:
        fig1 = a_plot.get_figure()
        if not have_python_tex:
            fig1.show()
    else:
        if not have_python_tex:
            plt.show()

    if not save:
        return

    if save:
        if have_python_tex:
            imagefolder = '../images'
        else:
            imagefolder = 'images'

        savename = get_savename(imagefolder, save = save, title = title)

        if not os.path.isdir(imagefolder):
            os.makedirs(imagefolder)

        # save image and get on with our lives
        fig1.savefig(savename, dpi=150, transparent=True)
        time = strftime("%H:%M:%S", localtime())
        if os.path.isfile(savename):
            print '\n' + time + ": " + savename + " created."
        else:
            raise ValueError("Error making %s." % savename)

    def get_savename(imagefolder, save = False, title = False):
        """Come up with the savename for the image."""

        def urlify(s):
            "Turn title into filename"
            import re
            s = s.lower()
            s = re.sub(r"[^\w\s]", '', s)
            s = re.sub(r"\s+", '-', s)
            return s

        # name as
        if type(save) == str:
            savename = os.path.join(imagefolder, urlify(save) + '.png')
        else:
            if title:
                savename = os.path.join(imagefolder, urlify(title) + '.png')
            # generic sequential naming
            else:
                list_images = [i for i in sorted(os.listdir(imagefolder)) if i.startswith('image-')]
                if len(list_images) > 0:
                    num = int(list_images[-1].split('-')[1].split('.')[0]) + 1
                    autoname = 'image-' + str(num).zfill(3)
                    savename = os.path.join(imagefolder, autoname + '.png')
                else:
                    savename = os.path.join(imagefolder, 'image-001.png')
        if savename.endswith('.png.png'):
            savename = savename[:-4]
        return savename<commit_msg>fixed save problems<commit_after>
def plotter(title,
            df,
            x_label = False,
            y_label = False,
            style = 'ggplot',
            figsize = (13, 6),
            save = False,
            legend = 'best',
            num_to_plot = 7,
            tex = 'try',
            colours = 'Paired',
            cumulative = False,
            **kwargs):
    """plot interrogator() or editor() output.

    **kwargs are for pandas first, which can then send them through to matplotlib.plot():

    http://pandas.pydata.org/pandas-docs/dev/generated/pandas.DataFrame.plot.html
    http://matplotlib.org/api/pyplot_api.html#matplotlib.pyplot.plot

    """

    import os
    import matplotlib.pyplot as plt
    from matplotlib import rc
    import pandas
    import pandas as pd
    from pandas import DataFrame

    import numpy
    from time import localtime, strftime
    from corpkit.tests import check_pytex

    have_python_tex = check_pytex()

    def get_savename(imagefolder, save = False, title = False):
        """Come up with the savename for the image."""
        import os
        def urlify(s):
            "Turn title into filename"
            import re
            s = s.lower()
            s = re.sub(r"[^\w\s]", '', s)
            s = re.sub(r"\s+", '-', s)
            return s
        # name as
        if type(save) == str:
            savename = os.path.join(imagefolder, (urlify(save) + '.png'))
        #this 'else' is redundant now that title is obligatory
        else:
            if title:
                filename = urlify(title) + '.png'
                savename = os.path.join(imagefolder, filename)
        #    # generic sequential naming
        #    else:
        #        list_images = [i for i in sorted(os.listdir(imagefolder)) if i.startswith('image-')]
        #        if len(list_images) > 0:
        #            num = int(list_images[-1].split('-')[1].split('.')[0]) + 1
        #            autoname = 'image-' + str(num).zfill(3)
        #            savename = os.path.join(imagefolder, autoname + '.png')
        #        else:
        #            savename = os.path.join(imagefolder, 'image-001.png')
        #if savename.endswith('.png.png'):
        #    savename = savename[:-4]
        return savename

    # are we doing subplots?
    sbplt = False
    if 'subplots' in kwargs:
        if kwargs['subplots'] is True:
            sbplt = True

    if colours is True:
        colours = 'Paired'

    styles = ['dark_background', 'bmh', 'grayscale', 'ggplot', 'fivethirtyeight', 'matplotlib']
    if style not in styles:
        raise ValueError('Style %s not found. Use %s' % (style, ', '.join(styles)))

    # try to use tex
    using_tex = False
    if tex == 'try' or tex is True:
        try:
            rc('text', usetex=True)
            rc('font', **{'family': 'serif', 'serif': ['Computer Modern']})
            using_tex = True
        except:
            rc('text', usetex=False)
    else:
        rc('text', usetex=False)

  # pie chart requires subplots
    piemode = False
    if 'kind' in kwargs:
        if kwargs['kind'] == 'pie':
            piemode = True
            sbplt = True

    kwargs['subplots'] = sbplt
    # copy data, make series into df
    dataframe = df.copy()
    was_series = False
    if type(dataframe) == pandas.core.series.Series:
        was_series = True
        dataframe = DataFrame(dataframe)
    # attempt to convert x axis to ints:
    try:
        dataframe.index = [int(i) for i in list(dataframe.index)]
    except:
        pass

    # remove totals if there ... maybe redundant
    try:
        dataframe = dataframe.drop('Total', axis = 0)
    except:
        pass

    if cumulative:
        dataframe = DataFrame(dataframe.cumsum())

    if num_to_plot == 'all':
        if was_series:
            num_to_plot = len(dataframe)
        else:
            num_to_plot = len(list(dataframe.columns))

    #  use colormap if need be:
    colormap = False
    if num_to_plot > 7:
        if 'kind' in kwargs:
            if kwargs['kind'] in ['pie', 'line', 'area']:
                if colours:
                    kwargs['colormap'] = colours
        else:
            if colours:
                kwargs['colormap'] = colours
    if 'kind' in kwargs:
        if colours:
            if kwargs['kind'].startswith('bar'):
                if len(list(dataframe.columns)) == 1:
                    import numpy as np
                    the_range = np.linspace(0, 1, num_to_plot)
                    cmap = plt.get_cmap(colours)
                    kwargs['colors'] = [cmap(n) for n in the_range]

    # no legend for bar chart if just one label
    if 'kind' in kwargs:
        if kwargs['kind'] in ['bar', 'barh', 'area', 'line']:
            if was_series:
                legend = False

    # the default legend placement
    if legend is True:
        legend = 'best'

    # cut dataframe if just_totals
    try:
        tst = dataframe['Combined total']
        dataframe = dataframe.head(num_to_plot)
    except:
        pass

    #rotate automatically
    #if len(max(list(dataframe.columns), key=len)) > 6:
        #kwargs['rot'] = 45

    # no title for subplots because ugly
    if sbplt and not piemode:
        title_to_show = ''
        figsize = (figsize[0], figsize[1] * 2.5)
    else:
        title_to_show = title
    if piemode:
        title_to_show = ''
        figsize = (figsize[0], figsize[0])
        legend = False

    # this gets tid of the y_label thing showing up for pie mode...
    if piemode:
        if len(dataframe.columns) == 1:
            dataframe.columns = ['']

    # use styles and plot
    with plt.style.context((style)):
        if type(legend) == bool:
            a_plot = DataFrame(dataframe[list(dataframe.columns)[:num_to_plot]]).plot(title = title_to_show, figsize = figsize, legend = legend, **kwargs)
        else:
            a_plot = DataFrame(dataframe[list(dataframe.columns)[:num_to_plot]]).plot(title = title_to_show, figsize = figsize, **kwargs)

    if x_label is False:
        check_x_axis = list(dataframe.index)[0] # get first entry# get second entry of first entry (year, count)
        try:
            check_x_axis = int(check_x_axis)
            if 1500 < check_x_axis < 2050:
                x_label = 'Year'
            else:
                x_label = 'Group'
        except:
            x_label = 'Group'

    if x_label is not None:
        if not sbplt:
            plt.xlabel(x_label)

    # make and set y label
    if y_label is False:
        try:
            if type(dataframe[list(dataframe.columns)[0]][list(dataframe.index)[0]]) == numpy.float64:
                y_label = 'Percentage'
            else:
                y_label = 'Absolute frequency'
        except:
            if type(dataframe['Total'][list(dataframe.index)[0]]) == numpy.float64:
                y_label = 'Percentage'
            else:
                y_label = 'Absolute frequency'

    if y_label is not None:
        if not sbplt:
            plt.ylabel(y_label)

    # hacky: turn legend into subplot titles :)
    if sbplt:
        for index, f in enumerate(a_plot):
            titletext = list(dataframe.columns)[index]
            if not piemode:
                f.legend_.remove()
                f.set_title(titletext)

    # legend values
    possible = {'best': 0, 'upper right': 1, 'upper left': 2, 'lower left': 3, 'lower right': 4,
                'right': 5, 'center left': 6, 'center right': 7, 'lower center': 8,
                'upper center': 9, 'center': 10}

    if legend:
        with plt.style.context((style)):
            if not legend.startswith('o'):
                if legend == 'outside right':
                    legend = 'outside upper right'
                if type(legend) == int:
                    the_loc = legend
                elif type(legend) == str:
                    try:
                        the_loc = possible[legend]
                    except KeyError:
                        raise KeyError('legend value must be one of:\n%s\n or an int between 0-10.' %', '.join(possible.keys()))
                else:
                    raise KeyError('legend value must be one of:\n%s' %', '.join(possible.keys()))
                if not sbplt:
                    if not sbplt:
                        plt.legend(loc=the_loc, framealpha=.8)
                    else:
                        plt.legend(loc=the_loc)
            elif legend.startswith('o'):
                if legend.startswith('outside r') or legend.startswith('o r'):
                    legend = 'outside upper right'
                if not sbplt:
                    os, plc = legend.split(' ', 1)
                    try:
                        if plc == 'upper right':
                        #the_loc = possible[plc]
                            plt.legend(bbox_to_anchor=(1.02, 1), loc=2, borderaxespad=1)
                        if plc == 'center right':
                            plt.legend(bbox_to_anchor=(1.02, 0.5), loc='center left', borderaxespad=1)
                        if plc == 'lower right':
                            plt.legend(bbox_to_anchor=(1.02, 0), loc='lower left', borderaxespad=1)
                    #if plc == 'upper left':
                        #plt.legend(bbox_to_anchor=(1, 0), loc=1, borderaxespad=1)
                    except KeyError:
                        raise KeyError('legend value must be one of: %s\n or an int between 0-10.' %', '.join(possible.keys()))
            else:
                raise KeyError('legend value must be one of: %s' %', '.join(possible.keys()))

    # make room at the bottom for label?
    plt.subplots_adjust(bottom=0.20)

    # make figure
    if not sbplt:
        fig1 = a_plot.get_figure()
        if not have_python_tex:
            fig1.show()
    else:
        if not have_python_tex:
            plt.show()

    if not save:
        return

    if save:
        import os
        if have_python_tex:
            imagefolder = '../images'
        else:
            imagefolder = 'images'

        savename = get_savename(imagefolder, save = save, title = title)

        if not os.path.isdir(imagefolder):
            os.makedirs(imagefolder)

        # save image and get on with our lives
        if not sbplt:
            fig1.savefig(savename, dpi=150, transparent=True)
        else:
            plt.savefig(savename, dpi=150, transparent=True)
        time = strftime("%H:%M:%S", localtime())
        if os.path.isfile(savename):
            print '\n' + time + ": " + savename + " created."
        else:
            raise ValueError("Error making %s." % savename)
<|endoftext|>


<commit_before>import argparse as ap
import gc
import jsonpickle
import numpy as pd
import os
import pandas as pd
import sys

from itertools import chain
from jsonpickle.ext import numpy as jsonpickle_numpy

from keras import callbacks
from keras.models import Sequential, load_model
from keras.layers import BatchNormalization, Dense, Dropout, PReLU
from keras.wrappers.scikit_learn import KerasClassifier

from tabulate import tabulate

from sklearn import preprocessing
from sklearn.externals import joblib
from sklearn.metrics import auc, precision_recall_fscore_support, roc_curve
from sklearn.model_selection import train_test_split
from sklearn.utils import shuffle

class Trainer(object):
    def __init__(self, **kwargs):
        if bool(kwargs):
            self.args = self.get_args(self.to_args_list(kwargs))

        else:
            self.args = self.get_args()

    def get_args(self, args=None):
        playdata = 'data/Xy.csv'

        parser = ap.ArgumentParser(description='Train wincast')

        parser.add_argument('--playdata', '-p', type=str, default=playdata)
        parser.add_argument('--nb-epoch', '-e', type=int, default=2)
        parser.add_argument('--verbose', '-v', default=False, action='store_true')
        parser.add_argument('--batch-size', '-b', type=int, default=128)
        parser.add_argument('--outdir', '-o', default=False)
        parser.add_argument('--indir', '-i', default=False)
        parser.add_argument('--monitor', '-m', action='store_true', default=False)
        parser.add_argument('--evaluate', action='store_true', default=False)
        parser.add_argument('--format', default='table', choices=('table', 'csv'))
        parser.add_argument('--headers', action='store_true', default=False)
        parser.add_argument('--board', action='store_true', default=False)

        return parser.parse_args(args=args)

    def to_args_list(self, args):
        args = [['--%s' % key, value] for key, value in args.items()]

        return chain.from_iterable(args)

    def create_baseline(self):

        # create model
        model = Sequential()

        model.add(Dense(30, input_dim=13, init='normal', activation='relu'))
        model.add(BatchNormalization())
        model.add(PReLU())
        model.add(Dropout(0.5))

        model.add(Dense(15, init='normal', activation='relu'))
        model.add(BatchNormalization())
        model.add(PReLU())
        model.add(Dropout(0.5))

        model.add(Dense(1, init='normal', activation='sigmoid'))

        # Compile model
        model.compile(loss='binary_crossentropy', optimizer='adam', metrics=[
            'accuracy', 'precision', 'recall'])

        return model

    def read(self):
        self.model = load_model(os.path.normpath(os.path.join(
            self.args.indir, 'wincast.model.h5')))
        self.scaler = joblib.load(os.path.normpath(os.path.join(
            self.args.indir, 'wincast.scaler.pkl')))

    def get_features(self, X):
        return X.loc[:, [
            'qtr',
            'min',
            'sec',
            'ptso',
            'ptsd',
            'timo',
            'timd',
            'dwn',
            'ytg',
            'yfog',
            'ou',
            'pts_s',
            'off_h',
        ]]

    def train(self):
        if self.args.indir:
            self.read()
            gc.collect()
            return

        self.scaler = preprocessing.StandardScaler()

        X = pd.read_csv(self.args.playdata)
        X_train = X.loc[~X['seas'].isin([2000, 2007, 2015])]
        X_test = X.loc[X['seas'].isin([2000, 2007, 2015])]

        y_train = X_train.loc[:, 'y']
        y_test = X_test.loc[:, 'y']

        X_train = self.get_features(X_train)
        X_test = self.get_features(X_test)

        X_train = self.scaler.fit_transform(X_train, y_train)
        X_test = self.scaler.transform(X_test)

        X_train, y_train = shuffle(X_train, y_train)

        self.X_train, self.y_train = X_train, y_train
        self.X_test, self.y_test = X_test, y_test

        self.model = self.create_baseline()

        train_callbacks = []

        if self.args.evaluate:
            train_callbacks.append(callbacks.CSVLogger('evaluate/training.csv'))

        if self.args.board:
            tb_cb = callbacks.TensorBoard(
                log_dir='evaluate/board/', histogram_freq=0,
                write_graph=True, write_images=False)
            train_callbacks.append(tb_cb)

        self.model.fit(X_train, y_train.as_matrix(),
                       batch_size=self.args.batch_size,
                       nb_epoch=self.args.nb_epoch,
                       verbose=self.args.verbose,
                       callbacks=train_callbacks)

        if self.args.evaluate:
            self.evaluate(X_test, y_test)

        if self.args.outdir:
            self.write()

        gc.collect()

    def evaluate(self, X, y):
        evaluation = self.model.evaluate(X, y.as_matrix())

        with open('evaluate/cv.csv', 'w+') as cv:
            cv.write(','.join(self.model.metrics_names))
            cv.write(os.linesep)
            cv.write(','.join(map(str, evaluation)))

    def _check_model(self):
        if not getattr(self, 'model', None):
            raise AttributeError('Model has not been trained, call train() before predicting')

    def predict(self, X, *args, **kwargs):
        self._check_model()

        return self.model.predict(self.scaler.transform(X), *args, **kwargs)

    def predict_proba(self, X, *args, **kwargs):
        self._check_model()

        return self.model.predict_proba(self.scaler.transform(X), **kwargs)

    def write(self):
        self.model.save(os.path.normpath(
            os.path.join(self.args.outdir, 'wincast.model.h5')))
        joblib.dump(self.scaler, os.path.normpath(
            os.path.join(self.args.outdir, 'wincast.scaler.pkl')))

if __name__ == '__main__':
    Trainer().train()
<commit_msg>Use a deeper and wider architecture<commit_after>import argparse as ap
import gc
import jsonpickle
import numpy as pd
import os
import pandas as pd
import sys

from itertools import chain
from jsonpickle.ext import numpy as jsonpickle_numpy

from keras import callbacks
from keras.models import Sequential, load_model
from keras.layers import BatchNormalization, Dense, Dropout, PReLU
from keras.wrappers.scikit_learn import KerasClassifier

from tabulate import tabulate

from sklearn import preprocessing
from sklearn.externals import joblib
from sklearn.metrics import auc, precision_recall_fscore_support, roc_curve
from sklearn.model_selection import train_test_split
from sklearn.utils import shuffle

class Trainer(object):
    def __init__(self, **kwargs):
        if bool(kwargs):
            self.args = self.get_args(self.to_args_list(kwargs))

        else:
            self.args = self.get_args()

    def get_args(self, args=None):
        playdata = 'data/Xy.csv'

        parser = ap.ArgumentParser(description='Train wincast')

        parser.add_argument('--playdata', '-p', type=str, default=playdata)
        parser.add_argument('--nb-epoch', '-e', type=int, default=2)
        parser.add_argument('--verbose', '-v', default=False, action='store_true')
        parser.add_argument('--batch-size', '-b', type=int, default=128)
        parser.add_argument('--outdir', '-o', default=False)
        parser.add_argument('--indir', '-i', default=False)
        parser.add_argument('--monitor', '-m', action='store_true', default=False)
        parser.add_argument('--evaluate', action='store_true', default=False)
        parser.add_argument('--format', default='table', choices=('table', 'csv'))
        parser.add_argument('--headers', action='store_true', default=False)
        parser.add_argument('--board', action='store_true', default=False)

        return parser.parse_args(args=args)

    def to_args_list(self, args):
        args = [['--%s' % key, value] for key, value in args.items()]

        return chain.from_iterable(args)

    def create_baseline(self):

        # create model
        model = Sequential()

        model.add(Dense(64, input_dim=13, init='normal', activation='relu'))
        model.add(BatchNormalization())
        model.add(PReLU())
        model.add(Dropout(0.5))

        model.add(Dense(32, init='normal', activation='relu'))
        model.add(BatchNormalization())
        model.add(PReLU())
        model.add(Dropout(0.5))

        model.add(Dense(16, init='normal', activation='relu'))
        model.add(BatchNormalization())
        model.add(PReLU())
        model.add(Dropout(0.5))

        model.add(Dense(1, init='normal', activation='sigmoid'))

        # Compile model
        model.compile(loss='binary_crossentropy', optimizer='adam', metrics=[
            'accuracy', 'precision', 'recall'])

        return model

    def read(self):
        self.model = load_model(os.path.normpath(os.path.join(
            self.args.indir, 'wincast.model.h5')))
        self.scaler = joblib.load(os.path.normpath(os.path.join(
            self.args.indir, 'wincast.scaler.pkl')))

    def get_features(self, X):
        return X.loc[:, [
            'qtr',
            'min',
            'sec',
            'ptso',
            'ptsd',
            'timo',
            'timd',
            'dwn',
            'ytg',
            'yfog',
            'ou',
            'pts_s',
            'off_h',
        ]]

    def train(self):
        if self.args.indir:
            self.read()
            gc.collect()
            return

        self.scaler = preprocessing.StandardScaler()

        X = pd.read_csv(self.args.playdata)
        X_train = X.loc[~X['seas'].isin([2001, 2008, 2016])]
        X_test = X.loc[X['seas'].isin([2001, 2008, 2016])]

        y_train = X_train.loc[:, 'y']
        y_test = X_test.loc[:, 'y']

        X_train = self.get_features(X_train)
        X_test = self.get_features(X_test)

        X_train = self.scaler.fit_transform(X_train, y_train)
        X_test = self.scaler.transform(X_test)

        X_train, y_train = shuffle(X_train, y_train)

        self.X_train, self.y_train = X_train, y_train
        self.X_test, self.y_test = X_test, y_test

        self.model = self.create_baseline()

        train_callbacks = []

        if self.args.evaluate:
            train_callbacks.append(callbacks.CSVLogger('evaluate/training.csv'))

        if self.args.board:
            tb_cb = callbacks.TensorBoard(
                log_dir='evaluate/board/', histogram_freq=0,
                write_graph=True, write_images=False)
            train_callbacks.append(tb_cb)

        self.model.fit(X_train, y_train.as_matrix(),
                       batch_size=self.args.batch_size,
                       nb_epoch=self.args.nb_epoch,
                       verbose=self.args.verbose,
                       callbacks=train_callbacks)

        if self.args.evaluate:
            self.evaluate(X_test, y_test)

        if self.args.outdir:
            self.write()

        gc.collect()

    def evaluate(self, X, y):
        evaluation = self.model.evaluate(X, y.as_matrix())

        with open('evaluate/cv.csv', 'w+') as cv:
            cv.write(','.join(self.model.metrics_names))
            cv.write(os.linesep)
            cv.write(','.join(map(str, evaluation)))

    def _check_model(self):
        if not getattr(self, 'model', None):
            raise AttributeError('Model has not been trained, call train() before predicting')

    def predict(self, X, *args, **kwargs):
        self._check_model()

        return self.model.predict(self.scaler.transform(X), *args, **kwargs)

    def predict_proba(self, X, *args, **kwargs):
        self._check_model()

        return self.model.predict_proba(self.scaler.transform(X), **kwargs)

    def write(self):
        self.model.save(os.path.normpath(
            os.path.join(self.args.outdir, 'wincast.model.h5')))
        joblib.dump(self.scaler, os.path.normpath(
            os.path.join(self.args.outdir, 'wincast.scaler.pkl')))

if __name__ == '__main__':
    Trainer().train()
<|endoftext|>


<commit_before>from datetime import datetime
import hashlib
from django.conf import settings
import uuid
from dimagi.utils.mixins import UnicodeMixIn
try:
    import simplejson
except ImportError:
    from django.utils import simplejson

from couchforms.models import XFormInstance, XFormDuplicate, XFormError, XFormDeprecated,\
    SubmissionErrorLog
from dimagi.utils.logging import log_exception
import logging
from couchdbkit.resource import RequestFailed
from couchforms.exceptions import CouchFormException
from couchforms.signals import xform_saved
from dimagi.utils.couch import uid
import re
from dimagi.utils.post import post_authenticated_data, post_unauthenticated_data
from restkit.errors import ResourceNotFound
from lxml import etree

class SubmissionError(Exception, UnicodeMixIn):
    """
    When something especially bad goes wrong during a submission, this
    exception gets raised.
    """

    def __init__(self, error_log, *args, **kwargs):
        super(SubmissionError, self).__init__(*args, **kwargs)
        self.error_log = error_log

    def __str__(self):
        return str(self.error_log)

def post_from_settings(instance, extras={}):
    url = settings.XFORMS_POST_URL if not extras else "%s?%s" % \
        (settings.XFORMS_POST_URL, "&".join(["%s=%s" % (k, v) for k, v in extras.items()]))
    if settings.COUCH_USERNAME:
        return post_authenticated_data(instance, url,
                                       settings.COUCH_USERNAME,
                                       settings.COUCH_PASSWORD)
    else:
        return post_unauthenticated_data(instance, url)

def post_xform_to_couch(instance, attachments={}):
    """
    Post an xform to couchdb, based on the settings.XFORMS_POST_URL.
    Returns the newly created document from couchdb, or raises an
    exception if anything goes wrong.

    attachments is a dictionary of the request.FILES that are not the xform.  Key is the parameter name, and the value is the django MemoryFile object stream.
    """
    # check settings for authentication credentials
    try:
        response, errors = post_from_settings(instance)
        if not _has_errors(response, errors):
            doc_id = response
            try:
                xform = XFormInstance.get(doc_id)
                #put attachments onto the saved xform instance
                for key, val in attachments.items():
                    res = xform.put_attachment(val, name=key, content_type=val.content_type, content_length=val.size)

                # fire signals
                # We don't trap any exceptions here. This is by design.
                # If something fails (e.g. case processing), we quarantine the
                # form into an error location.
                xform_saved.send(sender="couchforms", xform=xform)
                return xform
            except Exception, e:
                logging.error("Problem with form %s" % doc_id)
                logging.exception(e)
                # "rollback" by changing the doc_type to XFormError
                try:
                    bad = XFormError.get(doc_id)
                    bad.problem = "%s" % e
                    bad.save()
                    return bad
                except ResourceNotFound:
                    pass # no biggie, the failure must have been in getting it back
                raise
        else:
            raise CouchFormException("Problem POSTing form to couch! errors/response: %s/%s" % (errors, response))
    except RequestFailed, e:
        if e.status_int == 409:
            # this is an update conflict, i.e. the uid in the form was the same.
            return _handle_id_conflict(instance, attachments)
        else:
            doc = _log_hard_failure(instance, attachments, e)
            raise SubmissionError(doc)

def _has_errors(response, errors):
    return errors or "error" in response

def _handle_id_conflict(instance, attachments):
    """
    For id conflicts, we check if the files contain exactly the same content,
    If they do, we just log this as a dupe. If they don't, we deprecate the
    previous form and overwrite it with the new form's contents.
    """
    def _extract_id_from_raw_xml(xml):

        # this is the standard openrosa way of doing things
        parsed = etree.XML(xml)
        meta_ns = "http://openrosa.org/jr/xforms"
        val = parsed.find("{%(ns)s}meta/{%(ns)s}instanceID" % \
                          {"ns": meta_ns})
        if val is not None and val.text:
            return val.text

        # if we get here search more creatively for some of the older
        # formats
        _PATTERNS = (r"<instanceID>([\w-]+)</instanceID>",
                     r"<uid>([\w-]+)</uid>",
                     r"<uuid>([\w-]+)</uuid>")
        for pattern in _PATTERNS:
            if re.search(pattern, xml):
                return re.search(pattern, xml).groups()[0]

        logging.error("Unable to find conflicting matched uid in form: %s" % xml)
        return ""

    conflict_id = _extract_id_from_raw_xml(instance)

    # get old document
    existing_doc = XFormInstance.get(conflict_id)

    # compare md5s
    existing_md5 = existing_doc.xml_md5()
    new_md5 = hashlib.md5(instance).hexdigest()

    # if not same:
    # Deprecate old form (including changing ID)
    # to deprecate, copy new instance into a XFormDeprecated
    if existing_md5 != new_md5:
        doc_copy = XFormInstance.get_db().copy_doc(conflict_id)
        xfd = XFormDeprecated.get(doc_copy['id'])
        xfd.orig_id = conflict_id
        xfd.doc_type=XFormDeprecated.__name__
        xfd.save()

        # after that delete the original document and resubmit.
        XFormInstance.get_db().delete_doc(conflict_id)
        return post_xform_to_couch(instance, attachments=attachments)
    else:
        # follow standard dupe handling
        new_doc_id = uid.new()
        response, errors = post_from_settings(instance, {"uid": new_doc_id})
        if not _has_errors(response, errors):
            # create duplicate doc
            # get and save the duplicate to ensure the doc types are set correctly
            # so that it doesn't show up in our reports
            dupe = XFormDuplicate.get(response)
            dupe.problem = "Form is a duplicate of another! (%s)" % conflict_id
            dupe.save()
            return dupe
        else:
            # how badly do we care about this?
            raise CouchFormException("Problem POSTing form to couch! errors/response: %s/%s" % (errors, response))

def _log_hard_failure(instance, attachments, error):
    """
    Handle's a hard failure from posting a form to couch.

    Currently, it will save the raw payload to couch in a hard-failure doc
    and return that doc.
    """
    return SubmissionErrorLog.from_instance(instance, str(error))

def value_for_display(value, replacement_chars="_-"):
    """
    Formats an xform value for display, replacing the contents of the
    system characters with spaces
    """
    for char in replacement_chars:
        value = str(value).replace(char, " ")
    return value<commit_msg>delete dead comment<commit_after>from datetime import datetime
import hashlib
from django.conf import settings
import uuid
from dimagi.utils.mixins import UnicodeMixIn
try:
    import simplejson
except ImportError:
    from django.utils import simplejson

from couchforms.models import XFormInstance, XFormDuplicate, XFormError, XFormDeprecated,\
    SubmissionErrorLog
from dimagi.utils.logging import log_exception
import logging
from couchdbkit.resource import RequestFailed
from couchforms.exceptions import CouchFormException
from couchforms.signals import xform_saved
from dimagi.utils.couch import uid
import re
from dimagi.utils.post import post_authenticated_data, post_unauthenticated_data
from restkit.errors import ResourceNotFound
from lxml import etree

class SubmissionError(Exception, UnicodeMixIn):
    """
    When something especially bad goes wrong during a submission, this
    exception gets raised.
    """

    def __init__(self, error_log, *args, **kwargs):
        super(SubmissionError, self).__init__(*args, **kwargs)
        self.error_log = error_log

    def __str__(self):
        return str(self.error_log)

def post_from_settings(instance, extras={}):
    url = settings.XFORMS_POST_URL if not extras else "%s?%s" % \
        (settings.XFORMS_POST_URL, "&".join(["%s=%s" % (k, v) for k, v in extras.items()]))
    if settings.COUCH_USERNAME:
        return post_authenticated_data(instance, url,
                                       settings.COUCH_USERNAME,
                                       settings.COUCH_PASSWORD)
    else:
        return post_unauthenticated_data(instance, url)

def post_xform_to_couch(instance, attachments={}):
    """
    Post an xform to couchdb, based on the settings.XFORMS_POST_URL.
    Returns the newly created document from couchdb, or raises an
    exception if anything goes wrong.

    attachments is a dictionary of the request.FILES that are not the xform.  Key is the parameter name, and the value is the django MemoryFile object stream.
    """
    try:
        response, errors = post_from_settings(instance)
        if not _has_errors(response, errors):
            doc_id = response
            try:
                xform = XFormInstance.get(doc_id)
                #put attachments onto the saved xform instance
                for key, val in attachments.items():
                    res = xform.put_attachment(val, name=key, content_type=val.content_type, content_length=val.size)

                # fire signals
                # We don't trap any exceptions here. This is by design.
                # If something fails (e.g. case processing), we quarantine the
                # form into an error location.
                xform_saved.send(sender="couchforms", xform=xform)
                return xform
            except Exception, e:
                logging.error("Problem with form %s" % doc_id)
                logging.exception(e)
                # "rollback" by changing the doc_type to XFormError
                try:
                    bad = XFormError.get(doc_id)
                    bad.problem = "%s" % e
                    bad.save()
                    return bad
                except ResourceNotFound:
                    pass # no biggie, the failure must have been in getting it back
                raise
        else:
            raise CouchFormException("Problem POSTing form to couch! errors/response: %s/%s" % (errors, response))
    except RequestFailed, e:
        if e.status_int == 409:
            # this is an update conflict, i.e. the uid in the form was the same.
            return _handle_id_conflict(instance, attachments)
        else:
            doc = _log_hard_failure(instance, attachments, e)
            raise SubmissionError(doc)

def _has_errors(response, errors):
    return errors or "error" in response

def _handle_id_conflict(instance, attachments):
    """
    For id conflicts, we check if the files contain exactly the same content,
    If they do, we just log this as a dupe. If they don't, we deprecate the
    previous form and overwrite it with the new form's contents.
    """
    def _extract_id_from_raw_xml(xml):

        # this is the standard openrosa way of doing things
        parsed = etree.XML(xml)
        meta_ns = "http://openrosa.org/jr/xforms"
        val = parsed.find("{%(ns)s}meta/{%(ns)s}instanceID" % \
                          {"ns": meta_ns})
        if val is not None and val.text:
            return val.text

        # if we get here search more creatively for some of the older
        # formats
        _PATTERNS = (r"<instanceID>([\w-]+)</instanceID>",
                     r"<uid>([\w-]+)</uid>",
                     r"<uuid>([\w-]+)</uuid>")
        for pattern in _PATTERNS:
            if re.search(pattern, xml):
                return re.search(pattern, xml).groups()[0]

        logging.error("Unable to find conflicting matched uid in form: %s" % xml)
        return ""

    conflict_id = _extract_id_from_raw_xml(instance)

    # get old document
    existing_doc = XFormInstance.get(conflict_id)

    # compare md5s
    existing_md5 = existing_doc.xml_md5()
    new_md5 = hashlib.md5(instance).hexdigest()

    # if not same:
    # Deprecate old form (including changing ID)
    # to deprecate, copy new instance into a XFormDeprecated
    if existing_md5 != new_md5:
        doc_copy = XFormInstance.get_db().copy_doc(conflict_id)
        xfd = XFormDeprecated.get(doc_copy['id'])
        xfd.orig_id = conflict_id
        xfd.doc_type=XFormDeprecated.__name__
        xfd.save()

        # after that delete the original document and resubmit.
        XFormInstance.get_db().delete_doc(conflict_id)
        return post_xform_to_couch(instance, attachments=attachments)
    else:
        # follow standard dupe handling
        new_doc_id = uid.new()
        response, errors = post_from_settings(instance, {"uid": new_doc_id})
        if not _has_errors(response, errors):
            # create duplicate doc
            # get and save the duplicate to ensure the doc types are set correctly
            # so that it doesn't show up in our reports
            dupe = XFormDuplicate.get(response)
            dupe.problem = "Form is a duplicate of another! (%s)" % conflict_id
            dupe.save()
            return dupe
        else:
            # how badly do we care about this?
            raise CouchFormException("Problem POSTing form to couch! errors/response: %s/%s" % (errors, response))

def _log_hard_failure(instance, attachments, error):
    """
    Handle's a hard failure from posting a form to couch.

    Currently, it will save the raw payload to couch in a hard-failure doc
    and return that doc.
    """
    return SubmissionErrorLog.from_instance(instance, str(error))

def value_for_display(value, replacement_chars="_-"):
    """
    Formats an xform value for display, replacing the contents of the
    system characters with spaces
    """
    for char in replacement_chars:
        value = str(value).replace(char, " ")
    return value<|endoftext|>


<commit_before>from app import app
from flask.ext.sqlalchemy import SQLAlchemy
from flask.ext.script import Manager, prompt_bool
from datetime import datetime

db = SQLAlchemy(app)

manager = Manager(usage="Manage the database")

@manager.command
def create():
    """Create the database"""
    db.create_all()

@manager.command
def drop():
    """Empty the database"""
    if prompt_bool("Are you sure you want to drop all tables from the database?"):
        db.drop_all()

@manager.command
def recreate():
    """Recreate the database"""
    drop()
    create()

class Domain(db.Model):
    """
    A domain name
    """
    __tablename__ = 'domain'

    id = db.Column(db.Integer, primary_key=True)
    name = db.Column(db.Text(length=255), unique=True)

    def __init__(self, name):
        self.name = name

    def toObject(self):
        return {
            "id": self.id,
            "name": self.name
        }

class Address(db.Model):
    """
    An email address
    """
    __tablename__ = 'address'

    id = db.Column(db.Integer, primary_key=True)
    local = db.Column(db.Text)
    domain_id = db.Column(db.Integer, db.ForeignKey('domain.id'))
    domain = db.relationship('Domain')
    total_spam_score = db.Column(db.Float, default=0)
    total_received = db.Column(db.Integer, default=0)
    date_last_received = db.Column(db.DateTime(timezone=False), default=datetime.utcnow)
    created = db.Column(db.DateTime(timezone=False), default=datetime.utcnow)

    def __init__(self, local, domain_id):
        self.local = local
        self.domain_id = domain_id

    def toObject(self):
        return {
            "id": self.id,
            "local": self.local,
            "domain": self.domain.name,
            "created": self.created.strftime('%Y-%m-%d %H:%M:%S'),
            "avg_spam_score": self.total_spam_score / self.total_received,
            "date_last_received": self.date_last_received.strftime('%Y-%m-%d %H:%M:%S'),
            "total_received": self.total_received
        }
<commit_msg>Use MySQL VARCHAR<commit_after>from app import app
from flask.ext.sqlalchemy import SQLAlchemy
from sqlalchemy.dialects.mysql import VARCHAR
from flask.ext.script import Manager, prompt_bool
from datetime import datetime

db = SQLAlchemy(app)

manager = Manager(usage="Manage the database")

@manager.command
def create():
    """Create the database"""
    db.create_all()

@manager.command
def drop():
    """Empty the database"""
    if prompt_bool("Are you sure you want to drop all tables from the database?"):
        db.drop_all()

@manager.command
def recreate():
    """Recreate the database"""
    drop()
    create()

class Domain(db.Model):
    """
    A domain name
    """
    __tablename__ = 'domain'

    id = db.Column(db.Integer, primary_key=True)
    name = db.Column(VARCHAR(255), unique=True)

    def __init__(self, name):
        self.name = name

    def toObject(self):
        return {
            "id": self.id,
            "name": self.name
        }

class Address(db.Model):
    """
    An email address
    """
    __tablename__ = 'address'

    id = db.Column(db.Integer, primary_key=True)
    local = db.Column(db.Text)
    domain_id = db.Column(db.Integer, db.ForeignKey('domain.id'))
    domain = db.relationship('Domain')
    total_spam_score = db.Column(db.Float, default=0)
    total_received = db.Column(db.Integer, default=0)
    date_last_received = db.Column(db.DateTime(timezone=False), default=datetime.utcnow)
    created = db.Column(db.DateTime(timezone=False), default=datetime.utcnow)

    def __init__(self, local, domain_id):
        self.local = local
        self.domain_id = domain_id

    def toObject(self):
        return {
            "id": self.id,
            "local": self.local,
            "domain": self.domain.name,
            "created": self.created.strftime('%Y-%m-%d %H:%M:%S'),
            "avg_spam_score": self.total_spam_score / self.total_received,
            "date_last_received": self.date_last_received.strftime('%Y-%m-%d %H:%M:%S'),
            "total_received": self.total_received
        }
<|endoftext|>


<commit_before>import logging
from mwoauth import ConsumerToken, Handshaker, AccessToken
import re

from django.conf import settings
from django.contrib import messages
from django.contrib.auth import login, authenticate
from django.contrib.auth.models import User
from django.core.exceptions import PermissionDenied, DisallowedHost
from django.core.urlresolvers import reverse_lazy
from django.http import HttpResponseRedirect
from django.views.generic.base import View
from django.views.generic.edit import FormView
from django.utils.translation import ugettext as _

from .helpers.wiki_list import WIKI_DICT
from .forms import HomeWikiForm
from .models import Editor

logger = logging.getLogger(__name__)

try:
    ALLOWED_DOMAINS = settings.WP_CREDENTIALS.keys()
    # Construct a "consumer" from the key/secret provided by MediaWiki for all
    # available site domains. (Each domain requires a different set of
    # credentials.)
    CONSUMER_TOKENS = {domain: ConsumerToken(creds['key'], creds['secret'])
                       for domain, creds in settings.WP_CREDENTIALS.items()}
except AttributeError:
    # We'll get an AttributeError if the settings don't contain WP_CREDENTIALS.
    # This is the expected behavior for systems that don't have Wikipedia
    # OAuth credentials, like localhost. The system should not crash in that
    # case - it just shouldn't generate any tokens.
    ALLOWED_DOMAINS = []
    CONSUMER_TOKENS = {}

# Must be before handshaker construction
def _get_full_wiki_url(home_wiki):
    """
    Given something like 'en.wikipedia.org', return something like
    'https://en.wikipedia.org/w/index.php'.
    """
    return 'https://meta.wikimedia.org/w/index.php'

# Construct all conceivably needed handshakers. Will result in a dict like
# {domain: {wiki url: handshaker}}.
HANDSHAKERS = {}

WIKI_DOMAINS = WIKI_DICT.values()

for allowed_domain in ALLOWED_DOMAINS:
    tempdict = {}
    for wiki in WIKI_DOMAINS:
        tempdict[wiki] = Handshaker(
            _get_full_wiki_url(wiki), CONSUMER_TOKENS[allowed_domain])
    HANDSHAKERS[allowed_domain] = tempdict

def _get_token(domain):
    try:
        assert domain in ALLOWED_DOMAINS
    except AssertionError:
        logger.exception('Attempted to get a token for an unsupported domain')
        raise

    return CONSUMER_TOKENS[domain]

def _get_handshaker(domain, home_wiki):
    try:
        assert domain in ALLOWED_DOMAINS
    except AssertionError:
        logger.exception('Attempted to get a handshaker for an unsupported TWL domain')
        raise

    try:
        assert home_wiki in WIKI_DOMAINS
    except AssertionError:
        logger.exception('Attempted to get a handshaker for an unsupported wiki')
        raise

    return HANDSHAKERS[domain][home_wiki]

def _dehydrate_token(token):
    """
    Convert the request token into a dict suitable for storing in the session.
    """
    session_token = {}
    session_token['key'] = token.key
    session_token['secret'] = token.secret
    return session_token

def _rehydrate_token(token):
    """
    Convert the stored dict back into a request token that we can use for
    getting an access grant.
    """
    request_token = ConsumerToken(token['key'], token['secret'])
    return request_token

class OAuthBackend(object):

    def _get_language_code(self, identity):
        logger.info('Getting language code...')
        home_wiki_url = identity['iss']
        extractor = re.match(r'(https://)?(\w+).wikim|pedia.org', home_wiki_url)
        language_code = extractor.group(2)
        try:
            assert language_code in WIKI_DICT
        except AssertionError:
            logger.exception('Could not find code {lang} in WIKI_DICT; '
                'exiting'.format(lang=language_code))
            raise

        logger.info('Language code was {lang}.'.format(lang=language_code))

        return language_code

    def _get_username(self, identity):
        # The Wikipedia username is globally unique, but Wikipedia allows it to
        # have characters that the Django username system rejects. However,
        # wiki userID should be unique, and limited to ASCII.
        return '{sub}'.format(sub=identity['sub'])

    def _meets_minimum_requirement(self, identity, global_userinfo):
        """
        This needs to be reworked to actually check against global_userinfo.
        """
        if 'autoconfirmed' in identity['rights']:
            return True

        return False

    def _create_user_and_editor(self, identity, userinfo, global_userinfo):
        # This can't be super informative because we don't want to log
        # identities.
        logger.info('Creating user.')

        if not self._meets_minimum_requirement(identity, global_userinfo):
            # Don't create a User or Editor if this person does not meet the
            # minimum account quality requirement. It would be nice to provide
            # some user feedback here, but we can't; exception messages don't
            # get passed on as template context in Django 1.8. (They do in
            # 1.10, so this can be revisited in future.)
            logger.warning('User did not meet minimum requirements; not created.')
            raise PermissionDenied

        # This will assert that the language code is a real Wikipedia, which
        # is good - we want to verify that assumption before proceeding.
        language_code = self._get_language_code(identity)

        # -------------------------- Create the user ---------------------------
        try:
            email = userinfo['email']
        except KeyError:
            email = None

        username = self._get_username(identity)

        # Since we are not providing a password argument, this will call
        # set_unusable_password, which is exactly what we want; users created
        # via OAuth should only be allowed to log in via OAuth.
        user = User.objects.create_user(username=username, email=email)

        # ------------------------- Create the editor --------------------------
        logger.info('Creating editor'.format(username=username))
        editor = Editor()

        editor.user = user
        editor.wp_sub = identity['sub']
        editor.home_wiki = language_code
        editor.update_from_wikipedia(identity) # This call also saves the editor

        logger.info('User and editor successfully created.')
        return user, editor

    def _get_and_update_user_from_identity(self, identity, userinfo, global_userinfo):
        """
        If we have an Editor and User matching the identity returned by
        Wikipedia, update the editor with the identity parameters and return its
        associated user. If we don't, create an Editor and User, and return that
        user.

        If the wikipedia account does not meet our eligibility criteria, create
        a TWLight account if needed, but set it as inactive. Also deactivate
        any existing accounts that have become ineligible.

        Also return a boolean that is True if we created a user during this
        call and False if we did not.
        """
        logger.info('Attempting to update editor after OAuth login.')
        try:
            username = self._get_username(identity)
            user = User.objects.get(username=username)

            # This login path should only be used for accounts created via
            # Wikipedia login, which all have editor objects.
            assert hasattr(user, 'editor')
            editor = user.editor

            editor.update_from_wikipedia(identity)
            logger.info('Editor {editor} updated.'.format(editor=editor))

            created = False

        except User.DoesNotExist:
            logger.info("Can't find user; creating one.")
            user, editor = self._create_user_and_editor(identity, userinfo, global_userinfo)
            created = True

        except AttributeError:
            logger.warning('A user tried using the Wikipedia OAuth '
                'login path but does not have an attached editor'.format(
                    username=identity['username']))
            raise

        logger.info('User and editor created/updated after OAuth login.')
        return user, created

    def authenticate(self, request=None, access_token=None, handshaker=None):
        logger.info('Authenticating user...')
        if not request or not access_token or not handshaker:
            logger.info('Missing OAuth authentication elements; falling back'
                'to another authentication method.')
            # You must have meant to use a different authentication backend.
            # Returning None will make Django keep going down its list of
            # options.
            return None

        try:
            assert isinstance(access_token, AccessToken)
        except AssertionError:
            logger.exception('Did not have a properly formed AccessToken')
            return None

        # Get identifying information about the user. This doubles as a way
        # to authenticate the access token, which only Wikimedia can do,
        # and thereby to authenticate the user (which is hard for us to do as
        # we have no password.)
        logger.info('Identifying user...')
        try:
            identity = handshaker.identify(access_token)
        except:
            logger.warning('Someone tried to log in but presented an invalid '
                'access token.')
            raise PermissionDenied

        logger.info('Getting local user info...')
        try:
            editor = Editor()
            userinfo = editor.get_userinfo(identity)
        except:
            logger.warning('Unable to get local user information.')
            raise

        logger.info('Getting global user info...')
        try:
            editor = Editor()
            global_userinfo = editor.get_global_userinfo(identity)
        except:
            logger.warning('Unable to get global user information.')
            raise

        # Get or create the user.
        logger.info('User has been identified; getting or creating user.')
        user, created = self._get_and_update_user_from_identity(identity, userinfo, global_userinfo)

        if created:
            home_wiki = re.search(r'(\w+).wikim|pedia.org',
                                 handshaker.mw_uri).group(1)

            try:
                # It's actually a wiki, right?
                assert home_wiki in WIKI_DICT
                user.editor.home_wiki = home_wiki
                user.editor.save()
            except AssertionError:
                # Site functionality mostly works if people don't
                # declare a homewiki. There are some broken bits, like the SUL
                # link, but users can set their homewiki, or admins can do it
                # in the admin interface.
                pass
        else:
            logger.info('User has been updated.'.format(user=user))

        request.session['user_created'] = created

        # The authenticate() function of a Django auth backend must return
        # the user.
        return user

    def get_user(self, user_id):
        try:
            return User.objects.get(pk=user_id)
        except User.DoesNotExist:
            logger.warning('There is no user {user_id}'.format(user_id=user_id))
            return None

class OAuthInitializeView(FormView):
    """
    Ask Wikipedia for a temporary key/secret for the user, and redirect
    them to their home Wikipedia to confirm authorization.
    """
    template_name = 'users/oauth_login.html'
    form_class = HomeWikiForm

    def form_valid(self, form):
        # The form returns the code for the home wiki (the first element of the
        # tuple in WIKIS), but we want the URL, so we grab it from WIKI_DICT.
        home_wiki = WIKI_DICT[form.cleaned_data['home_wiki']]

        # The site might be running under multiple URLs, so find out the current
        # one (and make sure it's legit).
        # The Sites framework was designed for different URLs that correspond to
        # different databases or functionality - it's not a good fit here.
        try:
            domain = self.request.get_host()
            assert domain in settings.ALLOWED_HOSTS # safety first!
        except (AssertionError, DisallowedHost):
            logger.exception()
            raise PermissionDenied

        logger.info('home_wiki was {home_wiki}, domain was {domain}'.format(home_wiki=home_wiki, domain=domain))

        # Get handshaker matching this wiki URL from our dict.
        handshaker = _get_handshaker(domain, home_wiki)
        logger.info('handshaker gotten')

        try:
            redirect, request_token = handshaker.initiate()
        except:
            logger.exception('Handshaker not initiated')
            raise

        logger.info('handshaker initiated')
        self.request.session['request_token'] = _dehydrate_token(request_token)
        self.request.session['home_wiki'] = home_wiki
        return HttpResponseRedirect(redirect)

class OAuthCallbackView(View):
    """
    Receive the redirect from Wikipedia and parse the response token.
    """

    def get(self, request, *args, **kwargs):
        response_qs = request.META['QUERY_STRING']

        # Get the handshaker. It should have already been constructed by
        # OAuthInitializeView.
        try:
            domain = self.request.get_host()
            assert domain in settings.ALLOWED_HOSTS
        except (AssertionError, DisallowedHost):
            logger.exception()
            raise PermissionDenied

        home_wiki = request.session.pop('home_wiki', None)

        try:
            handshaker = _get_handshaker(domain, home_wiki)
        except AssertionError:
            # get_handshaker will throw AssertionErrors for invalid data.
            logger.exception('Could not find handshaker')
            raise PermissionDenied

        # Get the request token, placed in session by OAuthInitializeView.
        session_token = request.session.pop('request_token', None)
        request_token = _rehydrate_token(session_token)

        if not request_token:
            logger.info('no request token :(')
            raise PermissionDenied

        # See if we can complete the OAuth process.
        try:
            access_token = handshaker.complete(request_token, response_qs)
        except:
            logger.exception('Access token generation failed :(')
            raise PermissionDenied

        user = authenticate(request=request,
                            access_token=access_token,
                            handshaker=handshaker)
        created = request.session.pop('user_created', False)

        if not user.is_active:
            # Do NOT log in the user.

            if created:
                messages.add_message(request, messages.WARNING,
                    _('Your Wikipedia account does not meet the eligibility '
                      'criteria in the terms of use, so your Wikipedia Library '
                      'Card Platform account cannot be activated.'))
            else:
                messages.add_message(request, messages.WARNING,
                    _('Either your Wikipedia Library Card Platform account has '
                      'been deactivated or your Wikipedia account no longer '
                      'meets the eligibility criteria in the terms of use, so '
                      'you cannot be logged in.'))

            return_url = reverse_lazy('terms')

        else:
            login(request, user)

            if created:
                # Translators: this message is displayed to users with brand new accounts.
                messages.add_message(request, messages.INFO, _('Welcome! '
                    'Please agree to the terms of use.'))
                return_url = reverse_lazy('terms')
            else:
                messages.add_message(request, messages.INFO, _('Welcome back!'))
                return_url = reverse_lazy('users:editor_detail',
                    kwargs={'pk': user.editor.pk})

        return HttpResponseRedirect(return_url)
<commit_msg>Strip out unnecessary references to userinfo.<commit_after>import logging
from mwoauth import ConsumerToken, Handshaker, AccessToken
import re

from django.conf import settings
from django.contrib import messages
from django.contrib.auth import login, authenticate
from django.contrib.auth.models import User
from django.core.exceptions import PermissionDenied, DisallowedHost
from django.core.urlresolvers import reverse_lazy
from django.http import HttpResponseRedirect
from django.views.generic.base import View
from django.views.generic.edit import FormView
from django.utils.translation import ugettext as _

from .helpers.wiki_list import WIKI_DICT
from .forms import HomeWikiForm
from .models import Editor

logger = logging.getLogger(__name__)

try:
    ALLOWED_DOMAINS = settings.WP_CREDENTIALS.keys()
    # Construct a "consumer" from the key/secret provided by MediaWiki for all
    # available site domains. (Each domain requires a different set of
    # credentials.)
    CONSUMER_TOKENS = {domain: ConsumerToken(creds['key'], creds['secret'])
                       for domain, creds in settings.WP_CREDENTIALS.items()}
except AttributeError:
    # We'll get an AttributeError if the settings don't contain WP_CREDENTIALS.
    # This is the expected behavior for systems that don't have Wikipedia
    # OAuth credentials, like localhost. The system should not crash in that
    # case - it just shouldn't generate any tokens.
    ALLOWED_DOMAINS = []
    CONSUMER_TOKENS = {}

# Must be before handshaker construction
def _get_full_wiki_url(home_wiki):
    """
    Given something like 'en.wikipedia.org', return something like
    'https://en.wikipedia.org/w/index.php'.
    """
    return 'https://meta.wikimedia.org/w/index.php'

# Construct all conceivably needed handshakers. Will result in a dict like
# {domain: {wiki url: handshaker}}.
HANDSHAKERS = {}

WIKI_DOMAINS = WIKI_DICT.values()

for allowed_domain in ALLOWED_DOMAINS:
    tempdict = {}
    for wiki in WIKI_DOMAINS:
        tempdict[wiki] = Handshaker(
            _get_full_wiki_url(wiki), CONSUMER_TOKENS[allowed_domain])
    HANDSHAKERS[allowed_domain] = tempdict

def _get_token(domain):
    try:
        assert domain in ALLOWED_DOMAINS
    except AssertionError:
        logger.exception('Attempted to get a token for an unsupported domain')
        raise

    return CONSUMER_TOKENS[domain]

def _get_handshaker(domain, home_wiki):
    try:
        assert domain in ALLOWED_DOMAINS
    except AssertionError:
        logger.exception('Attempted to get a handshaker for an unsupported TWL domain')
        raise

    try:
        assert home_wiki in WIKI_DOMAINS
    except AssertionError:
        logger.exception('Attempted to get a handshaker for an unsupported wiki')
        raise

    return HANDSHAKERS[domain][home_wiki]

def _dehydrate_token(token):
    """
    Convert the request token into a dict suitable for storing in the session.
    """
    session_token = {}
    session_token['key'] = token.key
    session_token['secret'] = token.secret
    return session_token

def _rehydrate_token(token):
    """
    Convert the stored dict back into a request token that we can use for
    getting an access grant.
    """
    request_token = ConsumerToken(token['key'], token['secret'])
    return request_token

class OAuthBackend(object):

    def _get_language_code(self, identity):
        logger.info('Getting language code...')
        home_wiki_url = identity['iss']
        extractor = re.match(r'(https://)?(\w+).wikim|pedia.org', home_wiki_url)
        language_code = extractor.group(2)
        try:
            assert language_code in WIKI_DICT
        except AssertionError:
            logger.exception('Could not find code {lang} in WIKI_DICT; '
                'exiting'.format(lang=language_code))
            raise

        logger.info('Language code was {lang}.'.format(lang=language_code))

        return language_code

    def _get_username(self, identity):
        # The Wikipedia username is globally unique, but Wikipedia allows it to
        # have characters that the Django username system rejects. However,
        # wiki userID should be unique, and limited to ASCII.
        return '{sub}'.format(sub=identity['sub'])

    def _meets_minimum_requirement(self, identity):
        """
        This needs to be reworked to actually check against global_userinfo.
        """
        if 'autoconfirmed' in identity['rights']:
            return True

        return False

    def _create_user_and_editor(self, identity):
        # This can't be super informative because we don't want to log
        # identities.
        logger.info('Creating user.')

        if not self._meets_minimum_requirement(identity):
            # Don't create a User or Editor if this person does not meet the
            # minimum account quality requirement. It would be nice to provide
            # some user feedback here, but we can't; exception messages don't
            # get passed on as template context in Django 1.8. (They do in
            # 1.10, so this can be revisited in future.)
            logger.warning('User did not meet minimum requirements; not created.')
            raise PermissionDenied

        # This will assert that the language code is a real Wikipedia, which
        # is good - we want to verify that assumption before proceeding.
        language_code = self._get_language_code(identity)

        # -------------------------- Create the user ---------------------------
        try:
            email = identity['email']
        except KeyError:
            email = None

        username = self._get_username(identity)

        # Since we are not providing a password argument, this will call
        # set_unusable_password, which is exactly what we want; users created
        # via OAuth should only be allowed to log in via OAuth.
        user = User.objects.create_user(username=username, email=email)

        # ------------------------- Create the editor --------------------------
        logger.info('Creating editor'.format(username=username))
        editor = Editor()

        editor.user = user

        editor.wp_sub = identity['sub']
        editor.home_wiki = language_code
        editor.update_from_wikipedia(identity) # This call also saves the editor

        logger.info('User and editor successfully created.')
        return user, editor

    def _get_and_update_user_from_identity(self, identity):
        """
        If we have an Editor and User matching the identity returned by
        Wikipedia, update the editor with the identity parameters and return its
        associated user. If we don't, create an Editor and User, and return that
        user.

        If the wikipedia account does not meet our eligibility criteria, create
        a TWLight account if needed, but set it as inactive. Also deactivate
        any existing accounts that have become ineligible.

        Also return a boolean that is True if we created a user during this
        call and False if we did not.
        """
        logger.info('Attempting to update editor after OAuth login.')
        try:
            username = self._get_username(identity)
            user = User.objects.get(username=username)

            # This login path should only be used for accounts created via
            # Wikipedia login, which all have editor objects.
            assert hasattr(user, 'editor')
            editor = user.editor

            editor.update_from_wikipedia(identity)
            logger.info('Editor {editor} updated.'.format(editor=editor))

            created = False

        except User.DoesNotExist:
            logger.info("Can't find user; creating one.")
            user, editor = self._create_user_and_editor(identity)
            created = True

        except AttributeError:
            logger.warning('A user tried using the Wikipedia OAuth '
                'login path but does not have an attached editor'.format(
                    username=identity['username']))
            raise

        logger.info('User and editor created/updated after OAuth login.')
        return user, created

    def authenticate(self, request=None, access_token=None, handshaker=None):
        logger.info('Authenticating user...')
        if not request or not access_token or not handshaker:
            logger.info('Missing OAuth authentication elements; falling back'
                'to another authentication method.')
            # You must have meant to use a different authentication backend.
            # Returning None will make Django keep going down its list of
            # options.
            return None

        try:
            assert isinstance(access_token, AccessToken)
        except AssertionError:
            logger.exception('Did not have a properly formed AccessToken')
            return None

        # Get identifying information about the user. This doubles as a way
        # to authenticate the access token, which only Wikimedia can do,
        # and thereby to authenticate the user (which is hard for us to do as
        # we have no password.)
        logger.info('Identifying user...')
        try:
            identity = handshaker.identify(access_token)
        except:
            logger.warning('Someone tried to log in but presented an invalid '
                'access token.')
            raise PermissionDenied

        # Get or create the user.
        logger.info('User has been identified; getting or creating user.')
        user, created = self._get_and_update_user_from_identity(identity)

        if created:
            home_wiki = re.search(r'(\w+).wikim|pedia.org',
                                 handshaker.mw_uri).group(1)

            try:
                # It's actually a wiki, right?
                assert home_wiki in WIKI_DICT
                user.editor.home_wiki = home_wiki
                user.editor.save()
            except AssertionError:
                # Site functionality mostly works if people don't
                # declare a homewiki. There are some broken bits, like the SUL
                # link, but users can set their homewiki, or admins can do it
                # in the admin interface.
                pass
        else:
            logger.info('User has been updated.'.format(user=user))

        request.session['user_created'] = created

        # The authenticate() function of a Django auth backend must return
        # the user.
        return user

    def get_user(self, user_id):
        try:
            return User.objects.get(pk=user_id)
        except User.DoesNotExist:
            logger.warning('There is no user {user_id}'.format(user_id=user_id))
            return None

class OAuthInitializeView(FormView):
    """
    Ask Wikipedia for a temporary key/secret for the user, and redirect
    them to their home Wikipedia to confirm authorization.
    """
    template_name = 'users/oauth_login.html'
    form_class = HomeWikiForm

    def form_valid(self, form):
        # The form returns the code for the home wiki (the first element of the
        # tuple in WIKIS), but we want the URL, so we grab it from WIKI_DICT.
        home_wiki = WIKI_DICT[form.cleaned_data['home_wiki']]

        # The site might be running under multiple URLs, so find out the current
        # one (and make sure it's legit).
        # The Sites framework was designed for different URLs that correspond to
        # different databases or functionality - it's not a good fit here.
        try:
            domain = self.request.get_host()
            assert domain in settings.ALLOWED_HOSTS # safety first!
        except (AssertionError, DisallowedHost):
            logger.exception()
            raise PermissionDenied

        logger.info('home_wiki was {home_wiki}, domain was {domain}'.format(home_wiki=home_wiki, domain=domain))

        # Get handshaker matching this wiki URL from our dict.
        handshaker = _get_handshaker(domain, home_wiki)
        logger.info('handshaker gotten')

        try:
            redirect, request_token = handshaker.initiate()
        except:
            logger.exception('Handshaker not initiated')
            raise

        logger.info('handshaker initiated')
        self.request.session['request_token'] = _dehydrate_token(request_token)
        self.request.session['home_wiki'] = home_wiki
        return HttpResponseRedirect(redirect)

class OAuthCallbackView(View):
    """
    Receive the redirect from Wikipedia and parse the response token.
    """

    def get(self, request, *args, **kwargs):
        response_qs = request.META['QUERY_STRING']

        # Get the handshaker. It should have already been constructed by
        # OAuthInitializeView.
        try:
            domain = self.request.get_host()
            assert domain in settings.ALLOWED_HOSTS
        except (AssertionError, DisallowedHost):
            logger.exception()
            raise PermissionDenied

        home_wiki = request.session.pop('home_wiki', None)

        try:
            handshaker = _get_handshaker(domain, home_wiki)
        except AssertionError:
            # get_handshaker will throw AssertionErrors for invalid data.
            logger.exception('Could not find handshaker')
            raise PermissionDenied

        # Get the request token, placed in session by OAuthInitializeView.
        session_token = request.session.pop('request_token', None)
        request_token = _rehydrate_token(session_token)

        if not request_token:
            logger.info('no request token :(')
            raise PermissionDenied

        # See if we can complete the OAuth process.
        try:
            access_token = handshaker.complete(request_token, response_qs)
        except:
            logger.exception('Access token generation failed :(')
            raise PermissionDenied

        user = authenticate(request=request,
                            access_token=access_token,
                            handshaker=handshaker)
        created = request.session.pop('user_created', False)

        if not user.is_active:
            # Do NOT log in the user.

            if created:
                messages.add_message(request, messages.WARNING,
                    _('Your Wikipedia account does not meet the eligibility '
                      'criteria in the terms of use, so your Wikipedia Library '
                      'Card Platform account cannot be activated.'))
            else:
                messages.add_message(request, messages.WARNING,
                    _('Either your Wikipedia Library Card Platform account has '
                      'been deactivated or your Wikipedia account no longer '
                      'meets the eligibility criteria in the terms of use, so '
                      'you cannot be logged in.'))

            return_url = reverse_lazy('terms')

        else:
            login(request, user)

            if created:
                # Translators: this message is displayed to users with brand new accounts.
                messages.add_message(request, messages.INFO, _('Welcome! '
                    'Please agree to the terms of use.'))
                return_url = reverse_lazy('terms')
            else:
                messages.add_message(request, messages.INFO, _('Welcome back!'))
                return_url = reverse_lazy('users:editor_detail',
                    kwargs={'pk': user.editor.pk})

        return HttpResponseRedirect(return_url)
<|endoftext|>


<commit_before>from django.db import models
from django.conf import settings
from django.utils.translation import ugettext_lazy as _
from django.utils.encoding import python_2_unicode_compatible
from django.contrib.contenttypes.models import ContentType
from django.contrib.contenttypes import generic

from .settings import MODEL_CATEGORY, MODEL_FLAG
from .utils import get_flag_model

USER_MODEL = getattr(settings, 'AUTH_USER_MODEL', 'auth.User')

# This allows South to handle our custom 'CharFieldNullable' field.
if 'south' in settings.INSTALLED_APPS:
    from south.modelsinspector import add_introspection_rules
    add_introspection_rules([], ['^sitecats\.models\.CharFieldNullable'])

class CharFieldNullable(models.CharField):
    """We use custom char field to put nulls in CategoryBase 'alias' field.
    That allows 'unique' param to be handled properly, so we don't have
    two categories with the same alias.

    """
    def get_prep_value(self, value):
        if value is not None:
            if value.strip() == '':
                return None
        return self.to_python(value)

@python_2_unicode_compatible
class CategoryBase(models.Model):

    title = models.CharField(_('Title'), max_length=250, help_text=_('Category name.'))
    note = models.TextField(_('Note'), blank=True)

    alias = CharFieldNullable(_('Alias'), max_length=80, help_text=_('Short name to address category from a template.'), blank=True, null=True, unique=True)
    slug = models.SlugField(verbose_name=_('Slug'), unique=True, max_length=250, null=True, blank=True)
    status = models.IntegerField(_('Status'), help_text=_('.'), null=True, blank=True, db_index=True)

    creator = models.ForeignKey(USER_MODEL, related_name='%(class)s_creators', verbose_name=_('Creator'))
    time_created = models.DateTimeField(_('Date created'), auto_now_add=True)
    time_modified = models.DateTimeField(_('Date modified'), editable=False, auto_now=True)
    # The last two are for 'adjacency list' model.
    parent = models.ForeignKey('self', related_name='%(class)s_parents', verbose_name=_('Parent'), help_text=_('Parent category.'), db_index=True, null=True, blank=True)
    sort_order = models.PositiveIntegerField(_('Sort order'), help_text=_('Item position among other categories under the same parent.'), db_index=True, default=0)

    def save(self, force_insert=False, force_update=False, **kwargs):
        """We override parent save method to set category sort order to its primary key value."""
        super(CategoryBase, self).save(force_insert, force_update, **kwargs)
        if self.sort_order == 0:
            self.sort_order = self.id
            self.save()

    class Meta:
        abstract = True
        verbose_name = _('Category')
        verbose_name_plural = _('Categories')

    def __str__(self):
        return self.title

@python_2_unicode_compatible
class FlagBase(models.Model):

    category = models.ForeignKey(MODEL_CATEGORY, related_name='%(class)s_categories', verbose_name=_('Category'), null=True, blank=True)
    note = models.TextField(_('Note'), blank=True)
    status = models.IntegerField(_('Status'), help_text=_('.'), null=True, blank=True, db_index=True)

    creator = models.ForeignKey(USER_MODEL, related_name='%(class)s_creators', verbose_name=_('Creator'))
    time_created = models.DateTimeField(_('Date created'), auto_now_add=True)

    # Here follows a link to an object.
    object_id = models.PositiveIntegerField(verbose_name=_('Object ID'), db_index=True)
    content_type = models.ForeignKey(ContentType, verbose_name=_('Content type'), related_name='%(app_label)s_%(class)s_tags')

    linked_object = generic.GenericForeignKey()

    class Meta:
        abstract = True
        verbose_name = _('Flag')
        verbose_name_plural = _('Flags')

    def __str__(self):
        return _('%s:%s flagged %s' % (self.content_type, self.object_id, self.category))

class Category(CategoryBase):
    """Built-in category class. Default functionality."""

class Flag(FlagBase):
    """Built-in flag class. Default functionality."""

class ModelWithBookmark(models.Model):
    """Helper base class for models with bookmarks."""

    bookmarks = generic.GenericRelation(MODEL_FLAG)

    def add_bookmark(self, user, category=None, note=''):
        bookmark = get_flag_model()(category=category, creator=user, note=note, linked_object=self)
        bookmark.save()

    def delete_bookmark(self, user, category=None):
        get_flag_model().objects.filter(category=category, creator=user, linked_object=self).delete()

    def is_bookmarked(self, user, category=None):
        if not user.id:
            return None
        return self.bookmarks.filter(category=category, creator=user).count()

    class Meta:
        abstract = True
<commit_msg>Basic ModelWithTag implemented.<commit_after>from django.db import models
from django.conf import settings
from django.utils import six
from django.utils.translation import ugettext_lazy as _
from django.utils.encoding import python_2_unicode_compatible
from django.contrib.contenttypes.models import ContentType
from django.contrib.contenttypes import generic

from .settings import MODEL_CATEGORY, MODEL_FLAG
from .utils import get_flag_model, get_category_model

USER_MODEL = getattr(settings, 'AUTH_USER_MODEL', 'auth.User')

# This allows South to handle our custom 'CharFieldNullable' field.
if 'south' in settings.INSTALLED_APPS:
    from south.modelsinspector import add_introspection_rules
    add_introspection_rules([], ['^sitecats\.models\.CharFieldNullable'])

class CharFieldNullable(models.CharField):
    """We use custom char field to put nulls in CategoryBase 'alias' field.
    That allows 'unique' param to be handled properly, so we don't have
    two categories with the same alias.

    """
    def get_prep_value(self, value):
        if value is not None:
            if value.strip() == '':
                return None
        return self.to_python(value)

@python_2_unicode_compatible
class CategoryBase(models.Model):

    title = models.CharField(_('Title'), max_length=250, help_text=_('Category name.'))
    note = models.TextField(_('Note'), blank=True)

    alias = CharFieldNullable(_('Alias'), max_length=80, help_text=_('Short name to address category from a template.'), blank=True, null=True, unique=True)
    slug = models.SlugField(verbose_name=_('Slug'), unique=True, max_length=250, null=True, blank=True)
    status = models.IntegerField(_('Status'), help_text=_('.'), null=True, blank=True, db_index=True)

    creator = models.ForeignKey(USER_MODEL, related_name='%(class)s_creators', verbose_name=_('Creator'))
    time_created = models.DateTimeField(_('Date created'), auto_now_add=True)
    time_modified = models.DateTimeField(_('Date modified'), editable=False, auto_now=True)
    # The last two are for 'adjacency list' model.
    parent = models.ForeignKey('self', related_name='%(class)s_parents', verbose_name=_('Parent'), help_text=_('Parent category.'), db_index=True, null=True, blank=True)
    sort_order = models.PositiveIntegerField(_('Sort order'), help_text=_('Item position among other categories under the same parent.'), db_index=True, default=0)

    def save(self, force_insert=False, force_update=False, **kwargs):
        """We override parent save method to set category sort order to its primary key value."""
        super(CategoryBase, self).save(force_insert, force_update, **kwargs)
        if self.sort_order == 0:
            self.sort_order = self.id
            self.save()

    class Meta:
        abstract = True
        verbose_name = _('Category')
        verbose_name_plural = _('Categories')

    def __str__(self):
        return self.title

@python_2_unicode_compatible
class FlagBase(models.Model):

    category = models.ForeignKey(MODEL_CATEGORY, related_name='%(class)s_categories', verbose_name=_('Category'), null=True, blank=True)
    note = models.TextField(_('Note'), blank=True)
    status = models.IntegerField(_('Status'), help_text=_('.'), null=True, blank=True, db_index=True)

    creator = models.ForeignKey(USER_MODEL, related_name='%(class)s_creators', verbose_name=_('Creator'))
    time_created = models.DateTimeField(_('Date created'), auto_now_add=True)

    # Here follows a link to an object.
    object_id = models.PositiveIntegerField(verbose_name=_('Object ID'), db_index=True)
    content_type = models.ForeignKey(ContentType, verbose_name=_('Content type'), related_name='%(app_label)s_%(class)s_tags')

    linked_object = generic.GenericForeignKey()

    class Meta:
        abstract = True
        verbose_name = _('Flag')
        verbose_name_plural = _('Flags')

    def __str__(self):
        return _('%s:%s flagged %s' % (self.content_type, self.object_id, self.category))

class Category(CategoryBase):
    """Built-in category class. Default functionality."""

class Flag(FlagBase):
    """Built-in flag class. Default functionality."""

class ModelWithTag(models.Model):
    """Helper base class for models with tags."""

    tags = generic.GenericRelation(MODEL_FLAG, limit_choices_to={'category__isnull': False})  # TODO verify limit works or create Manager

    def add_tag(self, tag, parent_category, user, note=''):
        category = get_category_model()(title=tag, creator=user, parent=parent_category)
        category.save()
        if isinstance(tag, six.string_types):
            tag = get_flag_model()(category=category, creator=user, note=note, linked_object=self)
        tag.save()
        return tag

    class Meta:
        abstract = True

class ModelWithBookmark(models.Model):
    """Helper base class for models with bookmarks."""

    bookmarks = generic.GenericRelation(MODEL_FLAG)

    def add_bookmark(self, user, category=None, note=''):
        bookmark = get_flag_model()(category=category, creator=user, note=note, linked_object=self)
        bookmark.save()
        return bookmark

    def delete_bookmark(self, user, category=None):
        get_flag_model().objects.filter(category=category, creator=user, linked_object=self).delete()

    def is_bookmarked(self, user, category=None):
        if not user.id:
            return None
        return self.bookmarks.filter(category=category, creator=user).count()

    class Meta:
        abstract = True
<|endoftext|>


<commit_before>#!/usr/bin/python

import http.client
import logging
import datetime
import re
import urllib
from . import register, Site, Section, Chapter

logger = logging.getLogger(__name__)

@register
class RoyalRoad(Site):
    domain = r'royalroad'

    """Royal Road: a place where people write novels, mostly seeming to be light-novel in tone."""
    @classmethod
    def matches(cls, url):
        # e.g. https://royalroad.com/fiction/6752/lament-of-the-fallen
        match = re.match(r'^(https?://(?:www\.)?%s\.com/fiction/\d+)/?.*' % cls.domain, url)
        if match:
            return match.group(1) + '/'

    def extract(self, url):
        workid = re.match(r'^https?://(?:www\.)?%s\.com/fiction/(\d+)/?.*' % self.domain, url).group(1)
        soup = self._soup('https://www.{}.com/fiction/{}'.format(self.domain, workid))
        # should have gotten redirected, for a valid title

        original_maxheaders = http.client._MAXHEADERS
        http.client._MAXHEADERS = 1000

        story = Section(
            title=soup.find('h1', property='name').string.strip(),
            author=soup.find('meta', property='books:author').get('content').strip(),
            url=soup.find('meta', property='og:url').get('content').strip()
        )

        for chapter in soup.select('#chapters tbody tr[data-url]'):
            chapter_url = str(urllib.parse.urljoin(story.url, str(chapter.get('data-url'))))

            contents, updated = self._chapter(chapter_url)

            story.add(Chapter(title=chapter.find('a', href=True).string.strip(), contents=contents, date=updated))

        http.client._MAXHEADERS = original_maxheaders

        return story

    def _chapter(self, url):
        logger.info("Extracting chapter @ %s", url)
        soup = self._soup(url)
        content = soup.find('div', class_='chapter-content')

        # TODO: this could be more robust, and I don't know if there's post-chapter notes anywhere as well.
        author_note = soup.find('div', class_='author-note-portlet')

        updated = datetime.datetime.fromtimestamp(
            int(soup.find(class_="profile-info").find('time').get('unixtime'))
        )

        return (author_note and (author_note.prettify() + '
') or '') + content.prettify(), updated

@register
class RoyalRoadL(RoyalRoad):
    domain = 'royalroadl'
<commit_msg>Place post-chapter RR author notes at the end of the chapter<commit_after>#!/usr/bin/python

import http.client
import logging
import datetime
import re
import urllib
from . import register, Site, Section, Chapter

logger = logging.getLogger(__name__)

@register
class RoyalRoad(Site):
    domain = r'royalroad'

    """Royal Road: a place where people write novels, mostly seeming to be light-novel in tone."""
    @classmethod
    def matches(cls, url):
        # e.g. https://royalroad.com/fiction/6752/lament-of-the-fallen
        match = re.match(r'^(https?://(?:www\.)?%s\.com/fiction/\d+)/?.*' % cls.domain, url)
        if match:
            return match.group(1) + '/'

    def extract(self, url):
        workid = re.match(r'^https?://(?:www\.)?%s\.com/fiction/(\d+)/?.*' % self.domain, url).group(1)
        soup = self._soup('https://www.{}.com/fiction/{}'.format(self.domain, workid))
        # should have gotten redirected, for a valid title

        original_maxheaders = http.client._MAXHEADERS
        http.client._MAXHEADERS = 1000

        story = Section(
            title=soup.find('h1', property='name').string.strip(),
            author=soup.find('meta', property='books:author').get('content').strip(),
            url=soup.find('meta', property='og:url').get('content').strip()
        )

        for chapter in soup.select('#chapters tbody tr[data-url]'):
            chapter_url = str(urllib.parse.urljoin(story.url, str(chapter.get('data-url'))))

            contents, updated = self._chapter(chapter_url)

            story.add(Chapter(title=chapter.find('a', href=True).string.strip(), contents=contents, date=updated))

        http.client._MAXHEADERS = original_maxheaders

        return story

    def _chapter(self, url):
        logger.info("Extracting chapter @ %s", url)
        soup = self._soup(url)
        content = soup.find('div', class_='chapter-content')

        # TODO: this could be more robust.
        author_note = soup.find('div', class_='author-note-portlet')

        # Find the portlet-body and check if the first child div is the author note.
        if 'author-note-portlet' in soup.find('div', class_='portlet-body').find('div', recursive=False)['class']:
            output = (author_note and (author_note.prettify() + '
') or '') + content.prettify()
        else:  # Possible post chapter note
            output = content.prettify() + (author_note and ('
' + author_note.prettify()) or '')

        updated = datetime.datetime.fromtimestamp(
            int(soup.find(class_="profile-info").find('time').get('unixtime'))
        )

        return output, updated

@register
class RoyalRoadL(RoyalRoad):
    domain = 'royalroadl'
<|endoftext|>