<commit_before>"""Script to load part files into a sql database.

For the inherent alpha model

We assume that the host temperature is well known so we will resitrict the results to the host temperature. (within error bars)
For HD30501 and HD211847 this means  +- 50K so a fixed temperature.

"""
import argparse
import os
import sys

import matplotlib.pyplot as plt
import pandas as pd
import sqlalchemy as sa

from utilities.param_file import parse_paramfile

def _parser():
    """Take care of all the argparse stuff.

    :returns: the args
    """
    parser = argparse.ArgumentParser(description='Chisquare analysis.')
    parser.add_argument('database', help='Database name.')
    parser.add_argument("-e", "--echo", help="Echo the SQL queries", action="store_true")
    # parser.add_argument('-s', '--suffix', help='Suffix to add to database name.')
    # parser.add_argument('-v', '--verbose', help='Turn on Verbose.', action="store_true")
    return parser.parse_args()

def decompose_database_name(database):
    """Database names of form */Star_obsnum_chip...db ."""
    os.path.split(database)
    path, name = os.path.split(database)
    return path, name.split("_")[:3]

def get_host_params(star):
    """Find host star parameters from param file."""
    param_file = "/home/jneal/Phd/data/parameter_files/{}_params.txt".format(star)
    params = parse_paramfile(param_file)
    return params["temp"], params["logg"], params["fe_h"]

def savefig(name):
    plt.savefig(os.path.join(path, "plots", name), bbox="tight")

def main(database, echo=False):
    path, star, obs_num, chip = decompose_database_name(database)

    teff, logg, fe_h = closest_model_params(*get_host_params(star), path)

    engine = sa.create_engine('sqlite:///{}'.format(database), echo=echo)
    table_names = engine.table_names()
    print("Table names in database =", engine.table_names())
    if len(table_names) == 1:
        tb_name = table_names[0]
    else:
        raise ValueError("Database has two many tables {}".format(table_names))

    fix_host_parameters(engine, teff, logg, fe_h)

    get_column_limits(engine, tb_name)
    smallest_chi2_values(engine, tb_name)

    df = pd.read_sql_query('SELECT alpha, chi2 FROM {0} LIMIT 10000'.format(tb_name), engine)

    fig, ax = plt.subplots()
    ax.scatter(df["alpha"], df["chi2"], s=3, alpha=0.5)

    ax.set_xlabel('alpha', fontsize=15)
    ax.set_ylabel('chi2', fontsize=15)
    ax.set_title('alpha and Chi2.')

    ax.grid(True)
    fig.tight_layout()

    plt.show()

    alpha_rv_plot(engine, tb_name)
    # alpha_rv_contour(engine, tb_name)

def alpha_rv_plot(engine, tb_name=None):
    # df = pd.read_sql(sa.text('SELECT alpha, rvs, chi2, teff_2 FROM table WHERE teff_1=:teff1, logg_1=:logg1, feh_1=:feh1'), engine, params={'teff1': 5200, 'logg1': 4.5, 'feh1': 0.0})
    #df = pd.read_sql(sa.text('SELECT alpha, rvs, chi2, teff_2 FROM table WHERE teff_1=:teff_1'), engine, params={'teff_1': 5200})
    #df = pd.read_sql(sa.text('SELECT alpha, rvs, chi2, teff_2 FROM table WHERE teff_1=5200 and logg_1=4.5 and feh_1=0.0'), engine)
    df = pd.read_sql(sa.text('SELECT alpha, rv, chi2, teff_2 FROM {0}'.format(tb_name)), engine)
    #df = pd.read_sql_query('SELECT alpha  FROM table', engine)

    fig, ax = plt.subplots()
    ax.scatter(df["rv"], df["chi2"], c=df["alpha"], s=df["teff_2"] / 50, alpha=0.5)

    ax.set_xlabel('rv offset', fontsize=15)
    ax.set_ylabel('chi2', fontsize=15)
    ax.set_title('alpha (color) and companion temperature (size=Temp/50).')

    ax.grid(True)
    fig.tight_layout()

    plt.show()

def alpha_rv_contour(engine, tb_name=None):
    # df = pd.read_sql(sa.text('SELECT alpha, rvs, chi2, teff_2 FROM table WHERE teff_1=:teff1, logg_1=:logg1, feh_1=:feh1'), engine, params={'teff1': 5200, 'logg1': 4.5, 'feh1': 0.0})
    #df = pd.read_sql(sa.text('SELECT alpha, rvs, chi2, teff_2 FROM table WHERE teff_1=:teff_1'), engine, params={'teff_1': 5200})
    #df = pd.read_sql(sa.text('SELECT alpha, rvs, chi2, teff_2 FROM table WHERE teff_1=5200 and logg_1=4.5 and feh_1=0.0'), engine)
    df = pd.read_sql(sa.text('SELECT alpha, rv, chi2, teff_2 FROM {0}'.format(tb_name)), engine)
    #df = pd.read_sql_query('SELECT alpha  FROM table', engine)

    fig, ax = plt.subplots()
    ax.contourf(df["rv"], df["chi2"], c=df["alpha"], alpha=0.5)

    ax.set_xlabel('rv offset', fontsize=15)
    ax.set_ylabel('chi2', fontsize=15)
    ax.set_title('alpha (color)')

    ax.grid(True)
    fig.tight_layout()

    plt.show()

def smallest_chi2_values(engine, tb_name):
    """Find smallest chi2 in table."""
    df = pd.read_sql(sa.text('SELECT * FROM {0} ORDER BY chi2 ASC LIMIT 10'.format(tb_name)), engine)
    # df = pd.read_sql_query('SELECT alpha  FROM table', engine)

    print("Samllest Chi2 values in the database.")
    print(df.head(n=15))

    plt.show()

def get_column_limits(engine, tb_name):
    print("Column Value Ranges")
    for col in ["teff_1", "teff_2", "logg_1", "logg_2", "alpha", "gamma", "rv", "chi2"]:
        query = """
               SELECT * FROM (SELECT {1} FROM {0} ORDER BY {1} ASC LIMIT 1)
               UNION ALL
               SELECT * FROM (SELECT {1} FROM {0} ORDER BY {1} DESC LIMIT 1)
               """.format(tb_name, col)
        df = pd.read_sql(sa.text(query), engine)
        print(col, min(df[col]), max(df[col]))

def fix_host_parameters(engine, teff, logg, fe_h):
    print("Fixed host analysis")
    for col in ["teff_2", "logg_2", "alpha", "gamma", "rv"]
    query = """
           SELECT {} FROM {} WHERE teff_1 == {} logg_1 == {} feh_1 == {})""" .format(col, "chi2", teff, logg, fe_h)
    df = pd.read_sql(sa.text(query), engine)
    df.plot.scatter()
    savefig("fixed_host_params_")
    plt.show()

if __name__ == '__main__':
    args = vars(_parser())
    opts = {k: args[k] for k in args}

    sys.exit(main(**opts))
<commit_msg>Save plots to file, collect the parameters to pass around.<commit_after>"""Script to load part files into a sql database.

For the inherent alpha model

We assume that the host temperature is well known so we will resitrict the
results to the host temperature. (within error bars)
For HD30501 and HD211847 this means  +- 50K so a fixed temperature.

"""
import argparse
import os
import sys

import matplotlib.pyplot as plt
import pandas as pd
import sqlalchemy as sa

from utilities.param_file import parse_paramfile
from utilities.phoenix_utils import closest_model_params
from utilities.debug_utils import timeit2

def _parser():
    """Take care of all the argparse stuff.

    :returns: the args
    """
    parser = argparse.ArgumentParser(description='Chisquare analysis.')
    parser.add_argument('database', help='Database name.')
    parser.add_argument("-e", "--echo", help="Echo the SQL queries", action="store_true")
    # parser.add_argument('-s', '--suffix', help='Suffix to add to database name.')
    # parser.add_argument('-v', '--verbose', help='Turn on Verbose.', action="store_true")
    return parser.parse_args()

def decompose_database_name(database):
    """Database names of form */Star_obsnum_chip...db ."""
    os.path.split(database)
    path, name = os.path.split(database)
    name_split = name.split("_")
    star, obsnum = name_split[0].split("-")
    chip = name_split[1]
    return path, star, obsnum, chip

def get_host_params(star):
    """Find host star parameters from param file."""
    param_file = "/home/jneal/Phd/data/parameter_files/{}_params.dat".format(star)
    params = parse_paramfile(param_file)
    return params["temp"], params["logg"], params["fe_h"]

def main(database, echo=False):
    path, star, obs_num, chip = decompose_database_name(database)

    teff, logg, fe_h = closest_model_params(*get_host_params(star))
    params = {"path": path, "star": star, "obs_num": obs_num, "chip": chip, "teff": teff, "logg": logg, "fe_h": fe_h}

    engine = sa.create_engine('sqlite:///{}'.format(database), echo=echo)
    table_names = engine.table_names()
    print("Table names in database =", engine.table_names())
    if len(table_names) == 1:
        tb_name = table_names[0]
    else:
        raise ValueError("Database has two many tables {}".format(table_names))

    fix_host_parameters(engine, tb_name, params)

    get_column_limits(engine, tb_name, params)
    smallest_chi2_values(engine, tb_name, params)

    df = pd.read_sql_query('SELECT alpha, chi2 FROM {0} LIMIT 10000'.format(tb_name), engine)

    fig, ax = plt.subplots()
    ax.scatter(df["alpha"], df["chi2"], s=3, alpha=0.5)

    ax.set_xlabel('alpha', fontsize=15)
    ax.set_ylabel('chi2', fontsize=15)
    ax.set_title('alpha and Chi2.')

    ax.grid(True)
    fig.tight_layout()
    name = "{0}-{1}_{2}_test_test_figure1.pdf".format(
        params["star"], params["obs_num"], params["chip"])
    plt.savefig(os.path.join(params["path"], "plots", name))
    #$plt.show()

    alpha_rv_plot(engine, tb_name)
    # alpha_rv_contour(engine, tb_name)

def alpha_rv_plot(engine, tb_name=None):
    # df = pd.read_sql(sa.text('SELECT alpha, rvs, chi2, teff_2 FROM table WHERE teff_1=:teff1, logg_1=:logg1, feh_1=:feh1'), engine, params={'teff1': 5200, 'logg1': 4.5, 'feh1': 0.0})
    #df = pd.read_sql(sa.text('SELECT alpha, rvs, chi2, teff_2 FROM table WHERE teff_1=:teff_1'), engine, params={'teff_1': 5200})
    #df = pd.read_sql(sa.text('SELECT alpha, rvs, chi2, teff_2 FROM table WHERE teff_1=5200 and logg_1=4.5 and feh_1=0.0'), engine)
    df = pd.read_sql(sa.text('SELECT alpha, rv, chi2, teff_2 FROM {0}'.format(tb_name)), engine)
    #df = pd.read_sql_query('SELECT alpha  FROM table', engine)

    fig, ax = plt.subplots()
    ax.scatter(df["rv"], df["chi2"], c=df["alpha"], s=df["teff_2"] / 50, alpha=0.5)

    ax.set_xlabel('rv offset', fontsize=15)
    ax.set_ylabel('chi2', fontsize=15)
    ax.set_title('alpha (color) and companion temperature (size=Temp/50).')

    ax.grid(True)
    fig.tight_layout()
    name = "{0}-{1}_{2}_test_alpha_rv.pdf".format(
        params["star"], params["obs_num"], params["chip"])
    plt.savefig(os.path.join(params["path"], "plots", name))

    # plt.show()

def alpha_rv_contour(engine, tb_name=None):
    # df = pd.read_sql(sa.text('SELECT alpha, rvs, chi2, teff_2 FROM table WHERE teff_1=:teff1, logg_1=:logg1, feh_1=:feh1'), engine, params={'teff1': 5200, 'logg1': 4.5, 'feh1': 0.0})
    #df = pd.read_sql(sa.text('SELECT alpha, rvs, chi2, teff_2 FROM table WHERE teff_1=:teff_1'), engine, params={'teff_1': 5200})
    #df = pd.read_sql(sa.text('SELECT alpha, rvs, chi2, teff_2 FROM table WHERE teff_1=5200 and logg_1=4.5 and feh_1=0.0'), engine)
    df = pd.read_sql(sa.text('SELECT alpha, rv, chi2, teff_2 FROM {0}'.format(tb_name)), engine)
    #df = pd.read_sql_query('SELECT alpha  FROM table', engine)

    fig, ax = plt.subplots()
    ax.contourf(df["rv"], df["chi2"], c=df["alpha"], alpha=0.5)

    ax.set_xlabel('rv offset', fontsize=15)
    ax.set_ylabel('chi2', fontsize=15)
    ax.set_title('alpha (color)')

    ax.grid(True)
    fig.tight_layout()
    name = "{0}-{1}_{2}_test_alpha_rv_contour.pdf".format(
        params["star"], params["obs_num"], params["chip"])
    plt.savefig(os.path.join(params["path"], "plots", name))
    # plt.show()

def smallest_chi2_values(engine, tb_name, params):
    """Find smallest chi2 in table."""
    df = pd.read_sql(sa.text('SELECT * FROM {0} ORDER BY chi2 ASC LIMIT 10'.format(tb_name)), engine)
    # df = pd.read_sql_query('SELECT alpha  FROM table', engine)

    print("Samllest Chi2 values in the database.")
    print(df.head(n=15))
    name = "{0}-{1}_{2}_test_smallest_chi2.pdf".format(
        params["star"], params["obs_num"], params["chip"])
    plt.savefig(os.path.join(params["path"], "plots", name))
    # plt.show()

def get_column_limits(engine, tb_name, params):
    print("Column Value Ranges")
    for col in ["teff_1", "teff_2", "logg_1", "logg_2", "alpha", "gamma", "rv", "chi2"]:
        query = """
               SELECT * FROM (SELECT {1} FROM {0} ORDER BY {1} ASC LIMIT 1)
               UNION ALL
               SELECT * FROM (SELECT {1} FROM {0} ORDER BY {1} DESC LIMIT 1)
               """.format(tb_name, col)
        df = pd.read_sql(sa.text(query), engine)
        print(col, min(df[col]), max(df[col]))

@timeit2
def fix_host_parameters(engine, tb_name, params):
    print("Fixed host analysis.")
    for col in ["teff_2", "logg_2", "feh_2", "alpha", "gamma", "rv"]:
        query = """SELECT {}, {} FROM {} WHERE (teff_1 = {}  AND logg_1 = {} AND feh_1 = {})""" .format(
            col, "chi2", tb_name, params["teff"], params["logg"], params["fe_h"])
        df = pd.read_sql(sa.text(query), engine)
        print(df.columns)
        df.plot.scatter(col, "chi2")
        name = "{0}-{1}_{2}_fixed_host_params_{3}.pdf".format(
            params["star"], params["obs_num"], params["chip"], col)
        plt.savefig(os.path.join(params["path"], "plots", name))
        plt.close()

if __name__ == '__main__':
    args = vars(_parser())
    opts = {k: args[k] for k in args}

    sys.exit(main(**opts))
<|endoftext|>


<commit_before>"""Node Views"""
from __future__ import unicode_literals

import logging

import rest_framework.status as status
from django.http.response import Http404
from rest_framework.generics import GenericAPIView, ListAPIView
from rest_framework.response import Response

import mesos_api.api as mesos_api
import util.rest as rest_util
from node.models import Node
from node.serializers import NodeSerializer
from node.serializers_extra import NodeDetailsSerializer, NodeStatusSerializer
from scheduler.models import Scheduler

logger = logging.getLogger(__name__)

class NodesView(ListAPIView):
    """This view is the endpoint for viewing a list of nodes with metadata"""
    queryset = Node.objects.all()
    serializer_class = NodeSerializer

    def list(self, request):
        """Retrieves the list of all nodes and returns it in JSON form

        :param request: the HTTP GET request
        :type request: :class:`rest_framework.request.Request`
        :rtype: :class:`rest_framework.response.Response`
        :returns: the HTTP response to send back to the user
        """
        started = rest_util.parse_timestamp(request, 'started', required=False)
        ended = rest_util.parse_timestamp(request, 'ended', required=False)
        rest_util.check_time_range(started, ended)
        include_inactive = rest_util.parse_bool(request, 'include_inactive', False, False)

        order = rest_util.parse_string_list(request, 'order', required=False)

        nodes = Node.objects.get_nodes(started, ended, order, include_inactive=include_inactive)

        page = self.paginate_queryset(nodes)
        serializer = self.get_serializer(page, many=True)
        return self.get_paginated_response(serializer.data)

class NodeDetailsView(GenericAPIView):
    """This view is the endpoint for viewing and modifying a node"""
    queryset = Node.objects.all()
    serializer_class = NodeDetailsSerializer
    update_fields = ('pause_reason', 'is_paused', 'is_active')

    def get(self, request, node_id):
        """Gets node info

        :param request: the HTTP GET request
        :type request: :class:`rest_framework.request.Request`
        :param node_id: The ID for the node.
        :type node_id: str
        :rtype: :class:`rest_framework.response.Response`
        :returns: the HTTP response to send back to the user
        """

        return Response(self._get_node_details(node_id))

    def patch(self, request, node_id):
        """Modify node info with a subset of fields

        :param request: the HTTP GET request
        :type request: :class:`rest_framework.request.Request`
        :param node_id: The ID for the node.
        :type node_id: str
        :rtype: :class:`rest_framework.response.Response`
        :returns: the HTTP response to send back to the user
        """

        extra = filter(lambda x, y=self.update_fields: x not in y, request.data.keys())
        if len(extra) > 0:
            return Response('Unexpected fields: %s' % ', '.join(extra), status=status.HTTP_400_BAD_REQUEST)

        if not len(request.data):
            return Response('No fields specified for update.', status=status.HTTP_400_BAD_REQUEST)

        try:
            Node.objects.get(id=node_id)
        except Node.DoesNotExist:
            raise Http404

        Node.objects.update_node(dict(request.data), node_id=node_id)

        result = self._get_node_details(node_id)

        return Response(result, status=status.HTTP_200_OK,
                        headers={'Location': request.build_absolute_uri()})

    def _get_node_details(self, node_id):

        # Fetch the basic node attributes
        try:
            node = Node.objects.get_details(node_id)
        except Node.DoesNotExist:
            raise Http404
        serializer = self.get_serializer(node)
        result = serializer.data

        # Attempt to fetch resource usage for the node
        resources = None
        try:
            sched = Scheduler.objects.get_master()
            slave_info = mesos_api.get_slave(sched.master_hostname, sched.master_port, node.slave_id, True)
            if slave_info and slave_info.total:
                resources = slave_info.to_dict()['resources']
        except:
            logger.exception('Unable to fetch slave resource usage')
        if resources:
            result['resources'] = resources
        else:
            result['disconnected'] = True

        return result

class NodesStatusView(ListAPIView):
    """This view is the endpoint for retrieving overall node status information."""
    queryset = []
    serializer_class = NodeStatusSerializer

    def list(self, request):
        """Retrieves the list of all nodes with execution status and returns it in JSON form

        :param request: the HTTP GET request
        :type request: :class:`rest_framework.request.Request`
        :rtype: :class:`rest_framework.response.Response`
        :returns: the HTTP response to send back to the user
        """

        # Get a list of all node status counts
        started = rest_util.parse_timestamp(request, 'started', 'PT3H0M0S')
        ended = rest_util.parse_timestamp(request, 'ended', required=False)
        node_statuses = Node.objects.get_status(started, ended)

        # Get the online nodes
        try:
            sched = Scheduler.objects.get_master()
            slaves = mesos_api.get_slaves(sched.master_hostname, sched.master_port)
            slaves_dict = {s.hostname for s in slaves}
        except:
            logger.exception('Unable to fetch nodes online status')
            slaves_dict = dict()

        # Add the online status to each node
        for node_status in node_statuses:
            node_status.is_online = node_status.node.hostname in slaves_dict

        page = self.paginate_queryset(node_statuses)
        serializer = self.get_serializer(page, many=True)
        return self.get_paginated_response(serializer.data)
<commit_msg>Removed header  from PATCH Response return<commit_after>"""Node Views"""
from __future__ import unicode_literals

import logging

import rest_framework.status as status
from django.http.response import Http404
from rest_framework.generics import GenericAPIView, ListAPIView
from rest_framework.response import Response

import mesos_api.api as mesos_api
import util.rest as rest_util
from node.models import Node
from node.serializers import NodeSerializer
from node.serializers_extra import NodeDetailsSerializer, NodeStatusSerializer
from scheduler.models import Scheduler

logger = logging.getLogger(__name__)

class NodesView(ListAPIView):
    """This view is the endpoint for viewing a list of nodes with metadata"""
    queryset = Node.objects.all()
    serializer_class = NodeSerializer

    def list(self, request):
        """Retrieves the list of all nodes and returns it in JSON form

        :param request: the HTTP GET request
        :type request: :class:`rest_framework.request.Request`
        :rtype: :class:`rest_framework.response.Response`
        :returns: the HTTP response to send back to the user
        """
        started = rest_util.parse_timestamp(request, 'started', required=False)
        ended = rest_util.parse_timestamp(request, 'ended', required=False)
        rest_util.check_time_range(started, ended)
        include_inactive = rest_util.parse_bool(request, 'include_inactive', False, False)

        order = rest_util.parse_string_list(request, 'order', required=False)

        nodes = Node.objects.get_nodes(started, ended, order, include_inactive=include_inactive)

        page = self.paginate_queryset(nodes)
        serializer = self.get_serializer(page, many=True)
        return self.get_paginated_response(serializer.data)

class NodeDetailsView(GenericAPIView):
    """This view is the endpoint for viewing and modifying a node"""
    queryset = Node.objects.all()
    serializer_class = NodeDetailsSerializer
    update_fields = ('pause_reason', 'is_paused', 'is_active')

    def get(self, request, node_id):
        """Gets node info

        :param request: the HTTP GET request
        :type request: :class:`rest_framework.request.Request`
        :param node_id: The ID for the node.
        :type node_id: str
        :rtype: :class:`rest_framework.response.Response`
        :returns: the HTTP response to send back to the user
        """

        return Response(self._get_node_details(node_id))

    def patch(self, request, node_id):
        """Modify node info with a subset of fields

        :param request: the HTTP GET request
        :type request: :class:`rest_framework.request.Request`
        :param node_id: The ID for the node.
        :type node_id: str
        :rtype: :class:`rest_framework.response.Response`
        :returns: the HTTP response to send back to the user
        """

        extra = filter(lambda x, y=self.update_fields: x not in y, request.data.keys())
        if len(extra) > 0:
            return Response('Unexpected fields: %s' % ', '.join(extra), status=status.HTTP_400_BAD_REQUEST)

        if not len(request.data):
            return Response('No fields specified for update.', status=status.HTTP_400_BAD_REQUEST)

        try:
            Node.objects.get(id=node_id)
        except Node.DoesNotExist:
            raise Http404

        Node.objects.update_node(dict(request.data), node_id=node_id)

        result = self._get_node_details(node_id)

        return Response(result, status=status.HTTP_200_OK)

    def _get_node_details(self, node_id):

        # Fetch the basic node attributes
        try:
            node = Node.objects.get_details(node_id)
        except Node.DoesNotExist:
            raise Http404
        serializer = self.get_serializer(node)
        result = serializer.data

        # Attempt to fetch resource usage for the node
        resources = None
        try:
            sched = Scheduler.objects.get_master()
            slave_info = mesos_api.get_slave(sched.master_hostname, sched.master_port, node.slave_id, True)
            if slave_info and slave_info.total:
                resources = slave_info.to_dict()['resources']
        except:
            logger.exception('Unable to fetch slave resource usage')
        if resources:
            result['resources'] = resources
        else:
            result['disconnected'] = True

        return result

class NodesStatusView(ListAPIView):
    """This view is the endpoint for retrieving overall node status information."""
    queryset = []
    serializer_class = NodeStatusSerializer

    def list(self, request):
        """Retrieves the list of all nodes with execution status and returns it in JSON form

        :param request: the HTTP GET request
        :type request: :class:`rest_framework.request.Request`
        :rtype: :class:`rest_framework.response.Response`
        :returns: the HTTP response to send back to the user
        """

        # Get a list of all node status counts
        started = rest_util.parse_timestamp(request, 'started', 'PT3H0M0S')
        ended = rest_util.parse_timestamp(request, 'ended', required=False)
        node_statuses = Node.objects.get_status(started, ended)

        # Get the online nodes
        try:
            sched = Scheduler.objects.get_master()
            slaves = mesos_api.get_slaves(sched.master_hostname, sched.master_port)
            slaves_dict = {s.hostname for s in slaves}
        except:
            logger.exception('Unable to fetch nodes online status')
            slaves_dict = dict()

        # Add the online status to each node
        for node_status in node_statuses:
            node_status.is_online = node_status.node.hostname in slaves_dict

        page = self.paginate_queryset(node_statuses)
        serializer = self.get_serializer(page, many=True)
        return self.get_paginated_response(serializer.data)
<|endoftext|>


<commit_before>""" """

import re
from collections import Counter
from datetime import datetime
from datetime import timedelta

date_formatting = '%d/%b/%Y:%H:%M:%S %z'
request_pattern = '(?P<ip>\S+).*\s\[(?P<date>.+)\]\s*\"(?P<request>.+)\"\s*(?P<response_code>\d+)\s(?P<response_length>\d+|-)'
url_pattern = '((GET|HEAD|POST)\s+)?(?P<request_url>.+?)\s+(HTTP.*)?$'

def parse_log_line(log_line):
    """
    Parse log string into ip, date, url,
    protocol, response_code, and response_length.
    Return parsed information as a dictionary.
    Attributes:
        log_line (str): one request from log file
    Return:
        dictionary with following keys:
            - ip
            - date
            - url
            - response_code
            - response_length
    """

    request_regex = re.compile(request_pattern)
    url_regex = re.compile(url_pattern)

    parsed_request = request_regex.match(log_line)

    if parsed_request:
        result = parsed_request.groupdict()
        result["url"] = result["request"]
        result["date"] = datetime.strptime(result["date"], date_formatting)
        parsed_url = url_regex.match(result["request"])
        if parsed_url:
            result["url"] = parsed_url.groupdict()["request_url"].strip()
        return result
    else:
        write_information_to_file("../log_output/corrupted_data.txt", log_line)

def read_from_file(file):
    """
    Generator that yields line from log file.
    Attributes:
        file (str): url to log file
    """
    try:
        with open(file) as f_obj:
            for line in f_obj:
                parsed_line = parse_log_line(line)
                if parsed_line:
                    yield parse_log_line(line)
    except FileNotFoundError:
        print("read_from_file: File not found")

def count_host_frequency(counter, data):
    """
    Counts how many times hosts/IP addresses have accessed the site.
    Attributes:
        counter (Counter class): keeps track of all hosts
        data (dict): parsed information about request
    """
    if data:
        counter[data["ip"]] += 1

def count_bandwidth_resources(counter, data):
    """
    Counts how many bytes were sent for particular resources.
    Attributes:
        counter (Counter class): keeps overall bytes for particular resource
        data (dict): parsed information about request
    """
    if data:
        if data["response_length"].strip() != "-":
            counter[data["url"]] += int(data["response_length"])

def write_information_to_file(file, log_line):
    """
    Append information to a log file.
    Attributes:
        file (str): log file url
        log_line (str): information to append
    """
    with open(file, 'a') as f_obj:
        f_obj.write(log_line + "\n")

# filename = "../insight_testsuite/tests/test_features/log_input/log.txt"
filename = "../log_input/log.txt"
ip_frequency = Counter()
resources = Counter()

# Calculate feature 1 and feature 2
print("Feature 1 and 2 start")

for entry in read_from_file(filename):
    count_host_frequency(ip_frequency, entry)
    count_bandwidth_resources(resources, entry)

top_hosts = ip_frequency.most_common(10)
top_resources = resources.most_common(10)

# Write results to files
for host in top_hosts:
    information = host[0] + "," + str(host[1])
    write_information_to_file("../log_output/hosts.txt", information)

for resource in top_resources:
    write_information_to_file("../log_output/resources.txt", resource[0])

print("Feature 1 and 2 stop")
# FEATURE 3
# Need refactor
# Create array with size 3600 that will represent seconds in 60 minutes
# and save requests for each second

print("Feature 3 starts")

requests_per_hour = [0] * 3600
popular_dates = {}

generator = read_from_file(filename)
request = next(generator)
start = request["date"]
stop = start + timedelta(0, 3600)

def save_hour(key_date, top_10, requests_in_current_hour):
    """
    Save number of requests in current hour
    if it is on of the 10 busiest hours

    Attributes:
        key_date (str): time when current hour started
        top_10 (dict): 10 most busiest hours
        requests_in_current_hour (list): list of requests in this hour
    """
    if len(top_10) == 10:
        if sum(requests_in_current_hour) > min(top_10.values()):
            del top_10[min(top_10, key=top_10.get)]
            top_10[key_date] = sum(requests_in_current_hour)
    else:
        top_10[key_date] = sum(requests_in_current_hour)

def update_statistics(start_date, stop_date, top_10, requests_in_current_hour):
    """
    Update statistics about top 10 busiest hours and shift time by 1 second
    Attributes:

    """
    # Save only the top 10 popular hours.

while request:
    if request["date"] < stop:
        requests_per_hour[(request["date"] - start).seconds] += 1

    while request["date"] > stop:
        # Save only the top 10 popular hours.
        date_stamp = start.strftime(date_formatting)
        save_hour(date_stamp, popular_dates, requests_per_hour)
        # Save statistics and shift time by 1 second
        requests_per_hour.pop(0)
        requests_per_hour.append(0)
        start = start + timedelta(0, 1)
        stop = stop + timedelta(0, 1)
    try:
        request = next(generator)
    except StopIteration:
        request = None
        for i in range(0, 3600):
            # Count requests one more hour after the last request
            date_stamp = start.strftime(date_formatting)
            save_hour(date_stamp, popular_dates, requests_per_hour)
            # Shift by 1 second
            requests_per_hour.pop(0)
            requests_per_hour.append(0)
            start = start + timedelta(0, 1)

print(len(popular_dates.values()))
print(popular_dates)
<commit_msg>fixed logic error in feature 3: didn't write last hour if request date is greater than stop counter<commit_after>""" """

import re
from collections import Counter
from datetime import datetime
from datetime import timedelta

date_formatting = '%d/%b/%Y:%H:%M:%S %z'
request_pattern = '(?P<ip>\S+).*\s\[(?P<date>.+)\]\s*\"(?P<request>.+)\"\s*(?P<response_code>\d+)\s(?P<response_length>\d+|-)'
url_pattern = '((GET|HEAD|POST)\s+)?(?P<request_url>.+?)\s+(HTTP.*)?$'

def parse_log_line(log_line):
    """
    Parse log string into ip, date, url,
    protocol, response_code, and response_length.
    Return parsed information as a dictionary.
    Attributes:
        log_line (str): one request from log file
    Return:
        dictionary with following keys:
            - ip
            - date
            - url
            - response_code
            - response_length
    """

    request_regex = re.compile(request_pattern)
    url_regex = re.compile(url_pattern)

    parsed_request = request_regex.match(log_line)

    if parsed_request:
        result = parsed_request.groupdict()
        result["url"] = result["request"]
        result["date"] = datetime.strptime(result["date"], date_formatting)
        parsed_url = url_regex.match(result["request"])
        if parsed_url:
            result["url"] = parsed_url.groupdict()["request_url"].strip()
        return result
    else:
        write_information_to_file("../log_output/corrupted_data.txt", log_line)

def read_from_file(file):
    """
    Generator that yields line from log file.
    Attributes:
        file (str): url to log file
    """
    try:
        with open(file) as f_obj:
            for line in f_obj:
                parsed_line = parse_log_line(line)
                if parsed_line:
                    yield parse_log_line(line)
    except FileNotFoundError:
        print("read_from_file: File not found")

def count_host_frequency(counter, data):
    """
    Counts how many times hosts/IP addresses have accessed the site.
    Attributes:
        counter (Counter class): keeps track of all hosts
        data (dict): parsed information about request
    """
    if data:
        counter[data["ip"]] += 1

def count_bandwidth_resources(counter, data):
    """
    Counts how many bytes were sent for particular resources.
    Attributes:
        counter (Counter class): keeps overall bytes for particular resource
        data (dict): parsed information about request
    """
    if data:
        if data["response_length"].strip() != "-":
            counter[data["url"]] += int(data["response_length"])

def write_information_to_file(file, log_line):
    """
    Append information to a log file.
    Attributes:
        file (str): log file url
        log_line (str): information to append
    """
    with open(file, 'a') as f_obj:
        f_obj.write(log_line + "\n")

filename = "../insight_testsuite/tests/test_features/log_input/log.txt"
# filename = "../log_input/log.txt"
ip_frequency = Counter()
resources = Counter()

# Calculate feature 1 and feature 2
print("Feature 1 and 2 start")

for entry in read_from_file(filename):
    count_host_frequency(ip_frequency, entry)
    count_bandwidth_resources(resources, entry)

top_hosts = ip_frequency.most_common(10)

# Write results to a file
for host in top_hosts:
    information = host[0] + "," + str(host[1])
    write_information_to_file("../log_output/hosts.txt", information)

top_resources = resources.most_common(10)

# Write results to a file
for resource in top_resources:
    write_information_to_file("../log_output/resources.txt", resource[0])

print("Feature 1 and 2 stop")

# FEATURE 3
# Need refactor
# Create array with size 3600 that will represent seconds in 60 minutes
# and save requests for each second

print("Feature 3 starts")

requests_per_hour = [0] * 3600
popular_dates = {}

generator = read_from_file(filename)
request = next(generator)
start = request["date"]
stop = start + timedelta(0, 3600)

def save_hour(key_date, top_10, requests_in_current_hour):
    """
    Save number of requests in current hour
    if it is on of the 10 busiest hours

    Attributes:
        key_date (str): time when current hour started
        top_10 (dict): 10 most busiest hours
        requests_in_current_hour (list): list of requests in this hour
    """

    if len(top_10) == 10:
        if sum(requests_in_current_hour) > min(top_10.values()):
            del top_10[min(top_10, key=top_10.get)]
            top_10[key_date] = sum(requests_in_current_hour)
    else:
        top_10[key_date] = sum(requests_in_current_hour)

def update_statistics(start_date, stop_date, top_10, requests_in_current_hour):
    """
    Update statistics about top 10 busiest hours and shift time by 1 second
    Attributes:

    """
    # Save only the top 10 popular hours.

while request:
    while request["date"] >= stop:
        # Save only the top 10 popular hours.
        date_stamp = start.strftime(date_formatting)
        save_hour(date_stamp, popular_dates, requests_per_hour)
        # Save statistics and shift time by 1 second
        requests_per_hour.pop(0)
        requests_per_hour.append(0)
        start = start + timedelta(0, 1)
        stop = stop + timedelta(0, 1)

    requests_per_hour[(request["date"] - start).seconds] += 1

    try:
        request = next(generator)
    except StopIteration:
        request = None
        for i in range(0, 3600):
            # Count requests one more hour after the last request
            date_stamp = start.strftime(date_formatting)
            save_hour(date_stamp, popular_dates, requests_per_hour)
            # Shift by 1 second
            requests_per_hour.pop(0)
            requests_per_hour.append(0)
            start = start + timedelta(0, 1)

print(len(popular_dates.values()))
print(popular_dates)
<|endoftext|>


<commit_before># -*- coding: utf-8 -*-
""" This module contains ShortenedURL class and related classes """
from bisect import bisect_left
from math import log, floor
from random import randint

from cached_property import cached_property
from flask import url_for, abort
from sqlalchemy import types, inspect
from sqlalchemy.exc import IntegrityError

from . import db, app

class AliasValueError(ValueError):
    """The value of alias is incorrect """

class AliasLengthValueError(ValueError):
    """The value of alias-length related parameter is incorrect """

class RegistrationRetryLimitExceeded(Exception):
    """The maximum number of attempts at retrying to register
    a new short URL has been exceeded """

class AliasType(type):
    def __new__(cls, *args, **kwargs):
        alias_class = type.__new__(cls, *args, **kwargs)
        alias_class.init_random_factory(
            app.config['MIN_NEW_ALIAS_LENGTH'],
            app.config['MAX_NEW_ALIAS_LENGTH']
        )
        return alias_class

class Alias(metaclass=AliasType):
    """ An identifier for shortened URL

    In has two values used as its representations: a string value
    and an integer value, used in short URLs and in database,
    respectively.

    :cvar _chars: a string consisting of characters to be used
    in a string value of an alias.
    :cvar _base: a base of numeral system used to interpret string
    values as integers.
    :cvar _max_int_32: a maximum value of a 32 bit signed integer.

    This value is assumed as maximum allowed for the integers used
    in generation because we assume SQLAlchemy.types.Integer -
    a base type for alias property of ShortenedURL class - will
    translate into 32 bit signed integer type of underlying
    database engine used by the application.

    :cvar _max_allowed_length: maximum number of characters for
    newly generated alias strings for which there are at least
    some strings whose corresponding integer values are smaller
    or equal to _max_int_32.
    """

    _chars = '0123456789abcdefghijkmnopqrstuvwxyz'
    _base = len(_chars)
    _max_int_32 = 2**31 - 1
    _max_allowed_length = floor(log(_max_int_32, _base)) + 1

    def __init__(self, integer=None, string=None):
        """ Initialize new instance

        :param integer: a value representing the alias as an integer.
        It can not be None while string is None. If it is None, a
        corresponding property of the object will be based on
        the string parameter
        :param string: a value representing the alias as a string.
        It can not be None while integer is None, and it has to consist
        only of characters specified in self._chars string.
        If it is None, a value of corresponding property of the object
        will be based on the integer parameter
        :raises AliasValueError: if the alias contains characters that
        are not part of self._chars string, or if both string and
        integer params are None
        """
        self.integer = integer
        if integer is None:
            if string is None:
                raise AliasValueError(
                    'The string and integer arguments cannot both be None'
                )

            self.integer = 0
            for exponent, char in enumerate(reversed(string)):
                digit_value = bisect_left(self._chars, char)
                if (digit_value == self._base or
                        self._chars[digit_value] != char):
                    raise AliasValueError(
                        "The string '{}' contains an unexpected"
                        " character: '{}' ".format(string, char)
                    )
                self.integer += digit_value*self._base**exponent
        self._string = string

    @classmethod
    def _get_min_int(cls, length):
        """ Get minimum integer value for alias string of given length

        :param length: a number of digits used to write the value
        to be returned in a base cls._base system
        :returns: the smallest integer value possible to be written
        as a numeral of given length in the system
        """
        return cls._base**(length - 1)

    @classmethod
    def _get_max_int(cls, length):
        """ Get maximum integer value for alias string of given length

        :param length: a number of digits used to write the value
        to be returned in a base cls._base system
        :returns: the largest integer value possible to be written
        as a numeral of given length in the system
        """
        return cls._get_min_int(length + 1) - 1

    @classmethod
    def init_random_factory(cls, min_length, max_length):
        """ Configure factory of randomly generated instances
        of the class

        The arguments provide a range of lengths for string
        representations of instances of Alias to be returned by
        a random alias factory provided by the class. For these values,
        the method calculates a range of integer values corresponding
        to them.

        From that range, only integers smaller than cls._max_int_32
        can be used to generate new aliases.

        :param min_length: a minimum number of characters for string
        representation of the instances returned by the factory
        :param max_length: a maximum number of characters for string
        representation of the instances returned by the factory
        :raises AliasLengthValueError: if:

        * values of any of the parameters are less than zero, or
        * min_length > max_length, or
        * the whole range of alias string lengths, or just part
        of it, is not available for generation due to integer values
        corresponding to it being larger than cls._max_int_32
        """
        if not 0 < min_length <= max_length:
            raise AliasLengthValueError('The length limits are incorrect')
        if cls._max_int_32 < cls._get_min_int(max_length):
            raise AliasLengthValueError(
                'The maximum length of a new alias is too large. The maximum'
                ' acceptable length is: {}'.format(cls._max_allowed_length)
            )

        cls._min_new_int = cls._get_min_int(min_length)
        cls._max_new_int = min(
            cls._max_int_32,
            cls._get_max_int(max_length)
        )

    @classmethod
    def create_random(cls):
        """ Get a randomly generated instance of the class

        The method uses pre-configured class properties
        cls._min_new_int and cls._max_new_int as paramteres for
        pseudo-random integer generation.
        """
        random_integer = randint(cls._min_new_int, cls._max_new_int)
        return cls(integer=random_integer)

    def __str__(self):
        """ Get string representation of the alias

        :returns: a string representing value of the alias as a numeral
        in the base cls._base system.
        """
        if self._string is None:
            value = ''
            integer = self.integer
            while True:
                integer, remainder = divmod(integer, self._base)
                value = self._chars[remainder] + value
                if integer == 0:
                    break
            self._string = value
        return self._string

class IntegerAlias(types.TypeDecorator):
    """ Converts between database integers and
    instances of Alias
    """

    impl = types.Integer

    def process_bind_param(self, value, dialect):
        return value.integer

    process_literal_param = process_bind_param

    def process_result_value(self, value, dialect):
        return Alias(integer=value)

class ShortenedURL(db.Model):
    """ Represents a URL for which a short alias has been created

    :cvar alias: a value representing a registered URL in short URLs and
    in database
    """
    alias = db.Column(IntegerAlias, primary_key=True)
    target = db.Column(db.String(2083), unique=True)

    def __init__(self, target):
        """ Constructor

        :param target: URL represented by the instance
        """
        self.target = target

    def __str__(self):
        return self.target

    def _alternative_url(self, endpoint):
        return url_for(endpoint, _external=True, alias=self.alias)

    @cached_property
    def short_url(self):
        return self._alternative_url('redirect_for')

    @cached_property
    def preview_url(self):
        return self._alternative_url('preview')

    @classmethod
    def get_or_create(cls, target_url):
        """ Find an existing shortened URL, or
        create a new one

        :param target_url: the target of shortened URL
        :return: an instance of ShortenedURL, existing or one
        to be registered
        """
        shortened_url = cls.query.filter_by(target=target_url).one_or_none()
        if shortened_url is None:
            shortened_url = cls(target_url)
        return shortened_url

    @classmethod
    def get_or_404(cls, alias):
        """ Find an existing shortened URL, or abort
        with 404 error code

        :param alias: a string representation of alias
        :raises AliasValueError: if the string representation
        of alias is not valid
        :return: an instance of ShortenedURL representing an
        existing shortened URL
        """
        try:
            valid_alias = Alias(string=alias)
        except AliasValueError:
            abort(404)
        return cls.query.get_or_404(valid_alias)

def register_if_new(shortened_url):
    """ Register a shortened URL object if it is new

    The registration is performed by assigning a randomly generated
    alias to the shortened URL and then persisting it.

    :param shortened_url: an instance of ShortenedURL to be registered
    :raises RegistrationRetryLimitExceeded: if the application exceeded
    the maximum number of attempts at shortening a URL,
    without success.

    Registered shortened URLs get aliases chosen randomly from a set
    of values with length falling between configurable minimum
    and maximum values. If a significant number of aliases from
    this set is already in use, exceeding the retry limit becomes more
    and more likely.
    """
    state = inspect(shortened_url)
    if not state.transient:
        return

    retry_limit = app.config['REGISTRATION_RETRY_LIMIT']
    for _ in range(retry_limit):
        shortened_url.alias = Alias.create_random()
        try:
            db.session.add(shortened_url)
            db.session.commit()
            return
        except IntegrityError as ex:
            db.session.rollback()
            msg = (
                'An integrity error occured during registration of'
                ' shortened URL: {}'.format(ex)
            )
            app.logger.warning(msg)
    msg_tpl = 'Registration retry limit of {} has been reached'
    raise RegistrationRetryLimitExceeded(msg_tpl.format(retry_limit))

def configure_random_factory(app_object):
    min_length = app_object.config['MIN_NEW_ALIAS_LENGTH']
    max_length = app_object.config['MAX_NEW_ALIAS_LENGTH']
    Alias.init_random_factory(min_length, max_length)

    msg_tpl = 'Length of newly generated aliases: from {} to {} characters.'
    app_object.logger.info(msg_tpl.format(min_length, max_length))
<commit_msg>Remove dependency of AliasType on application object<commit_after># -*- coding: utf-8 -*-
""" This module contains ShortenedURL class and related classes """
from bisect import bisect_left
from math import log, floor
from random import randint

from cached_property import cached_property
from flask import url_for, abort
from sqlalchemy import types, inspect
from sqlalchemy.exc import IntegrityError

from . import db, app

class AliasValueError(ValueError):
    """The value of alias is incorrect """

class AliasLengthValueError(ValueError):
    """The value of alias-length related parameter is incorrect """

class RegistrationRetryLimitExceeded(Exception):
    """The maximum number of attempts at retrying to register
    a new short URL has been exceeded """

class AliasType(type):
    def __new__(cls, *args, **kwargs):
        alias_class = type.__new__(cls, *args, **kwargs)
        alias_class.init_random_factory(1, 2)
        return alias_class

class Alias(metaclass=AliasType):
    """ An identifier for shortened URL

    In has two values used as its representations: a string value
    and an integer value, used in short URLs and in database,
    respectively.

    :cvar _chars: a string consisting of characters to be used
    in a string value of an alias.
    :cvar _base: a base of numeral system used to interpret string
    values as integers.
    :cvar _max_int_32: a maximum value of a 32 bit signed integer.

    This value is assumed as maximum allowed for the integers used
    in generation because we assume SQLAlchemy.types.Integer -
    a base type for alias property of ShortenedURL class - will
    translate into 32 bit signed integer type of underlying
    database engine used by the application.

    :cvar _max_allowed_length: maximum number of characters for
    newly generated alias strings for which there are at least
    some strings whose corresponding integer values are smaller
    or equal to _max_int_32.
    """

    _chars = '0123456789abcdefghijkmnopqrstuvwxyz'
    _base = len(_chars)
    _max_int_32 = 2**31 - 1
    _max_allowed_length = floor(log(_max_int_32, _base)) + 1

    def __init__(self, integer=None, string=None):
        """ Initialize new instance

        :param integer: a value representing the alias as an integer.
        It can not be None while string is None. If it is None, a
        corresponding property of the object will be based on
        the string parameter
        :param string: a value representing the alias as a string.
        It can not be None while integer is None, and it has to consist
        only of characters specified in self._chars string.
        If it is None, a value of corresponding property of the object
        will be based on the integer parameter
        :raises AliasValueError: if the alias contains characters that
        are not part of self._chars string, or if both string and
        integer params are None
        """
        self.integer = integer
        if integer is None:
            if string is None:
                raise AliasValueError(
                    'The string and integer arguments cannot both be None'
                )

            self.integer = 0
            for exponent, char in enumerate(reversed(string)):
                digit_value = bisect_left(self._chars, char)
                if (digit_value == self._base or
                        self._chars[digit_value] != char):
                    raise AliasValueError(
                        "The string '{}' contains an unexpected"
                        " character: '{}' ".format(string, char)
                    )
                self.integer += digit_value*self._base**exponent
        self._string = string

    @classmethod
    def _get_min_int(cls, length):
        """ Get minimum integer value for alias string of given length

        :param length: a number of digits used to write the value
        to be returned in a base cls._base system
        :returns: the smallest integer value possible to be written
        as a numeral of given length in the system
        """
        return cls._base**(length - 1)

    @classmethod
    def _get_max_int(cls, length):
        """ Get maximum integer value for alias string of given length

        :param length: a number of digits used to write the value
        to be returned in a base cls._base system
        :returns: the largest integer value possible to be written
        as a numeral of given length in the system
        """
        return cls._get_min_int(length + 1) - 1

    @classmethod
    def init_random_factory(cls, min_length, max_length):
        """ Configure factory of randomly generated instances
        of the class

        The arguments provide a range of lengths for string
        representations of instances of Alias to be returned by
        a random alias factory provided by the class. For these values,
        the method calculates a range of integer values corresponding
        to them.

        From that range, only integers smaller than cls._max_int_32
        can be used to generate new aliases.

        :param min_length: a minimum number of characters for string
        representation of the instances returned by the factory
        :param max_length: a maximum number of characters for string
        representation of the instances returned by the factory
        :raises AliasLengthValueError: if:

        * values of any of the parameters are less than zero, or
        * min_length > max_length, or
        * the whole range of alias string lengths, or just part
        of it, is not available for generation due to integer values
        corresponding to it being larger than cls._max_int_32
        """
        if not 0 < min_length <= max_length:
            raise AliasLengthValueError('The length limits are incorrect')
        if cls._max_int_32 < cls._get_min_int(max_length):
            raise AliasLengthValueError(
                'The maximum length of a new alias is too large. The maximum'
                ' acceptable length is: {}'.format(cls._max_allowed_length)
            )

        cls._min_new_int = cls._get_min_int(min_length)
        cls._max_new_int = min(
            cls._max_int_32,
            cls._get_max_int(max_length)
        )

    @classmethod
    def create_random(cls):
        """ Get a randomly generated instance of the class

        The method uses pre-configured class properties
        cls._min_new_int and cls._max_new_int as paramteres for
        pseudo-random integer generation.
        """
        random_integer = randint(cls._min_new_int, cls._max_new_int)
        return cls(integer=random_integer)

    def __str__(self):
        """ Get string representation of the alias

        :returns: a string representing value of the alias as a numeral
        in the base cls._base system.
        """
        if self._string is None:
            value = ''
            integer = self.integer
            while True:
                integer, remainder = divmod(integer, self._base)
                value = self._chars[remainder] + value
                if integer == 0:
                    break
            self._string = value
        return self._string

class IntegerAlias(types.TypeDecorator):
    """ Converts between database integers and
    instances of Alias
    """

    impl = types.Integer

    def process_bind_param(self, value, dialect):
        return value.integer

    process_literal_param = process_bind_param

    def process_result_value(self, value, dialect):
        return Alias(integer=value)

class ShortenedURL(db.Model):
    """ Represents a URL for which a short alias has been created

    :cvar alias: a value representing a registered URL in short URLs and
    in database
    """
    alias = db.Column(IntegerAlias, primary_key=True)
    target = db.Column(db.String(2083), unique=True)

    def __init__(self, target):
        """ Constructor

        :param target: URL represented by the instance
        """
        self.target = target

    def __str__(self):
        return self.target

    def _alternative_url(self, endpoint):
        return url_for(endpoint, _external=True, alias=self.alias)

    @cached_property
    def short_url(self):
        return self._alternative_url('redirect_for')

    @cached_property
    def preview_url(self):
        return self._alternative_url('preview')

    @classmethod
    def get_or_create(cls, target_url):
        """ Find an existing shortened URL, or
        create a new one

        :param target_url: the target of shortened URL
        :return: an instance of ShortenedURL, existing or one
        to be registered
        """
        shortened_url = cls.query.filter_by(target=target_url).one_or_none()
        if shortened_url is None:
            shortened_url = cls(target_url)
        return shortened_url

    @classmethod
    def get_or_404(cls, alias):
        """ Find an existing shortened URL, or abort
        with 404 error code

        :param alias: a string representation of alias
        :raises AliasValueError: if the string representation
        of alias is not valid
        :return: an instance of ShortenedURL representing an
        existing shortened URL
        """
        try:
            valid_alias = Alias(string=alias)
        except AliasValueError:
            abort(404)
        return cls.query.get_or_404(valid_alias)

def register_if_new(shortened_url):
    """ Register a shortened URL object if it is new

    The registration is performed by assigning a randomly generated
    alias to the shortened URL and then persisting it.

    :param shortened_url: an instance of ShortenedURL to be registered
    :raises RegistrationRetryLimitExceeded: if the application exceeded
    the maximum number of attempts at shortening a URL,
    without success.

    Registered shortened URLs get aliases chosen randomly from a set
    of values with length falling between configurable minimum
    and maximum values. If a significant number of aliases from
    this set is already in use, exceeding the retry limit becomes more
    and more likely.
    """
    state = inspect(shortened_url)
    if not state.transient:
        return

    retry_limit = app.config['REGISTRATION_RETRY_LIMIT']
    for _ in range(retry_limit):
        shortened_url.alias = Alias.create_random()
        try:
            db.session.add(shortened_url)
            db.session.commit()
            return
        except IntegrityError as ex:
            db.session.rollback()
            msg = (
                'An integrity error occured during registration of'
                ' shortened URL: {}'.format(ex)
            )
            app.logger.warning(msg)
    msg_tpl = 'Registration retry limit of {} has been reached'
    raise RegistrationRetryLimitExceeded(msg_tpl.format(retry_limit))

def configure_random_factory(app_object):
    min_length = app_object.config['MIN_NEW_ALIAS_LENGTH']
    max_length = app_object.config['MAX_NEW_ALIAS_LENGTH']
    Alias.init_random_factory(min_length, max_length)

    msg_tpl = 'Length of newly generated aliases: from {} to {} characters.'
    app_object.logger.info(msg_tpl.format(min_length, max_length))
<|endoftext|>


<commit_before>from canvas_sdk import client, utils

def list_accounts(request_ctx, per_page=None, as_user_id=None, **request_kwargs):
    """
    List accounts that the current user can view or manage.  Typically,
    students and even teachers will get an empty list in response, only
    account admins can view the accounts that they are in.

        :param request_ctx: The request context
        :type request_ctx: :class:RequestContext
        :param per_page: (optional) Set how many results canvas should return, defaults to config.LIMIT_PER_PAGE
        :type per_page: integer or None
        :param as_user_id: (optional) Masquerade as the given canvas user
        :type as_user_id: integer or None
        :return: List accounts
        :rtype: requests.Response (with array data)

    """

    if per_page is None:
        per_page = request_ctx.per_page
    path = '/v1/accounts'
    payload = {
        'per_page': per_page,
    }
    if as_user_id:
        payload['as_user_id'] = as_user_id
    url = request_ctx.base_api_url + path.format()
    response = client.get(request_ctx, url, payload=payload, **request_kwargs)

    return response

def get_single_account(request_ctx, id, **request_kwargs):
    """
    Retrieve information on an individual account, given by id or sis
    sis_account_id.

        :param request_ctx: The request context
        :type request_ctx: :class:RequestContext
        :param id: (required) ID
        :type id: string
        :return: Get a single account
        :rtype: requests.Response (with Account data)

    """

    path = '/v1/accounts/{id}'
    url = request_ctx.base_api_url + path.format(id=id)
    response = client.get(request_ctx, url, **request_kwargs)

    return response

def get_sub_accounts_of_account(request_ctx, account_id, recursive=None, per_page=None, as_user_id=None, **request_kwargs):
    """
    List accounts that are sub-accounts of the given account.

        :param request_ctx: The request context
        :type request_ctx: :class:RequestContext
        :param account_id: (required) ID
        :type account_id: string
        :param recursive: (optional) If true, the entire account tree underneath this account will be returned (though still paginated). If false, only direct sub-accounts of this account will be returned. Defaults to false.
        :type recursive: boolean or None
        :param per_page: (optional) Set how many results canvas should return, defaults to config.LIMIT_PER_PAGE
        :type per_page: integer or None
        :return: Get the sub-accounts of an account
        :rtype: requests.Response (with array data)

    """

    if per_page is None:
        per_page = request_ctx.per_page
    path = '/v1/accounts/{account_id}/sub_accounts'
    payload = {
        'recursive': recursive,
        'per_page': per_page,
    }
    if as_user_id:
        payload['as_user_id'] = as_user_id
    url = request_ctx.base_api_url + path.format(account_id=account_id)
    response = client.get(request_ctx, url, payload=payload, **request_kwargs)

    return response

def list_active_courses_in_account(request_ctx, account_id, with_enrollments=None, published=None, completed=None, by_teachers=None, by_subaccounts=None, hide_enrollmentless_courses=None, state=None, enrollment_term_id=None, search_term=None, include=None, per_page=None, **request_kwargs):
    """
    Retrieve the list of courses in this account.

        :param request_ctx: The request context
        :type request_ctx: :class:RequestContext
        :param account_id: (required) ID
        :type account_id: string
        :param with_enrollments: (optional) If true, include only courses with at least one enrollment. If false, include only courses with no enrollments. If not present, do not filter on course enrollment status.
        :type with_enrollments: boolean or None
        :param published: (optional) If true, include only published courses. If false, exclude published courses. If not present, do not filter on published status.
        :type published: boolean or None
        :param completed: (optional) If true, include only completed courses (these may be in state 'completed', or their enrollment term may have ended). If false, exclude completed courses. If not present, do not filter on completed status.
        :type completed: boolean or None
        :param by_teachers: (optional) List of User IDs of teachers; if supplied, include only courses taught by one of the referenced users.
        :type by_teachers: array or None
        :param by_subaccounts: (optional) List of Account IDs; if supplied, include only courses associated with one of the referenced subaccounts.
        :type by_subaccounts: array or None
        :param hide_enrollmentless_courses: (optional) If present, only return courses that have at least one enrollment. Equivalent to 'with_enrollments=true'; retained for compatibility.
        :type hide_enrollmentless_courses: boolean or None
        :param state: (optional) If set, only return courses that are in the given state(s). By default, all states but "deleted" are returned.
        :type state: array or None
        :param enrollment_term_id: (optional) If set, only includes courses from the specified term.
        :type enrollment_term_id: integer or None
        :param search_term: (optional) The partial course name, code, or full ID to match and return in the results list. Must be at least 3 characters.
        :type search_term: string or None
        :param include: (optional) - All explanations can be seen in the {api:CoursesController#index Course API index documentation}
        :type include: array or None
        :param per_page: (optional) Set how many results canvas should return, defaults to config.LIMIT_PER_PAGE
        :type per_page: integer or None
        :return: List active courses in an account
        :rtype: requests.Response (with array data)

    """

    if per_page is None:
        per_page = request_ctx.per_page
    state_types = ('created', 'claimed', 'available', 'completed', 'deleted', 'all')
    include_types = ('syllabus_body', 'term', 'course_progress', 'storage_quota_used_mb', 'total_students', 'teachers')
    utils.validate_attr_is_acceptable(state, state_types)
    utils.validate_attr_is_acceptable(include, include_types)
    path = '/v1/accounts/{account_id}/courses'
    payload = {
        'with_enrollments' : with_enrollments,
        'published' : published,
        'completed' : completed,
        'by_teachers' : by_teachers,
        'by_subaccounts' : by_subaccounts,
        'hide_enrollmentless_courses' : hide_enrollmentless_courses,
        'state' : state,
        'enrollment_term_id' : enrollment_term_id,
        'search_term' : search_term,
        'include' : include,
        'per_page' : per_page,
    }
    url = request_ctx.base_api_url + path.format(account_id=account_id)
    response = client.get(request_ctx, url, payload=payload, **request_kwargs)

    return response

def update_account(request_ctx, id, account_name=None, account_default_time_zone=None, account_default_storage_quota_mb=None, account_default_user_storage_quota_mb=None, account_default_group_storage_quota_mb=None, **request_kwargs):
    """
    Update an existing account.

        :param request_ctx: The request context
        :type request_ctx: :class:RequestContext
        :param id: (required) ID
        :type id: string
        :param account_name: (optional) Updates the account name
        :type account_name: string or None
        :param account_default_time_zone: (optional) The default time zone of the account. Allowed time zones are {http://www.iana.org/time-zones IANA time zones} or friendlier {http://api.rubyonrails.org/classes/ActiveSupport/TimeZone.html Ruby on Rails time zones}.
        :type account_default_time_zone: string or None
        :param account_default_storage_quota_mb: (optional) The default course storage quota to be used, if not otherwise specified.
        :type account_default_storage_quota_mb: integer or None
        :param account_default_user_storage_quota_mb: (optional) The default user storage quota to be used, if not otherwise specified.
        :type account_default_user_storage_quota_mb: integer or None
        :param account_default_group_storage_quota_mb: (optional) The default group storage quota to be used, if not otherwise specified.
        :type account_default_group_storage_quota_mb: integer or None
        :return: Update an account
        :rtype: requests.Response (with Account data)

    """

    path = '/v1/accounts/{id}'
    payload = {
        'account[name]' : account_name,
        'account[default_time_zone]' : account_default_time_zone,
        'account[default_storage_quota_mb]' : account_default_storage_quota_mb,
        'account[default_user_storage_quota_mb]' : account_default_user_storage_quota_mb,
        'account[default_group_storage_quota_mb]' : account_default_group_storage_quota_mb,
    }
    url = request_ctx.base_api_url + path.format(id=id)
    response = client.put(request_ctx, url, payload=payload, **request_kwargs)

    return response

def create_new_sub_account(request_ctx, account_id, account_name, account_default_storage_quota_mb=None, account_default_user_storage_quota_mb=None, account_default_group_storage_quota_mb=None, per_page=None, **request_kwargs):
    """
    Add a new sub-account to a given account.

        :param request_ctx: The request context
        :type request_ctx: :class:RequestContext
        :param account_id: (required) ID
        :type account_id: string
        :param account_name: (required) The name of the new sub-account.
        :type account_name: string
        :param account_default_storage_quota_mb: (optional) The default course storage quota to be used, if not otherwise specified.
        :type account_default_storage_quota_mb: integer or None
        :param account_default_user_storage_quota_mb: (optional) The default user storage quota to be used, if not otherwise specified.
        :type account_default_user_storage_quota_mb: integer or None
        :param account_default_group_storage_quota_mb: (optional) The default group storage quota to be used, if not otherwise specified.
        :type account_default_group_storage_quota_mb: integer or None
        :param per_page: (optional) Set how many results canvas should return, defaults to config.LIMIT_PER_PAGE
        :type per_page: integer or None
        :return: Create a new sub-account
        :rtype: requests.Response (with array data)

    """

    if per_page is None:
        per_page = request_ctx.per_page
    path = '/v1/accounts/{account_id}/sub_accounts'
    payload = {
        'account[name]' : account_name,
        'account[default_storage_quota_mb]' : account_default_storage_quota_mb,
        'account[default_user_storage_quota_mb]' : account_default_user_storage_quota_mb,
        'account[default_group_storage_quota_mb]' : account_default_group_storage_quota_mb,
        'per_page' : per_page,
    }
    url = request_ctx.base_api_url + path.format(account_id=account_id)
    response = client.post(request_ctx, url, payload=payload, **request_kwargs)

    return response
<commit_msg>Adds the sis_account_id option to the create_new_sub_account payload<commit_after>from canvas_sdk import client, utils

def list_accounts(request_ctx, per_page=None, as_user_id=None, **request_kwargs):
    """
    List accounts that the current user can view or manage.  Typically,
    students and even teachers will get an empty list in response, only
    account admins can view the accounts that they are in.

        :param request_ctx: The request context
        :type request_ctx: :class:RequestContext
        :param per_page: (optional) Set how many results canvas should return, defaults to config.LIMIT_PER_PAGE
        :type per_page: integer or None
        :param as_user_id: (optional) Masquerade as the given canvas user
        :type as_user_id: integer or None
        :return: List accounts
        :rtype: requests.Response (with array data)

    """

    if per_page is None:
        per_page = request_ctx.per_page
    path = '/v1/accounts'
    payload = {
        'per_page': per_page,
    }
    if as_user_id:
        payload['as_user_id'] = as_user_id
    url = request_ctx.base_api_url + path.format()
    response = client.get(request_ctx, url, payload=payload, **request_kwargs)

    return response

def get_single_account(request_ctx, id, **request_kwargs):
    """
    Retrieve information on an individual account, given by id or sis
    sis_account_id.

        :param request_ctx: The request context
        :type request_ctx: :class:RequestContext
        :param id: (required) ID
        :type id: string
        :return: Get a single account
        :rtype: requests.Response (with Account data)

    """

    path = '/v1/accounts/{id}'
    url = request_ctx.base_api_url + path.format(id=id)
    response = client.get(request_ctx, url, **request_kwargs)

    return response

def get_sub_accounts_of_account(request_ctx, account_id, recursive=None, per_page=None, as_user_id=None, **request_kwargs):
    """
    List accounts that are sub-accounts of the given account.

        :param request_ctx: The request context
        :type request_ctx: :class:RequestContext
        :param account_id: (required) ID
        :type account_id: string
        :param recursive: (optional) If true, the entire account tree underneath this account will be returned (though still paginated). If false, only direct sub-accounts of this account will be returned. Defaults to false.
        :type recursive: boolean or None
        :param per_page: (optional) Set how many results canvas should return, defaults to config.LIMIT_PER_PAGE
        :type per_page: integer or None
        :return: Get the sub-accounts of an account
        :rtype: requests.Response (with array data)

    """

    if per_page is None:
        per_page = request_ctx.per_page
    path = '/v1/accounts/{account_id}/sub_accounts'
    payload = {
        'recursive': recursive,
        'per_page': per_page,
    }
    if as_user_id:
        payload['as_user_id'] = as_user_id
    url = request_ctx.base_api_url + path.format(account_id=account_id)
    response = client.get(request_ctx, url, payload=payload, **request_kwargs)

    return response

def list_active_courses_in_account(request_ctx, account_id, with_enrollments=None, published=None, completed=None, by_teachers=None, by_subaccounts=None, hide_enrollmentless_courses=None, state=None, enrollment_term_id=None, search_term=None, include=None, per_page=None, **request_kwargs):
    """
    Retrieve the list of courses in this account.

        :param request_ctx: The request context
        :type request_ctx: :class:RequestContext
        :param account_id: (required) ID
        :type account_id: string
        :param with_enrollments: (optional) If true, include only courses with at least one enrollment. If false, include only courses with no enrollments. If not present, do not filter on course enrollment status.
        :type with_enrollments: boolean or None
        :param published: (optional) If true, include only published courses. If false, exclude published courses. If not present, do not filter on published status.
        :type published: boolean or None
        :param completed: (optional) If true, include only completed courses (these may be in state 'completed', or their enrollment term may have ended). If false, exclude completed courses. If not present, do not filter on completed status.
        :type completed: boolean or None
        :param by_teachers: (optional) List of User IDs of teachers; if supplied, include only courses taught by one of the referenced users.
        :type by_teachers: array or None
        :param by_subaccounts: (optional) List of Account IDs; if supplied, include only courses associated with one of the referenced subaccounts.
        :type by_subaccounts: array or None
        :param hide_enrollmentless_courses: (optional) If present, only return courses that have at least one enrollment. Equivalent to 'with_enrollments=true'; retained for compatibility.
        :type hide_enrollmentless_courses: boolean or None
        :param state: (optional) If set, only return courses that are in the given state(s). By default, all states but "deleted" are returned.
        :type state: array or None
        :param enrollment_term_id: (optional) If set, only includes courses from the specified term.
        :type enrollment_term_id: integer or None
        :param search_term: (optional) The partial course name, code, or full ID to match and return in the results list. Must be at least 3 characters.
        :type search_term: string or None
        :param include: (optional) - All explanations can be seen in the {api:CoursesController#index Course API index documentation}
        :type include: array or None
        :param per_page: (optional) Set how many results canvas should return, defaults to config.LIMIT_PER_PAGE
        :type per_page: integer or None
        :return: List active courses in an account
        :rtype: requests.Response (with array data)

    """

    if per_page is None:
        per_page = request_ctx.per_page
    state_types = ('created', 'claimed', 'available', 'completed', 'deleted', 'all')
    include_types = ('syllabus_body', 'term', 'course_progress', 'storage_quota_used_mb', 'total_students', 'teachers')
    utils.validate_attr_is_acceptable(state, state_types)
    utils.validate_attr_is_acceptable(include, include_types)
    path = '/v1/accounts/{account_id}/courses'
    payload = {
        'with_enrollments' : with_enrollments,
        'published' : published,
        'completed' : completed,
        'by_teachers' : by_teachers,
        'by_subaccounts' : by_subaccounts,
        'hide_enrollmentless_courses' : hide_enrollmentless_courses,
        'state' : state,
        'enrollment_term_id' : enrollment_term_id,
        'search_term' : search_term,
        'include' : include,
        'per_page' : per_page,
    }
    url = request_ctx.base_api_url + path.format(account_id=account_id)
    response = client.get(request_ctx, url, payload=payload, **request_kwargs)

    return response

def update_account(request_ctx, id, account_name=None, account_default_time_zone=None, account_default_storage_quota_mb=None, account_default_user_storage_quota_mb=None, account_default_group_storage_quota_mb=None, **request_kwargs):
    """
    Update an existing account.

        :param request_ctx: The request context
        :type request_ctx: :class:RequestContext
        :param id: (required) ID
        :type id: string
        :param account_name: (optional) Updates the account name
        :type account_name: string or None
        :param account_default_time_zone: (optional) The default time zone of the account. Allowed time zones are {http://www.iana.org/time-zones IANA time zones} or friendlier {http://api.rubyonrails.org/classes/ActiveSupport/TimeZone.html Ruby on Rails time zones}.
        :type account_default_time_zone: string or None
        :param account_default_storage_quota_mb: (optional) The default course storage quota to be used, if not otherwise specified.
        :type account_default_storage_quota_mb: integer or None
        :param account_default_user_storage_quota_mb: (optional) The default user storage quota to be used, if not otherwise specified.
        :type account_default_user_storage_quota_mb: integer or None
        :param account_default_group_storage_quota_mb: (optional) The default group storage quota to be used, if not otherwise specified.
        :type account_default_group_storage_quota_mb: integer or None
        :return: Update an account
        :rtype: requests.Response (with Account data)

    """

    path = '/v1/accounts/{id}'
    payload = {
        'account[name]' : account_name,
        'account[default_time_zone]' : account_default_time_zone,
        'account[default_storage_quota_mb]' : account_default_storage_quota_mb,
        'account[default_user_storage_quota_mb]' : account_default_user_storage_quota_mb,
        'account[default_group_storage_quota_mb]' : account_default_group_storage_quota_mb,
    }
    url = request_ctx.base_api_url + path.format(id=id)
    response = client.put(request_ctx, url, payload=payload, **request_kwargs)

    return response

def create_new_sub_account(request_ctx, account_id, account_name, sis_account_id, account_default_storage_quota_mb=None, account_default_user_storage_quota_mb=None, account_default_group_storage_quota_mb=None, per_page=None, **request_kwargs):
    """
    Add a new sub-account to a given account.

        :param request_ctx: The request context
        :type request_ctx: :class:RequestContext
        :param account_id: (required) ID
        :type account_id: string
        :param account_name: (required) The name of the new sub-account.
        :type account_name: string
        :param account_default_storage_quota_mb: (optional) The default course storage quota to be used, if not otherwise specified.
        :type account_default_storage_quota_mb: integer or None
        :param account_default_user_storage_quota_mb: (optional) The default user storage quota to be used, if not otherwise specified.
        :type account_default_user_storage_quota_mb: integer or None
        :param account_default_group_storage_quota_mb: (optional) The default group storage quota to be used, if not otherwise specified.
        :type account_default_group_storage_quota_mb: integer or None
        :param per_page: (optional) Set how many results canvas should return, defaults to config.LIMIT_PER_PAGE
        :type per_page: integer or None
        :return: Create a new sub-account
        :rtype: requests.Response (with array data)

    """

    if per_page is None:
        per_page = request_ctx.per_page
    path = '/v1/accounts/{account_id}/sub_accounts'
    payload = {
        'account[name]' : account_name,
        'account[sis_account_id]': sis_account_id,
        'account[default_storage_quota_mb]' : account_default_storage_quota_mb,
        'account[default_user_storage_quota_mb]' : account_default_user_storage_quota_mb,
        'account[default_group_storage_quota_mb]' : account_default_group_storage_quota_mb,
        'per_page' : per_page,
    }
    url = request_ctx.base_api_url + path.format(account_id=account_id)
    response = client.post(request_ctx, url, payload=payload, **request_kwargs)

    return response
<|endoftext|>


<commit_before>from __future__ import unicode_literals
from bst import BST
import pytest

tree_val = [15, 7, 14, 16, 1, 5]

second_tree_val = [15, 7, 16, 5, 14, 3]

@pytest.fixture()
def empty_tree():
    tree = BST()
    return tree

@pytest.fixture()
def full_tree():
    tree = BST()
    for val in tree_val:
        tree.insert(val)
    return tree

@pytest.fixture()
def second_full_tree():
    tree = BST()
    for val in second_tree_val:
        tree.insert(val)
    return tree

def test_contains(full_tree):
    for val in tree_val:
        assert full_tree.contains(val)

def test_size(full_tree, empty_tree):
    assert full_tree.size() == (len(tree_val))
    assert empty_tree.size() == 0

def test_insert(full_tree, empty_tree):
    for val in tree_val:
        assert full_tree.contains(val)
    full_tree.insert(22)
    assert full_tree.contains(22)
    assert full_tree.size() == (len(tree_val) + 1)
    full_tree.insert(22)
    assert full_tree.size() == (len(tree_val) + 1)
    empty_tree.insert(5)
    assert empty_tree.size() == 1
    assert empty_tree.contains(5)

def test_depth(full_tree, empty_tree):
    assert full_tree.depth() == 4
    assert empty_tree.depth() == 0

def test_balance(full_tree, empty_tree):
    assert full_tree.balance() > 0
    assert empty_tree.balance() == 0
    empty_tree.insert(5)
    assert empty_tree.balance() == 0
    empty_tree.insert(8)
    assert empty_tree.balance() < 0
    empty_tree.insert(4)
    assert empty_tree.balance() == 0
    empty_tree.insert(1)
    assert empty_tree.balance() > 0

def test_pre_order(full_tree):
    g = full_tree.pre_order()
    assert list(g) == [15, 7, 1, 5, 14, 16]

def test_in_order(full_tree):
    g = full_tree.in_order()
    assert list(g) == [1, 5, 7, 14, 15, 16]

def test_post_order(full_tree):
    g = full_tree.post_order()
    assert list(g) == [5, 1, 14, 7, 16, 15]

def test_breadth_first(full_tree):
    g = full_tree.breadth_first()
    assert list(g) == [15, 7, 16, 1, 14, 5]

def test_remove_leaf(full_tree):
    full_tree.remove(16)
    assert full_tree.size() == 5
    g = full_tree.in_order()
    assert list(g) == [1, 5, 7, 14, 15]
    assert full_tree.contains(16) is False

def test_remove_node_with_one_right_child(full_tree):
    full_tree.remove(1)
    assert full_tree.size() == 5
    g = full_tree.in_order()
    assert list(g) == [5, 7, 14, 15, 16]
    assert full_tree.contains(1) is False

def test_remove_node_with_one_left_child(second_full_tree):
    second_full_tree.remove(5)
    assert second_full_tree.size() == 5
    g = second_full_tree.in_order()
    assert list(g) == [3, 7, 14, 15, 16]
    assert second_full_tree.contains(5) is False

def test_remove_node_with_two_child(full_tree):
    full_tree.remove(7)
    assert full_tree.size() == 5
    g = full_tree.in_order()
    assert list(g) == [1, 5, 14, 15, 16]
    assert full_tree.contains(7) is False

def test_remove_root_node(full_tree):
    full_tree.remove(15)
    assert full_tree.size() == 5
    g = full_tree.in_order()
    assert list(g) == [1, 5, 7, 14, 16]
    assert full_tree.contains(15) is False
<commit_msg>add remove root node test<commit_after>from __future__ import unicode_literals
from bst import BST
import pytest

tree_val = [15, 7, 14, 16, 1, 5]

second_tree_val = [15, 7, 16, 5, 14, 3]

@pytest.fixture()
def empty_tree():
    tree = BST()
    return tree

@pytest.fixture()
def full_tree():
    tree = BST()
    for val in tree_val:
        tree.insert(val)
    return tree

@pytest.fixture()
def second_full_tree():
    tree = BST()
    for val in second_tree_val:
        tree.insert(val)
    return tree

def test_contains(full_tree):
    for val in tree_val:
        assert full_tree.contains(val)

def test_size(full_tree, empty_tree):
    assert full_tree.size() == (len(tree_val))
    assert empty_tree.size() == 0

def test_insert(full_tree, empty_tree):
    for val in tree_val:
        assert full_tree.contains(val)
    full_tree.insert(22)
    assert full_tree.contains(22)
    assert full_tree.size() == (len(tree_val) + 1)
    full_tree.insert(22)
    assert full_tree.size() == (len(tree_val) + 1)
    empty_tree.insert(5)
    assert empty_tree.size() == 1
    assert empty_tree.contains(5)

def test_depth(full_tree, empty_tree):
    assert full_tree.depth() == 4
    assert empty_tree.depth() == 0

def test_balance(full_tree, empty_tree):
    assert full_tree.balance() > 0
    assert empty_tree.balance() == 0
    empty_tree.insert(5)
    assert empty_tree.balance() == 0
    empty_tree.insert(8)
    assert empty_tree.balance() < 0
    empty_tree.insert(4)
    assert empty_tree.balance() == 0
    empty_tree.insert(1)
    assert empty_tree.balance() > 0

def test_pre_order(full_tree):
    g = full_tree.pre_order()
    assert list(g) == [15, 7, 1, 5, 14, 16]

def test_in_order(full_tree):
    g = full_tree.in_order()
    assert list(g) == [1, 5, 7, 14, 15, 16]

def test_post_order(full_tree):
    g = full_tree.post_order()
    assert list(g) == [5, 1, 14, 7, 16, 15]

def test_breadth_first(full_tree):
    g = full_tree.breadth_first()
    assert list(g) == [15, 7, 16, 1, 14, 5]

def test_remove_leaf(full_tree):
    full_tree.remove(16)
    assert full_tree.size() == 5
    g = full_tree.in_order()
    assert list(g) == [1, 5, 7, 14, 15]
    assert full_tree.contains(16) is False

def test_remove_node_with_one_right_child(full_tree):
    full_tree.remove(1)
    assert full_tree.size() == 5
    g = full_tree.in_order()
    assert list(g) == [5, 7, 14, 15, 16]
    assert full_tree.contains(1) is False

def test_remove_node_with_one_left_child(second_full_tree):
    second_full_tree.remove(5)
    assert second_full_tree.size() == 5
    g = second_full_tree.in_order()
    assert list(g) == [3, 7, 14, 15, 16]
    assert second_full_tree.contains(5) is False

def test_remove_node_with_two_child(full_tree):
    full_tree.remove(7)
    assert full_tree.size() == 5
    g = full_tree.in_order()
    assert list(g) == [1, 5, 14, 15, 16]
    assert full_tree.contains(7) is False
<|endoftext|>