<commit_before>""" shRUnk - Rutgers University URL Shortener

Sets up a Flask application for the main web server.
"""
from flask import Flask, render_template, make_response, request, redirect, g
from flask_login import LoginManager, login_required, current_user, logout_user
from flask_auth import Auth

from shrunk.client import ShrunkCursor
from shrunk.forms import BlockLinksForm
from shrunk.forms import LinkForm, RULoginForm, BlacklistUserForm, AddAdminForm
from shrunk.user import User, get_user, admin_required
from shrunk.util import get_db_client, set_logger, formattime

# Create application
app = Flask(__name__)

# Import settings in config.py
app.config.from_pyfile("config.py", silent=True)
app.secret_key = app.config['SECRET_KEY']

# Initialize logging
set_logger(app)

# Initialize login manager
login_manager = LoginManager()
login_manager.init_app(app)
login_manager.login_view = '/login'

# Allows us to use the function in our templates
app.jinja_env.globals.update(formattime=formattime)

@login_manager.user_loader
def load_user(userid):
    """Loads user object for login.

    :Parameters:
      - `userid`: An id for the user (typically a NetID).
    """
    return User(userid)

def render_login(**kwargs):
    """Renders the login template.

    Takes a WTForm in the keyword arguments.
    """
    return render_template('login.html', **kwargs)

def login_success(user):
    """Function executed on successful login.

    Redirects the user to the homepage.

    :Parameters:
      - `user`: The user that has logged in.
    """
    return redirect('/')

def unauthorized_admin():
    return redirect("/")

### Views ###
try:
    if app.config["DUAL_SERVER"]:
        @app.route("/<short_url>")
        def redirect_link(short_url):
            """Redirects to the short URL's true destination.

            This looks up the short URL's destination in the database and performs a
            redirect, logging some information at the same time. If no such link
            exists, a not found page is shown.

            :Parameters:
              - `short_url`: A string containing a shrunk-ified URL.
            """
            client = get_db_client(app, g)
            app.logger.info("{} requests {}".format(request.remote_addr, short_url))

            # Perform a lookup and redirect
            long_url = client.get_long_url(short_url)
            if long_url is None:
                return render_template("link-404.html", short_url=short_url)
            else:
                client.visit(short_url, request.remote_addr)
                # Check if a protocol exists
                if "://" in long_url:
                    return redirect(long_url)
                else:
                    return redirect("http://{}".format(long_url))
except KeyError:
    # No setting in the config file
    pass

@app.route("/")
def render_index(**kwargs):
    """Renders the homepage.

    Renders the homepage for the current user. By default, this renders all of
    the links owned by them. If a search has been made, then only the links
    matching their search query are shown.
    """
    client = get_db_client(app, g)
    if not hasattr(current_user, "netid"):
        # Anonymous user
        return redirect("/login")

    # Grab the current page number
    try:
        page = int(request.args["p"])
    except:
        page = 0

    # If this exists, execute a search query
    try:
        query = request.args["search"]
    except:
        query = ""

    # Display all users or just the current administrator?
    try:
        all_users = request.args["all_users"]
    except:
        if "all_users" in request.cookies:
            all_users = request.cookies["all_users"]
        else:
            all_users = "1"

    # Change sorting preferences
    if "sortby" in request.args:
        sortby = request.args["sortby"]
    elif "sortby" in request.cookies:
        sortby = request.cookies["sortby"]
    else:
        sortby = "0"

    # Depending on the type of user, get info from the database
    is_admin = not current_user.is_anonymous() and current_user.is_admin()
    if is_admin:
        netid = current_user.netid
        if query:
            cursor = client.search(query)
        elif all_users == "1":
            cursor = client.get_all_urls(query)
        else:
            cursor = client.get_urls(current_user.netid)
    else:
        netid = current_user.netid
        if query:
            cursor = client.search(query, netid=netid)
            app.logger.info("search: {}, '{}'".format(netid, query))
        else:
            cursor = client.get_urls(current_user.netid)
            app.logger.info("render index: {}".format(netid))

    # Perform sorting, pagination and get the results
    cursor.sort(sortby)
    page, lastpage = cursor.paginate(page, app.config["MAX_DISPLAY_LINKS"])
    links = cursor.get_results()

    resp = make_response(
            render_template("index.html",
                            admin=is_admin,
                            all_users=all_users,
                            lastpage=lastpage,
                            links=links,
                            linkserver_url=app.config["LINKSERVER_URL"],
                            netid=netid,
                            page=page,
                            query=query,
                            sortby=sortby,
                            **kwargs))
    resp.set_cookie("all_users", all_users)
    resp.set_cookie("sortby", sortby)
    return resp

@app.route("/login", methods=['GET', 'POST'])
def login():
    """Handles authentication."""
    a = Auth(app.config['AUTH'], get_user)
    return a.login(request, RULoginForm, render_login, login_success)

@app.route("/logout")
@login_required
def logout():
    """Handles logging out."""
    logout_user()
    return redirect('/')

@app.route("/add", methods=["GET", "POST"])
@login_required
def add_link():
    """Adds a new link for the current user."""
    form = LinkForm(request.form)
    client = get_db_client(app, g)

    if request.method == "POST":
        # Validate the form

        if form.validate():
            # TODO Handle an error on db insert
            kwargs = form.to_json()
            response = client.create_short_url(
                netid=current_user.netid,
                **kwargs
            )
            if not response:
                # Specifically, there is no response from the database
                return render_template("add.html",
                                       errors=["Blocked Link"],
                                       admin=current_user.is_admin()
                )
            else:
                # Success
                return redirect("/")
        else: # WTForms detects a form validation error
            form.long_url.data = "http://" + form.long_url.data

            if form.validate():
                kwargs = form.to_json()
                response = client.create_short_url(
                        netid=current_user.netid,
                        **kwargs
                )
                if not response:
                    # Specifically, there is no response from the database
                    return render_template("add.html",
                                        errors=["Blocked Link"],
                                        admin=current_user.is_admin()
                    )
                else:
                    # Success
                    return redirect("/")
            else:
                return render_template("add.html",
                                    errors=form.errors,
                                    netid=current_user.netid,
                                    admin=current_user.is_admin())
    else: # GET request
        if not request.form:
            form = LinkForm()

        return render_template("add.html",
                               netid=current_user.netid,
                               admin=current_user.is_admin())

@app.route("/delete", methods=["GET", "POST"])
@login_required
def delete_link():
    """Deletes a link."""
    client = get_db_client(app, g)

    # TODO Handle the response intelligently, or put that logic somewhere else
    if request.method == "POST":
        app.logger.info("Deleting URL: {}".format(request.form["short_url"]))
        client.delete_url(request.form["short_url"])
    return redirect("/")

@app.route("/edit", methods=["GET", "POST"])
@login_required
def edit_link():
    """Edits a link.

    On POST, this route expects a form that contains the unique short URL that
    will be edited.
    """
    client = get_db_client(app, g)
    form = LinkForm(request.form)

    if request.method == "POST":
        # Validate form before continuing
        if form.validate():
            # Success - make the edits in the database
            kwargs = form.to_json()
            response = client.modify_url(
                request.form["short_url"],
                **kwargs
            )
            return redirect("/")
        else:
            form.long_url.data = "http://" + form.long_url.data
            if form.validate():
                kwargs = form.to_json()
                response = client.modify_url(
                    request.form["short_url"],
                    **kwargs
                )
                return redirect("/")
            else:
                # Validation error
                short_url = request.form["short_url"]
                info = client.get_url_info(short_url)
                long_url = info["long_url"]
                title = info["title"]
                return render_template("edit.html",
                                    errors=form.errors,
                                    netid=current_user.netid,
                                    title=title,
                                    short_url=short_url,
                                    long_url=long_url)
    else: # GET request
        # Hit the database to get information
        short_url = request.args["url"]
        info = client.get_url_info(short_url)
        owner = info["netid"]
        if owner != current_user.netid:
            return render_index(wrong_owner=True)

        long_url = info["long_url"]
        title = info["title"]
        # Render the edit template
        return render_template("edit.html", netid=current_user.netid,
                                            title=title,
                                            short_url=short_url,
                                            long_url=long_url)

@app.route("/admin/")
@app.route("/admin/<action>", methods=["GET", "POST"])
@login_required
@admin_required(unauthorized_admin)
def admin_sub(action=None):
    """Renders the admin interface.
    :Parameters:
      - `action`: Which action to take. This can be one of the following:
        1. blacklist - Go to blacklist panel used for blacklisting users
        2. add - Go to add panel used for adding additional admins
        3. blocklink - Go to block link panel, used for blacklisting long urls
    """

    netid = current_user.netid
    if action == None:
        return render_template("admin.html", netid=netid)
    client = get_db_client(app, g)
    if action == 'blacklist':
        form = BlacklistUserForm(request.form)
        if request.method == "POST" and form.validate():
            if form.action.data == 'ban':
                res = client.blacklist_user(form.netid.data, netid)
            else:
                res = client.allow_user(form.netid.data)
            return render_template('admin_blacklist.html', form=form,
                                   netid=netid, msg='Success!')
        return render_template('admin_blacklist.html', netid=netid, form=form)

    elif action == 'add':
        form = AddAdminForm(request.form)
        if request.method == "POST" and form.validate():
            res = client.add_admin(form.netid.data, current_user.netid)
            return render_template('admin_add.html', form=form, netid=netid,
                                   msg='Success!')
        return render_template('admin_add.html', form=form, netid=netid)

    elif action == 'blocklink':
        form = BlockLinksForm(request.form)
        if request.method == "POST" and form.validate():
            if form.action.data == 'block':
                res = client.block_link(form.link.data, netid)
            else:
                res = client.allow_link(form.link.data)
            return render_template('admin_block_links.html', form=form,
                                   msg='Success!', netid=netid)
        return render_template('admin_block_links.html', form=form,
                               netid=netid)

    else:
        return redirect('/')
<commit_msg>Add prevention of adding/editing linkserver URLs<commit_after>""" shRUnk - Rutgers University URL Shortener

Sets up a Flask application for the main web server.
"""
from flask import Flask, render_template, make_response, request, redirect, g
from flask_login import LoginManager, login_required, current_user, logout_user
from flask_auth import Auth

from shrunk.client import ShrunkCursor
from shrunk.forms import BlockLinksForm
from shrunk.forms import LinkForm, RULoginForm, BlacklistUserForm, AddAdminForm
from shrunk.user import User, get_user, admin_required
from shrunk.util import get_db_client, set_logger, formattime

# Create application
app = Flask(__name__)

# Import settings in config.py
app.config.from_pyfile("config.py", silent=True)
app.secret_key = app.config['SECRET_KEY']

# Initialize logging
set_logger(app)

# Initialize login manager
login_manager = LoginManager()
login_manager.init_app(app)
login_manager.login_view = '/login'

# Allows us to use the function in our templates
app.jinja_env.globals.update(formattime=formattime)

@login_manager.user_loader
def load_user(userid):
    """Loads user object for login.

    :Parameters:
      - `userid`: An id for the user (typically a NetID).
    """
    return User(userid)

def render_login(**kwargs):
    """Renders the login template.

    Takes a WTForm in the keyword arguments.
    """
    return render_template('login.html', **kwargs)

def login_success(user):
    """Function executed on successful login.

    Redirects the user to the homepage.

    :Parameters:
      - `user`: The user that has logged in.
    """
    return redirect('/')

def unauthorized_admin():
    return redirect("/")

### Views ###
try:
    if app.config["DUAL_SERVER"]:
        @app.route("/<short_url>")
        def redirect_link(short_url):
            """Redirects to the short URL's true destination.

            This looks up the short URL's destination in the database and performs a
            redirect, logging some information at the same time. If no such link
            exists, a not found page is shown.

            :Parameters:
              - `short_url`: A string containing a shrunk-ified URL.
            """
            client = get_db_client(app, g)
            app.logger.info("{} requests {}".format(request.remote_addr, short_url))

            # Perform a lookup and redirect
            long_url = client.get_long_url(short_url)
            if long_url is None:
                return render_template("link-404.html", short_url=short_url)
            else:
                client.visit(short_url, request.remote_addr)
                # Check if a protocol exists
                if "://" in long_url:
                    return redirect(long_url)
                else:
                    return redirect("http://{}".format(long_url))
except KeyError:
    # No setting in the config file
    pass

@app.route("/")
def render_index(**kwargs):
    """Renders the homepage.

    Renders the homepage for the current user. By default, this renders all of
    the links owned by them. If a search has been made, then only the links
    matching their search query are shown.
    """
    client = get_db_client(app, g)
    if not hasattr(current_user, "netid"):
        # Anonymous user
        return redirect("/login")

    # Grab the current page number
    try:
        page = int(request.args["p"])
    except:
        page = 0

    # If this exists, execute a search query
    try:
        query = request.args["search"]
    except:
        query = ""

    # Display all users or just the current administrator?
    try:
        all_users = request.args["all_users"]
    except:
        if "all_users" in request.cookies:
            all_users = request.cookies["all_users"]
        else:
            all_users = "1"

    # Change sorting preferences
    if "sortby" in request.args:
        sortby = request.args["sortby"]
    elif "sortby" in request.cookies:
        sortby = request.cookies["sortby"]
    else:
        sortby = "0"

    # Depending on the type of user, get info from the database
    is_admin = not current_user.is_anonymous() and current_user.is_admin()
    if is_admin:
        netid = current_user.netid
        if query:
            cursor = client.search(query)
        elif all_users == "1":
            cursor = client.get_all_urls(query)
        else:
            cursor = client.get_urls(current_user.netid)
    else:
        netid = current_user.netid
        if query:
            cursor = client.search(query, netid=netid)
            app.logger.info("search: {}, '{}'".format(netid, query))
        else:
            cursor = client.get_urls(current_user.netid)
            app.logger.info("render index: {}".format(netid))

    # Perform sorting, pagination and get the results
    cursor.sort(sortby)
    page, lastpage = cursor.paginate(page, app.config["MAX_DISPLAY_LINKS"])
    links = cursor.get_results()

    resp = make_response(
            render_template("index.html",
                            admin=is_admin,
                            all_users=all_users,
                            lastpage=lastpage,
                            links=links,
                            linkserver_url=app.config["LINKSERVER_URL"],
                            netid=netid,
                            page=page,
                            query=query,
                            sortby=sortby,
                            **kwargs))
    resp.set_cookie("all_users", all_users)
    resp.set_cookie("sortby", sortby)
    return resp

@app.route("/login", methods=['GET', 'POST'])
def login():
    """Handles authentication."""
    a = Auth(app.config['AUTH'], get_user)
    return a.login(request, RULoginForm, render_login, login_success)

@app.route("/logout")
@login_required
def logout():
    """Handles logging out."""
    logout_user()
    return redirect('/')

@app.route("/add", methods=["GET", "POST"])
@login_required
def add_link():
    """Adds a new link for the current user."""
    form = LinkForm(request.form)
    client = get_db_client(app, g)

    if request.method == "POST":
        # Validate the form

        # long url may not be the linkserver
        if form.long_url.data.startswith(app.config['LINKSERVER_URL']) or form.long_url.data.startswith(app.config['LINKSERVER_URL'][:7]):
            return render_template("add.html",
                                    errors=["Blocked Link"],
                                    admin=current_user.is_admin()
            )
        if form.validate():
            # TODO Handle an error on db insert
            kwargs = form.to_json()
            response = client.create_short_url(
                netid=current_user.netid,
                **kwargs
            )
            if not response:
                # Specifically, there is no response from the database
                return render_template("add.html",
                                       errors=["Blocked Link"],
                                       admin=current_user.is_admin()
                )
            else:
                # Success
                return redirect("/")
        else: # WTForms detects a form validation error
            form.long_url.data = "http://" + form.long_url.data

            if form.validate():
                kwargs = form.to_json()
                response = client.create_short_url(
                        netid=current_user.netid,
                        **kwargs
                )
                if not response:
                    # Specifically, there is no response from the database
                    return render_template("add.html",
                                        errors=["Blocked Link"],
                                        admin=current_user.is_admin()
                    )
                else:
                    # Success
                    return redirect("/")
            else:
                return render_template("add.html",
                                    errors=form.errors,
                                    netid=current_user.netid,
                                    admin=current_user.is_admin())
    else: # GET request
        if not request.form:
            form = LinkForm()

        return render_template("add.html",
                               netid=current_user.netid,
                               admin=current_user.is_admin())

@app.route("/delete", methods=["GET", "POST"])
@login_required
def delete_link():
    """Deletes a link."""
    client = get_db_client(app, g)

    # TODO Handle the response intelligently, or put that logic somewhere else
    if request.method == "POST":
        app.logger.info("Deleting URL: {}".format(request.form["short_url"]))
        client.delete_url(request.form["short_url"])
    return redirect("/")

@app.route("/edit", methods=["GET", "POST"])
@login_required
def edit_link():
    """Edits a link.

    On POST, this route expects a form that contains the unique short URL that
    will be edited.
    """
    client = get_db_client(app, g)
    form = LinkForm(request.form)

    if request.method == "POST":
        # Validate form before continuing
        if form.long_url.data.startswith(app.config['LINKSERVER_URL']) or form.long_url.data.startswith(app.config['LINKSERVER_URL'][:7]):
            return render_template("add.html",
                                    errors=["Blocked Link"],
                                    admin=current_user.is_admin()
            )
        if form.validate():
            # Success - make the edits in the database
            kwargs = form.to_json()
            response = client.modify_url(
                request.form["short_url"],
                **kwargs
            )
            return redirect("/")
        else:
            form.long_url.data = "http://" + form.long_url.data
            if form.validate():
                kwargs = form.to_json()
                response = client.modify_url(
                    request.form["short_url"],
                    **kwargs
                )
                return redirect("/")
            else:
                # Validation error
                short_url = request.form["short_url"]
                info = client.get_url_info(short_url)
                long_url = info["long_url"]
                title = info["title"]
                return render_template("edit.html",
                                    errors=form.errors,
                                    netid=current_user.netid,
                                    title=title,
                                    short_url=short_url,
                                    long_url=long_url)
    else: # GET request
        # Hit the database to get information
        short_url = request.args["url"]
        info = client.get_url_info(short_url)
        owner = info["netid"]
        if owner != current_user.netid:
            return render_index(wrong_owner=True)

        long_url = info["long_url"]
        title = info["title"]
        # Render the edit template
        return render_template("edit.html", netid=current_user.netid,
                                            title=title,
                                            short_url=short_url,
                                            long_url=long_url)

@app.route("/admin/")
@app.route("/admin/<action>", methods=["GET", "POST"])
@login_required
@admin_required(unauthorized_admin)
def admin_sub(action=None):
    """Renders the admin interface.
    :Parameters:
      - `action`: Which action to take. This can be one of the following:
        1. blacklist - Go to blacklist panel used for blacklisting users
        2. add - Go to add panel used for adding additional admins
        3. blocklink - Go to block link panel, used for blacklisting long urls
    """

    netid = current_user.netid
    if action == None:
        return render_template("admin.html", netid=netid)
    client = get_db_client(app, g)
    if action == 'blacklist':
        form = BlacklistUserForm(request.form)
        if request.method == "POST" and form.validate():
            if form.action.data == 'ban':
                res = client.blacklist_user(form.netid.data, netid)
            else:
                res = client.allow_user(form.netid.data)
            return render_template('admin_blacklist.html', form=form,
                                   netid=netid, msg='Success!')
        return render_template('admin_blacklist.html', netid=netid, form=form)

    elif action == 'add':
        form = AddAdminForm(request.form)
        if request.method == "POST" and form.validate():
            res = client.add_admin(form.netid.data, current_user.netid)
            return render_template('admin_add.html', form=form, netid=netid,
                                   msg='Success!')
        return render_template('admin_add.html', form=form, netid=netid)

    elif action == 'blocklink':
        form = BlockLinksForm(request.form)
        if request.method == "POST" and form.validate():
            if form.action.data == 'block':
                res = client.block_link(form.link.data, netid)
            else:
                res = client.allow_link(form.link.data)
            return render_template('admin_block_links.html', form=form,
                                   msg='Success!', netid=netid)
        return render_template('admin_block_links.html', form=form,
                               netid=netid)

    else:
        return redirect('/')
<|endoftext|>


<commit_before><commit_msg>Enable WAL<commit_after>from django.db.backends.signals import connection_created

def activate_wal(sender, connection, **kwargs):
    """Enable WAL mode."""
    if connection.vendor == 'sqlite':
        cursor = connection.cursor()
        cursor.execute('PRAGMA journal_mode=WAL;')

connection_created.connect(activate_foreign_keys)<|endoftext|>


<commit_before>#!/usr/bin/env python
# encoding: UTF-8
# Copyright Kapsi Internet-käyttäjät ry
# Idea from http://code.google.com/p/python-iptools ; code completely new

import socket
import struct

def ipv6_to_long(ip):
    """Return the IPv6 address string as a long

    >>> ipv6_to_long("2001:db8::1")
    42540766411282592856903984951653826561L
    >>> ipv6_to_long("::1")
    1L
    """
    ip_bytes_n = socket.inet_pton(socket.AF_INET6, ip)
    ip_parts = struct.unpack('!QQ', ip_bytes_n)
    return 2**64 * ip_parts[0] + ip_parts[1]

def ipv4_to_long(ip):
    """Return the IPv4 address string as a long

    >>> ipv4_to_long("127.0.0.1")
    2130706433L
    >>> ipv4_to_long("1.2.3.4")
    16909060L
    """
    ip_bytes_n = socket.inet_pton(socket.AF_INET, ip)
    return int(struct.unpack('!I', ip_bytes_n)[0])

def long_to_ipv4(ip_long):
    """Convert a long representation to a human readable IPv4 string

    >>> long_to_ipv4(2130706433L)
    '127.0.0.1'
    >>> long_to_ipv4(16909060L)
    '1.2.3.4'
    """
    ip_bytes = struct.pack("!I", ip_long)
    return socket.inet_ntop(socket.AF_INET, ip_bytes)

def long_to_ipv6(ip_long):
    """Convert a long representation to a human readable IPv6 string

    >>> long_to_ipv6(42540766411282592856903984951653826561L)
    '2001:db8::1'
    >>> long_to_ipv6(1L)
    '::1'
    """
    ip_parts = (int(ip_long/2**64), int(ip_long % 2**64))
    ip_bytes = struct.pack("!QQ", *ip_parts)
    return socket.inet_ntop(socket.AF_INET6, ip_bytes)

def ipv4_mask_to_long(mask):
    """Convert an IPv4 netmax (prefixlen) to long

    >>> ipv4_mask_to_long(24)
    4294967040L
    >>> ipv4_mask_to_long(8)
    4278190080L
    """
    mask_binary = ("1" * mask) + ("0" * (32-mask))
    return int(mask_binary, 2)

def ipv6_mask_to_long(mask):
    """Convert an IPv6 prefixlen to long

    >>> ipv6_mask_to_long(64)
    340282366920938463444927863358058659840L
    >>> ipv6_mask_to_long(128)
    340282366920938463463374607431768211455L
    """
    mask_binary = ("1" * mask) + ("0" * (128-mask))
    return int(mask_binary, 2)

def cidr_to_network(ip, prefixlen):
    """Calculate the network address for a CIDR notation network

    >>> cidr_to_network("127.1.1.1", 8)
    '127.0.0.0'
    >>> cidr_to_network("172.16.32.6", 24)
    '172.16.32.0'
    """
    if ':' in ip: # IPv6
        netmask = ipv6_mask_to_long(prefixlen)
        ip_long = ipv6_to_long(ip)
        network_long = ip_long & netmask
        network = long_to_ipv6(network_long)
    else:
        netmask = ipv4_mask_to_long(prefixlen)
        ip_long = ipv4_to_long(ip)
        network_long = ip_long & netmask
        network = long_to_ipv4(network_long)
    return network

class IpRangeList:
    """IP range list that supports CIDR notation for IPv4 and IPv6

    >>> r = IpRangeList("10.0.0.1", "172.16.4.254/24", "2001:db8::1/64")
    >>> "10.0.0.1" in r
    True
    >>> "10.0.0.2" in r
    False
    >>> "172.16.4.0" in r
    True
    >>> "172.16.4.255" in r
    True
    >>> "172.16.5.0" in r
    False
    >>> "2001:db8::ffff:ffff:ffff:ffff" in r
    True
    >>> "2001:db8:1::1" in r
    False
    >>> "2001:db8::" in r
    True
    """

    def __init__(self, *args):
        self._masks = []
        for cidr in args:
            if not '/' in cidr:
                ip = cidr
                if ':' in ip:
                    prefixlen = 128
                else:
                    prefixlen = 32
            else:
                (ip, prefixlen) = cidr.split("/")
            prefixlen = int(prefixlen)
            network = cidr_to_network(ip, prefixlen)
            self._masks.append((network, prefixlen))

    def __contains__(self, x):
        for (network, prefixlen) in self._masks:
            x_network = None
            if ':' in x and ':' in network:
                x_network = cidr_to_network(x, prefixlen)
            elif not ':' in network and not ':' in x:
                x_network = cidr_to_network(x, prefixlen)
            if x_network == network:
                return True
        return False

def load_tests(loader, tests, pattern):
    import doctest
    suite = doctest.DocTestSuite()
    tests.addTests(suite)
    return suite

def main():
    import unittest
    import doctest
    suite = doctest.DocTestSuite()
    unittest.TextTestRunner().run(suite)

if __name__ == '__main__':
    main()
<commit_msg>Use // instead of int<commit_after>#!/usr/bin/env python
# encoding: UTF-8
# Copyright Kapsi Internet-käyttäjät ry
# Idea from http://code.google.com/p/python-iptools ; code completely new

import socket
import struct

def ipv6_to_long(ip):
    """Return the IPv6 address string as a long

    >>> ipv6_to_long("2001:db8::1")
    42540766411282592856903984951653826561L
    >>> ipv6_to_long("::1")
    1L
    """
    ip_bytes_n = socket.inet_pton(socket.AF_INET6, ip)
    ip_parts = struct.unpack('!QQ', ip_bytes_n)
    return 2**64 * ip_parts[0] + ip_parts[1]

def ipv4_to_long(ip):
    """Return the IPv4 address string as a long

    >>> ipv4_to_long("127.0.0.1")
    2130706433L
    >>> ipv4_to_long("1.2.3.4")
    16909060L
    """
    ip_bytes_n = socket.inet_pton(socket.AF_INET, ip)
    return int(struct.unpack('!I', ip_bytes_n)[0])

def long_to_ipv4(ip_long):
    """Convert a long representation to a human readable IPv4 string

    >>> long_to_ipv4(2130706433L)
    '127.0.0.1'
    >>> long_to_ipv4(16909060L)
    '1.2.3.4'
    """
    ip_bytes = struct.pack("!I", ip_long)
    return socket.inet_ntop(socket.AF_INET, ip_bytes)

def long_to_ipv6(ip_long):
    """Convert a long representation to a human readable IPv6 string

    >>> long_to_ipv6(42540766411282592856903984951653826561)
    '2001:db8::1'
    >>> long_to_ipv6(1)
    '::1'
    """
    ip_parts = (ip_long//2**64, ip_long % 2**64)
    ip_bytes = struct.pack("!QQ", *ip_parts)
    return socket.inet_ntop(socket.AF_INET6, ip_bytes)

def ipv4_mask_to_long(mask):
    """Convert an IPv4 netmax (prefixlen) to long

    >>> ipv4_mask_to_long(24)
    4294967040L
    >>> ipv4_mask_to_long(8)
    4278190080L
    """
    mask_binary = ("1" * mask) + ("0" * (32-mask))
    return int(mask_binary, 2)

def ipv6_mask_to_long(mask):
    """Convert an IPv6 prefixlen to long

    >>> ipv6_mask_to_long(64)
    340282366920938463444927863358058659840L
    >>> ipv6_mask_to_long(128)
    340282366920938463463374607431768211455L
    """
    mask_binary = ("1" * mask) + ("0" * (128-mask))
    return int(mask_binary, 2)

def cidr_to_network(ip, prefixlen):
    """Calculate the network address for a CIDR notation network

    >>> cidr_to_network("127.1.1.1", 8)
    '127.0.0.0'
    >>> cidr_to_network("172.16.32.6", 24)
    '172.16.32.0'
    """
    if ':' in ip: # IPv6
        netmask = ipv6_mask_to_long(prefixlen)
        ip_long = ipv6_to_long(ip)
        network_long = ip_long & netmask
        network = long_to_ipv6(network_long)
    else:
        netmask = ipv4_mask_to_long(prefixlen)
        ip_long = ipv4_to_long(ip)
        network_long = ip_long & netmask
        network = long_to_ipv4(network_long)
    return network

class IpRangeList:
    """IP range list that supports CIDR notation for IPv4 and IPv6

    >>> r = IpRangeList("10.0.0.1", "172.16.4.254/24", "2001:db8::1/64")
    >>> "10.0.0.1" in r
    True
    >>> "10.0.0.2" in r
    False
    >>> "172.16.4.0" in r
    True
    >>> "172.16.4.255" in r
    True
    >>> "172.16.5.0" in r
    False
    >>> "2001:db8::ffff:ffff:ffff:ffff" in r
    True
    >>> "2001:db8:1::1" in r
    False
    >>> "2001:db8::" in r
    True
    """

    def __init__(self, *args):
        self._masks = []
        for cidr in args:
            if not '/' in cidr:
                ip = cidr
                if ':' in ip:
                    prefixlen = 128
                else:
                    prefixlen = 32
            else:
                (ip, prefixlen) = cidr.split("/")
            prefixlen = int(prefixlen)
            network = cidr_to_network(ip, prefixlen)
            self._masks.append((network, prefixlen))

    def __contains__(self, x):
        for (network, prefixlen) in self._masks:
            x_network = None
            if ':' in x and ':' in network:
                x_network = cidr_to_network(x, prefixlen)
            elif not ':' in network and not ':' in x:
                x_network = cidr_to_network(x, prefixlen)
            if x_network == network:
                return True
        return False

def load_tests(loader, tests, pattern):
    import doctest
    suite = doctest.DocTestSuite()
    tests.addTests(suite)
    return suite

def main():
    import unittest
    import doctest
    suite = doctest.DocTestSuite()
    unittest.TextTestRunner().run(suite)

if __name__ == '__main__':
    main()
<|endoftext|>


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

# Copyright(C) 2010-2012 Romain Bignon, Florent Fourcot
#
# This file is part of weboob.
#
# weboob is free software: you can redistribute it and/or modify
# it under the terms of the GNU Affero General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
# weboob is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
# GNU Affero General Public License for more details.
#
# You should have received a copy of the GNU Affero General Public License
# along with weboob. If not, see <http://www.gnu.org/licenses/>.

from __future__ import with_statement

from .browser import SachsenBrowser
from weboob.capabilities.gauge import ICapGauge, GaugeSensor, Gauge, SensorNotFound
from weboob.tools.backend import BaseBackend

__all__ = ['SachsenLevelBackend']

class SachsenLevelBackend(BaseBackend, ICapGauge):
    NAME = 'sachsen'
    MAINTAINER = u'Florent Fourcot'
    EMAIL = 'weboob@flo.fourcot.fr'
    VERSION = '0.e'
    LICENSE = 'AGPLv3+'
    DESCRIPTION = u"Level of Sachsen river"
    BROWSER = SachsenBrowser

    def iter_gauges(self, pattern=None):
        if pattern is None:
            for gauge in self.browser.get_rivers_list():
                yield gauge
        else:
            lowpattern = pattern.lower()
            for gauge in self.get_rivers_list():
                if lowpattern in gauge.name.lower() or lowpattern in gauge.river.lower():
                    yield gauge

    def _get_gauge_by_id(self, id):
        for gauge in self.browser.get_rivers_list():
            if id == gauge.id:
                return gauge
        return None

    def _get_sensor_by_id(self, id):
        for gauge in self.browser.get_rivers_list():
            for sensor in gauge.sensors:
                if id == sensor.id:
                    return sensor
        return None

    def iter_sensors(self, gauge, pattern=None):
        if not isinstance(gauge, Gauge):
            gauge = self._get_gauge_by_id(gauge)
        if pattern is None:
            for sensor in gauge.sensors:
                yield sensor
        else:
            lowpattern = pattern.lower()
            for sensor in gauge.sensors:
                if lowpattern in sensor.name.lower():
                    yield sensor

    def iter_gauge_history(self, sensor):
        if not isinstance(sensor, GaugeSensor):
            sensor = self._get_sensor_by_id(sensor)
        if sensor is None:
            raise SensorNotFound()
        return self.browser.iter_history(sensor)

    def get_last_measure(self, sensor):
        if not isinstance(sensor, GaugeSensor):
            sensor = self._get_sensor_by_id(sensor)
        if sensor is None:
            raise SensorNotFound()
        return sensor.lastvalue
<commit_msg>Fix typo error<commit_after># -*- coding: utf-8 -*-

# Copyright(C) 2010-2012 Romain Bignon, Florent Fourcot
#
# This file is part of weboob.
#
# weboob is free software: you can redistribute it and/or modify
# it under the terms of the GNU Affero General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
# weboob is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
# GNU Affero General Public License for more details.
#
# You should have received a copy of the GNU Affero General Public License
# along with weboob. If not, see <http://www.gnu.org/licenses/>.

from __future__ import with_statement

from .browser import SachsenBrowser
from weboob.capabilities.gauge import ICapGauge, GaugeSensor, Gauge, SensorNotFound
from weboob.tools.backend import BaseBackend

__all__ = ['SachsenLevelBackend']

class SachsenLevelBackend(BaseBackend, ICapGauge):
    NAME = 'sachsen'
    MAINTAINER = u'Florent Fourcot'
    EMAIL = 'weboob@flo.fourcot.fr'
    VERSION = '0.e'
    LICENSE = 'AGPLv3+'
    DESCRIPTION = u"Level of Sachsen river"
    BROWSER = SachsenBrowser

    def iter_gauges(self, pattern=None):
        if pattern is None:
            for gauge in self.browser.get_rivers_list():
                yield gauge
        else:
            lowpattern = pattern.lower()
            for gauge in self.browser.get_rivers_list():
                if lowpattern in gauge.name.lower() or lowpattern in gauge.object.lower():
                    yield gauge

    def _get_gauge_by_id(self, id):
        for gauge in self.browser.get_rivers_list():
            if id == gauge.id:
                return gauge
        return None

    def _get_sensor_by_id(self, id):
        for gauge in self.browser.get_rivers_list():
            for sensor in gauge.sensors:
                if id == sensor.id:
                    return sensor
        return None

    def iter_sensors(self, gauge, pattern=None):
        if not isinstance(gauge, Gauge):
            gauge = self._get_gauge_by_id(gauge)
        if pattern is None:
            for sensor in gauge.sensors:
                yield sensor
        else:
            lowpattern = pattern.lower()
            for sensor in gauge.sensors:
                if lowpattern in sensor.name.lower():
                    yield sensor

    def iter_gauge_history(self, sensor):
        if not isinstance(sensor, GaugeSensor):
            sensor = self._get_sensor_by_id(sensor)
        if sensor is None:
            raise SensorNotFound()
        return self.browser.iter_history(sensor)

    def get_last_measure(self, sensor):
        if not isinstance(sensor, GaugeSensor):
            sensor = self._get_sensor_by_id(sensor)
        if sensor is None:
            raise SensorNotFound()
        return sensor.lastvalue
<|endoftext|>


<commit_before>import re

from ._errors import STAFError

class STAFUnmarshalError(STAFError):
    pass

unmarshal_recursive     = 'unmarshal-recursive'
unmarshal_non_recursive = 'unmarshal-non-recursive'
unmarshal_none          = 'unmarshal-none'

marker = '@SDT/'

def unmarshal(data, mode=unmarshal_recursive):
    if mode == unmarshal_none or not data.startswith(marker):
        return data

    (obj, remainder) = unmarshal_internal(data, mode)

    if remainder != '':
        raise STAFUnmarshalError('unexpected trailing data')

    return obj

# Special classes for dealing with STAF Map Class types. Unmarshalled map
# classes will instances of MapClass, which is derived from dict and can be used
# as a dict. It also provides display_name() and display_short_name() to access
# those attributes if desired.

class MapClassDefinition(object):
    def __init__(self, name):
        self.name = name
        self.keys = []
        self._names = {} # {'key' : ('display name', 'short display name')}

    def add_item(self, key, display_name, display_short_name=None):
        self.keys.append(key)
        self._names[key] = (display_name, display_short_name)

    def display_name(self, key):
        return self._names[key][0]

    def display_short_name(self, key):
        return self._names[key][1]

class MapClass(dict):
    def __init__(self, definition, *args, **kwargs):
        super(MapClass, self).__init__(*args, **kwargs)

        self.definition = definition

    def display_name(self, key):
        return self.definition.display_name(key)

    def display_short_name(self, key):
        return self.definition.display_short_name(key)

def unmarshal_internal(data, mode, context=None):
    if context is None:
        context = {}

    if not data.startswith(marker):
        raise STAFUnmarshalError('missing marshalled data marker')

    sym_index = len(marker)
    try:
        symbol = data[sym_index]
        rest = data[sym_index+1:]
    except IndexError:
        raise STAFUnmarshalError('incomplete marshalled data')

    unmarshaller = get_unmarshaller(symbol)
    if unmarshaller is None:
        raise STAFUnmarshalError('unrecognized data type indicator')

    return unmarshaller.unmarshal(rest, mode, context)

def get_unmarshaller(symbol):
    if symbol == '$':
        return ScalarUnmarshaller
    elif symbol == '{':
        return MapUnmarshaller
    elif symbol == '[':
        return ListUnmarshaller
    elif symbol == '%':
        return MapClassUnmarshaller
    elif symbol == '*':
        return ContextUnmarshaller
    else:
        return None

class Unmarshaller(object):
    clc_matcher = re.compile(
        r'''^
        :(?P<len>\d+):  # Colon-length-colon
        (?P<rest>.*)    # Data
        $''', re.VERBOSE | re.DOTALL)

    @classmethod
    def unmarshal(cls, data, mode, context):
        raise NotImplementedError()

    @classmethod
    def read_clc_obj(cls, data):
        m = cls.clc_matcher.match(data)
        if m is None:
            raise STAFUnmarshalError('bad format for colon-length-colon object')

        (length, rest) = m.group('len', 'rest')
        length = int(length)

        if length > len(rest):
            raise STAFUnmarshalError('specified length exceeds available data')

        obj = rest[:length]
        rest = rest[length:]

        return (obj, rest)

class ScalarUnmarshaller(Unmarshaller):
    @classmethod
    def unmarshal(cls, data, mode, context):
        if not data.startswith('0') and not data.startswith('S'):
            raise STAFUnmarshalError('bad format for scalar object')

        typ = data[0]
        (obj, rest) = cls.read_clc_obj(data[1:])

        if typ == '0':
            if obj != '':
                raise STAFUnmarshalError('bad format for none object')
            return (None, rest)

        else: # typ == 'S'
            # Possibly do recursive unmarshalling.
            if mode == unmarshal_recursive:
                obj = unmarshal(obj, mode)

            return (obj, rest)

class MapUnmarshaller(Unmarshaller):
    @classmethod
    def unmarshal(cls, data, mode, context):
        (items, remainder) = cls.read_clc_obj(data)

        result = {}

        while items:
            (key, items) = cls.read_clc_obj(items)
            (val, items) = unmarshal_internal(items, mode, context)

            result[key] = val

        return (result, remainder)

class ListUnmarshaller(Unmarshaller):
    count_matcher = re.compile(r'''^
        (?P<count>\d+)  # Number of elements
        (?P<rest>.*)    # List items
        $''', re.VERBOSE | re.DOTALL)

    @classmethod
    def unmarshal(cls, data, mode, context):
        m = cls.count_matcher.match(data)
        if m is None:
            raise STAFUnmarshalError('bad format for list object')

        (count, rest) = m.group('count', 'rest')
        count = int(count)

        (items, remainder) = cls.read_clc_obj(rest)

        result = []
        for i in range(count):
            (obj, items) = unmarshal_internal(items, mode, context)
            result.append(obj)

        if items != '':
            raise STAFUnmarshalError('unexpected trailing data')

        return (result, remainder)

class MapClassUnmarshaller(Unmarshaller):
    @classmethod
    def unmarshal(cls, data, mode, context):
        (content, remainder) = cls.read_clc_obj(data)
        (class_name, values) = cls.read_clc_obj(content)

        class_def = context.get(class_name)
        if class_def is None:
            raise STAFUnmarshalError('missing map class definition for %r' %
                                     class_name)

        result = MapClass(class_def)
        # The ordering and names of the keys are given by class_def.keys.
        for key in class_def.keys:
            (value, values) = unmarshal_internal(values, mode, context)
            result[key] = value

        if values != '':
            raise STAFUnmarshalError('unexpected trailing data')

        return (result, remainder)

class ContextUnmarshaller(Unmarshaller):
    @classmethod
    def unmarshal(cls, data, mode, context):
        (content, remainder) = cls.read_clc_obj(data)
        (context_map, root_data) = unmarshal_internal(content, mode, context)

        class_map = context_map.get('map-class-map', {})

        # Build a map of names to MapClassDefinitions.
        new_context = {}
        for (name, info) in class_map.iteritems():
            class_def = MapClassDefinition(name)
            for item in info['keys']:
                class_def.add_item(item['key'], item['display-name'],
                                   item.get('display-short-name'))

            new_context[name] = class_def

        # Note that we may be forsaking an existing class map in 'context' in
        # favor of new_context. Nested contexts probably shouldn't happen, but
        # if they do this means the objects in the inner context won't be able
        # to reference objects in the outer context. This is probably fine.
        (root_obj, trailing) = unmarshal_internal(root_data, mode, new_context)

        if trailing != '':
            raise STAFUnmarshalError('unexpected trailing data')

        return (root_obj, remainder)
<commit_msg>Make MapClass present items in the original order.<commit_after>import re

from ._errors import STAFError

class STAFUnmarshalError(STAFError):
    pass

unmarshal_recursive     = 'unmarshal-recursive'
unmarshal_non_recursive = 'unmarshal-non-recursive'
unmarshal_none          = 'unmarshal-none'

marker = '@SDT/'

def unmarshal(data, mode=unmarshal_recursive):
    if mode == unmarshal_none or not data.startswith(marker):
        return data

    (obj, remainder) = unmarshal_internal(data, mode)

    if remainder != '':
        raise STAFUnmarshalError('unexpected trailing data')

    return obj

# Special classes for dealing with STAF Map Class types. Unmarshalled map
# classes will instances of MapClass, which is derived from dict and can be used
# as a dict. It also provides display_name() and display_short_name() to access
# those attributes if desired.

class MapClassDefinition(object):
    def __init__(self, name):
        self.name = name
        self.keys = []
        self._names = {} # {'key' : ('display name', 'short display name')}

    def add_item(self, key, display_name, display_short_name=None):
        self.keys.append(key)
        self._names[key] = (display_name, display_short_name)

    def display_name(self, key):
        return self._names[key][0]

    def display_short_name(self, key):
        return self._names[key][1]

class MapClass(dict):
    def __init__(self, definition, *args, **kwargs):
        super(MapClass, self).__init__(*args, **kwargs)

        self.definition = definition

    def display_name(self, key):
        return self.definition.display_name(key)

    def display_short_name(self, key):
        return self.definition.display_short_name(key)

    # Overrides to preserve ordering on item access.

    def __iter__(self):
        return iter(self.definition.keys)

    def iterkeys(self):
        return iter(self.definition.keys)

    def keys(self):
        return list(self.iterkeys())

    def itervalues(self):
        for key in self.definition.keys:
            yield self[key]

    def values(self):
        return list(self.itervalues())

    def iteritems(self):
        for key in self.definition.keys:
            yield (key, self[key])

    def items(self):
        return list(self.iteritems())

    def __repr__(self):
        items = []
        for (k, v) in self.iteritems():
            items.append('%r: %r' % (k, v))

        return '{%s}' % ', '.join(items)

    # NOTE: Python 2.7 adds viewitems(), viewkeys(), and viewvalues(). These
    # aren't supported here because it would be tricky to make them work and
    # they aren't really useful in MapClasses that aren't expected to have items
    # added/removed. Because the inherited versions from dict won't work as
    # expected, disable them.

    def viewitems(self):
        raise NotImplementedError('viewitems is not supported in MapClass')

    def viewkeys(self):
        raise NotImplementedError('viewkeys is not supported in MapClass')

    def viewvalues(self):
        raise NotImplementedError('viewvalues is not supported in MapClass')

def unmarshal_internal(data, mode, context=None):
    if context is None:
        context = {}

    if not data.startswith(marker):
        raise STAFUnmarshalError('missing marshalled data marker')

    sym_index = len(marker)
    try:
        symbol = data[sym_index]
        rest = data[sym_index+1:]
    except IndexError:
        raise STAFUnmarshalError('incomplete marshalled data')

    unmarshaller = get_unmarshaller(symbol)
    if unmarshaller is None:
        raise STAFUnmarshalError('unrecognized data type indicator')

    return unmarshaller.unmarshal(rest, mode, context)

def get_unmarshaller(symbol):
    if symbol == '$':
        return ScalarUnmarshaller
    elif symbol == '{':
        return MapUnmarshaller
    elif symbol == '[':
        return ListUnmarshaller
    elif symbol == '%':
        return MapClassUnmarshaller
    elif symbol == '*':
        return ContextUnmarshaller
    else:
        return None

class Unmarshaller(object):
    clc_matcher = re.compile(
        r'''^
        :(?P<len>\d+):  # Colon-length-colon
        (?P<rest>.*)    # Data
        $''', re.VERBOSE | re.DOTALL)

    @classmethod
    def unmarshal(cls, data, mode, context):
        raise NotImplementedError()

    @classmethod
    def read_clc_obj(cls, data):
        m = cls.clc_matcher.match(data)
        if m is None:
            raise STAFUnmarshalError('bad format for colon-length-colon object')

        (length, rest) = m.group('len', 'rest')
        length = int(length)

        if length > len(rest):
            raise STAFUnmarshalError('specified length exceeds available data')

        obj = rest[:length]
        rest = rest[length:]

        return (obj, rest)

class ScalarUnmarshaller(Unmarshaller):
    @classmethod
    def unmarshal(cls, data, mode, context):
        if not data.startswith('0') and not data.startswith('S'):
            raise STAFUnmarshalError('bad format for scalar object')

        typ = data[0]
        (obj, rest) = cls.read_clc_obj(data[1:])

        if typ == '0':
            if obj != '':
                raise STAFUnmarshalError('bad format for none object')
            return (None, rest)

        else: # typ == 'S'
            # Possibly do recursive unmarshalling.
            if mode == unmarshal_recursive:
                obj = unmarshal(obj, mode)

            return (obj, rest)

class MapUnmarshaller(Unmarshaller):
    @classmethod
    def unmarshal(cls, data, mode, context):
        (items, remainder) = cls.read_clc_obj(data)

        result = {}

        while items:
            (key, items) = cls.read_clc_obj(items)
            (val, items) = unmarshal_internal(items, mode, context)

            result[key] = val

        return (result, remainder)

class ListUnmarshaller(Unmarshaller):
    count_matcher = re.compile(r'''^
        (?P<count>\d+)  # Number of elements
        (?P<rest>.*)    # List items
        $''', re.VERBOSE | re.DOTALL)

    @classmethod
    def unmarshal(cls, data, mode, context):
        m = cls.count_matcher.match(data)
        if m is None:
            raise STAFUnmarshalError('bad format for list object')

        (count, rest) = m.group('count', 'rest')
        count = int(count)

        (items, remainder) = cls.read_clc_obj(rest)

        result = []
        for i in range(count):
            (obj, items) = unmarshal_internal(items, mode, context)
            result.append(obj)

        if items != '':
            raise STAFUnmarshalError('unexpected trailing data')

        return (result, remainder)

class MapClassUnmarshaller(Unmarshaller):
    @classmethod
    def unmarshal(cls, data, mode, context):
        (content, remainder) = cls.read_clc_obj(data)
        (class_name, values) = cls.read_clc_obj(content)

        class_def = context.get(class_name)
        if class_def is None:
            raise STAFUnmarshalError('missing map class definition for %r' %
                                     class_name)

        result = MapClass(class_def)
        # The ordering and names of the keys are given by class_def.keys.
        for key in class_def.keys:
            (value, values) = unmarshal_internal(values, mode, context)
            result[key] = value

        if values != '':
            raise STAFUnmarshalError('unexpected trailing data')

        return (result, remainder)

class ContextUnmarshaller(Unmarshaller):
    @classmethod
    def unmarshal(cls, data, mode, context):
        (content, remainder) = cls.read_clc_obj(data)
        (context_map, root_data) = unmarshal_internal(content, mode, context)

        class_map = context_map.get('map-class-map', {})

        # Build a map of names to MapClassDefinitions.
        new_context = {}
        for (name, info) in class_map.iteritems():
            class_def = MapClassDefinition(name)
            for item in info['keys']:
                class_def.add_item(item['key'], item['display-name'],
                                   item.get('display-short-name'))

            new_context[name] = class_def

        # Note that we may be forsaking an existing class map in 'context' in
        # favor of new_context. Nested contexts probably shouldn't happen, but
        # if they do this means the objects in the inner context won't be able
        # to reference objects in the outer context. This is probably fine.
        (root_obj, trailing) = unmarshal_internal(root_data, mode, new_context)

        if trailing != '':
            raise STAFUnmarshalError('unexpected trailing data')

        return (root_obj, remainder)
<|endoftext|>


<commit_before># Copyright (c) 2017 The Regents of the University of Michigan.
# All Rights Reserved. Licensed according to the terms of the Revised
# BSD License. See LICENSE.txt for details.
from hamcrest import *
import unittest

from ..generate_pageview import Pagetags

class GivenEmptyPagetags (unittest.TestCase):

    def setUp (self):
        self.tags = Pagetags()

    def test_can_init_pagetags (self):
        assert_that(self.tags.generate_pageview(), is_(equal_to("")))

    def test_default_confidence_is_100 (self):
        assert_that(self.tags.default_confidence, is_(equal_to(100)))

    def test_can_alter_default_confidence (self):
        for confid in (900, 100, 444):
            self.tags.default_confidence = confid
            assert_that(self.tags.default_confidence,
                        is_(equal_to(confid)))

    def test_cannot_set_confidence_to_weird_values (self):
        for confid in (0, 99, 901, 1000, 500.5):
            assert_that(calling(setattr).with_args(self.tags,
                                                   "default_confidence",
                                                   confid),
                        raises(ValueError))

    def test_can_add_tags (self):
        self.tags.add_raw_tags({"tags": []})
<commit_msg>:white_check_mark: Confidence must be an int<commit_after># Copyright (c) 2017 The Regents of the University of Michigan.
# All Rights Reserved. Licensed according to the terms of the Revised
# BSD License. See LICENSE.txt for details.
from hamcrest import *
import unittest

from ..generate_pageview import Pagetags

class GivenEmptyPagetags (unittest.TestCase):

    def setUp (self):
        self.tags = Pagetags()

    def test_can_init_pagetags (self):
        assert_that(self.tags.generate_pageview(), is_(equal_to("")))

    def test_default_confidence_is_100 (self):
        assert_that(self.tags.default_confidence, is_(equal_to(100)))

    def test_can_alter_default_confidence (self):
        for confid in (900, 100, 444):
            self.tags.default_confidence = confid
            assert_that(self.tags.default_confidence,
                        is_(equal_to(confid)))

    def test_cannot_set_confidence_to_weird_values (self):
        for confid in (0, 99, 901, 1000, 500.5, "500"):
            assert_that(calling(setattr).with_args(self.tags,
                                                   "default_confidence",
                                                   confid),
                        raises(ValueError))

    def test_can_add_tags (self):
        self.tags.add_raw_tags({"tags": []})
<|endoftext|>


<commit_before>from __future__ import unicode_literals
import unittest

from embedly.client import Embedly
from embedly.models import Url

class EmbedlyTestCase(unittest.TestCase):

    def __init__(self, *args, **kwargs):
        self.key = 'internal'

        if not self.key:
            raise ValueError('Set envirnomental varible EMBEDLY_API_KEY ' +
                             'before running these tests like so: $ export ' +
                             'EMBEDLY_API_KEY=key')

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

    def test_model(self):
        data = {
            'provider_url': 'http://www.google.com/',
            'safe': True,
            'description': 'Google',
            'original_url': 'http://google.com/',
            'url': 'http://www.google.com/',
            'type': 'html',
            'object': {},
            'provider_display': 'www.google.com',
            'author_name': None,
            'favicon_url': 'http://www.google.com/favicon.ico',
            'place': {},
            'author_url': None,
            'images': [
                {'url': 'http://www.google.com/intl/en_ALL/images/srpr/logo1w.png',
                 'width': 275,
                 'height': 95}],
            'title': 'Google',
            'provider_name': 'Google',
            'cache_age': 86400,
            'embeds': []
        }

        obj = Url(data, 'preview', 'http://google.com/')

        self.assertTrue(len(obj) is 17)
        self.assertTrue(len(obj.values()) is 17)
        self.assertTrue(len(obj.keys()) is 17)
        self.assertTrue(len(obj.items()) is 17)

        self.assertTrue('type' in obj.keys())
        self.assertTrue('html' in obj.values())

        #Get the object
        self.assertTrue(obj.type == 'html')
        self.assertTrue(obj['type'] == 'html')
        self.assertTrue(obj.get('type') == 'html')

        #nope
        self.assertTrue(obj.nothing is None)

        obj.nothing = 'something'
        self.assertTrue(obj.nothing == 'something')

        obj['nothing'] = 'maybe'
        self.assertTrue(obj['nothing'] == 'maybe')

        del obj['nothing']
        self.assertTrue(obj.nothing is None)

        #Deep Get attrs
        self.assertTrue(obj.images[0].width is 275)
        self.assertTrue(obj.images[0].nothing is None)
        self.assertTrue(obj.object.type is None)

    def test_provider(self):
        http = Embedly(self.key)

        obj = http.oembed('http://www.scribd.com/doc/13994900/Easter')
        self.assertTrue(obj.provider_url == 'http://www.scribd.com/')

        obj = http.oembed('http://www.scribd.com/doc/28452730/Easter-Cards')
        self.assertTrue(obj.provider_url == 'http://www.scribd.com/')

        obj = http.oembed('http://www.youtube.com/watch?v=Zk7dDekYej0')
        self.assertTrue(obj.provider_url == 'http://www.youtube.com/')

        obj = http.oembed('http://yfrog.com/h22eu4j')
        self.assertTrue(obj.provider_url == 'http://yfrog.com')

    def test_providers(self):
        http = Embedly(self.key)

        objs = list(http.oembed(['http://www.scribd.com/doc/13994900/Easter',
                                 'http://www.scribd.com/doc/28452730/Easter-Cards']))
        self.assertTrue(objs[0].provider_url == 'http://www.scribd.com/')
        self.assertTrue(objs[1].provider_url == 'http://www.scribd.com/')

        objs = list(http.oembed(['http://www.youtube.com/watch?v=Zk7dDekYej0',
                                 'http://yfrog.com/h22eu4']))
        self.assertTrue(objs[0].provider_url == 'http://www.youtube.com/')
        self.assertTrue(objs[1].provider_url == 'http://yfrog.com')

    def test_error(self):
        http = Embedly(self.key)

        obj = http.oembed('http://www.embedly.com/this/is/a/bad/url')
        self.assertTrue(obj.error is True, obj.dict)
        obj = http.oembed('http://blog.embed.ly/lsbsdlfldsf/asdfkljlas/klajsdlfkasdf')
        self.assertTrue(obj.error is True, obj.dict)
        obj = http.oembed('http://twitpic/nothing/to/see/here')
        self.assertTrue(obj.error is True, obj.dict)

    def test_multi_errors(self):
        http = Embedly(self.key)

        objs = list(http.oembed(['http://www.embedly.com/this/is/a/bad/url',
                                 'http://blog.embed.ly/alsd/slsdlf/asdlfj']))
        self.assertTrue(objs[0].type == 'error', objs[0].dict)
        self.assertTrue(objs[1].type == 'error', objs[1].dict)

        objs = list(http.oembed(['http://blog.embed.ly/lsbsdlfldsf/asdf/kl',
                                 'http://twitpic.com/nothing/to/see/here']))
        self.assertTrue(objs[0].type == 'error',objs[0].dict)
        self.assertTrue(objs[1].type == 'error',objs[1].dict)

        objs = list(http.oembed(['http://blog.embed.ly/lsbsdlfldsf/asdf/kl',
                            'http://yfrog.com/h22eu4j']))
        self.assertTrue(objs[0].type == 'error',objs[0].dict)
        self.assertTrue(objs[1].type == 'photo',objs[1].dict)

        objs = list(http.oembed(['http://yfrog.com/h22eu4j',
                            'http://www.scribd.com/asdf/asdf/asdfasdf']))
        self.assertTrue(objs[0].type == 'photo',objs[0].dict)
        self.assertTrue(objs[1].type == 'error',objs[1].dict)

    def test_too_many_urls(self):
        http = Embedly(self.key)

        urls = ['http://embed.ly'] * 21
        try:
            http.oembed(urls)
            self.fail('too many urls, should have thrown an error')
        except Exception as e:
            self.assertTrue(type(e), ValueError)

if __name__ == '__main__':
    unittest.main()<commit_msg>More pep8 fixes<commit_after>from __future__ import unicode_literals
import unittest

from embedly.client import Embedly
from embedly.models import Url

class EmbedlyTestCase(unittest.TestCase):

    def __init__(self, *args, **kwargs):
        self.key = 'internal'

        if not self.key:
            raise ValueError('Set envirnomental varible EMBEDLY_API_KEY ' +
                             'before running these tests like so: $ export ' +
                             'EMBEDLY_API_KEY=key')

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

    def test_model(self):
        data = {
            'provider_url': 'http://www.google.com/',
            'safe': True,
            'description': 'Google',
            'original_url': 'http://google.com/',
            'url': 'http://www.google.com/',
            'type': 'html',
            'object': {},
            'provider_display': 'www.google.com',
            'author_name': None,
            'favicon_url': 'http://www.google.com/favicon.ico',
            'place': {},
            'author_url': None,
            'images': [
                {'url': 'http://www.google.com/intl/en_ALL/images/srpr/logo1w.png',
                 'width': 275,
                 'height': 95}],
            'title': 'Google',
            'provider_name': 'Google',
            'cache_age': 86400,
            'embeds': []
        }

        obj = Url(data, 'preview', 'http://google.com/')

        self.assertTrue(len(obj) is 17)
        self.assertTrue(len(obj.values()) is 17)
        self.assertTrue(len(obj.keys()) is 17)
        self.assertTrue(len(obj.items()) is 17)

        self.assertTrue('type' in obj.keys())
        self.assertTrue('html' in obj.values())

        #Get the object
        self.assertTrue(obj.type == 'html')
        self.assertTrue(obj['type'] == 'html')
        self.assertTrue(obj.get('type') == 'html')

        #nope
        self.assertTrue(obj.nothing is None)

        obj.nothing = 'something'
        self.assertTrue(obj.nothing == 'something')

        obj['nothing'] = 'maybe'
        self.assertTrue(obj['nothing'] == 'maybe')

        del obj['nothing']
        self.assertTrue(obj.nothing is None)

        #Deep Get attrs
        self.assertTrue(obj.images[0].width is 275)
        self.assertTrue(obj.images[0].nothing is None)
        self.assertTrue(obj.object.type is None)

    def test_provider(self):
        http = Embedly(self.key)

        obj = http.oembed('http://www.scribd.com/doc/13994900/Easter')
        self.assertTrue(obj.provider_url == 'http://www.scribd.com/')

        obj = http.oembed('http://www.scribd.com/doc/28452730/Easter-Cards')
        self.assertTrue(obj.provider_url == 'http://www.scribd.com/')

        obj = http.oembed('http://www.youtube.com/watch?v=Zk7dDekYej0')
        self.assertTrue(obj.provider_url == 'http://www.youtube.com/')

        obj = http.oembed('http://yfrog.com/h22eu4j')
        self.assertTrue(obj.provider_url == 'http://yfrog.com')

    def test_providers(self):
        http = Embedly(self.key)

        objs = list(http.oembed(['http://www.scribd.com/doc/13994900/Easter',
                                 'http://www.scribd.com/doc/28452730/Easter-Cards']))
        self.assertTrue(objs[0].provider_url == 'http://www.scribd.com/')
        self.assertTrue(objs[1].provider_url == 'http://www.scribd.com/')

        objs = list(http.oembed(['http://www.youtube.com/watch?v=Zk7dDekYej0',
                                 'http://yfrog.com/h22eu4']))
        self.assertTrue(objs[0].provider_url == 'http://www.youtube.com/')
        self.assertTrue(objs[1].provider_url == 'http://yfrog.com')

    def test_error(self):
        http = Embedly(self.key)

        obj = http.oembed('http://www.embedly.com/this/is/a/bad/url')
        self.assertTrue(obj.error is True, obj.dict)
        obj = http.oembed('http://blog.embed.ly/lsbsdlfldsf/asdfkljlas/klajsdlfkasdf')
        self.assertTrue(obj.error is True, obj.dict)
        obj = http.oembed('http://twitpic/nothing/to/see/here')
        self.assertTrue(obj.error is True, obj.dict)

    def test_multi_errors(self):
        http = Embedly(self.key)

        objs = list(http.oembed(['http://www.embedly.com/this/is/a/bad/url',
                                 'http://blog.embed.ly/alsd/slsdlf/asdlfj']))
        self.assertTrue(objs[0].type == 'error', objs[0].dict)
        self.assertTrue(objs[1].type == 'error', objs[1].dict)

        objs = list(http.oembed(['http://blog.embed.ly/lsbsdlfldsf/asdf/kl',
                                 'http://twitpic.com/nothing/to/see/here']))
        self.assertTrue(objs[0].type == 'error',objs[0].dict)
        self.assertTrue(objs[1].type == 'error',objs[1].dict)

        objs = list(http.oembed(['http://blog.embed.ly/lsbsdlfldsf/asdf/kl',
                                 'http://yfrog.com/h22eu4j']))
        self.assertTrue(objs[0].type == 'error',objs[0].dict)
        self.assertTrue(objs[1].type == 'photo',objs[1].dict)

        objs = list(http.oembed(['http://yfrog.com/h22eu4j',
                                 'http://www.scribd.com/asdf/asdf/asdfasdf']))
        self.assertTrue(objs[0].type == 'photo',objs[0].dict)
        self.assertTrue(objs[1].type == 'error',objs[1].dict)

    def test_too_many_urls(self):
        http = Embedly(self.key)

        urls = ['http://embed.ly'] * 21
        try:
            http.oembed(urls)
            self.fail('too many urls, should have thrown an error')
        except Exception as e:
            self.assertTrue(type(e), ValueError)

if __name__ == '__main__':
    unittest.main()<|endoftext|>


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

from arch.api import session
from arch.api.utils.log_utils import LoggerFactory
from fate_flow.entity.metric import MetricType, MetricMeta, Metric
from federatedml.framework.homo.procedure import aggregator
from federatedml.model_base import ModelBase
from federatedml.nn.homo_nn import nn_model
from federatedml.nn.homo_nn.nn_model import restore_nn_model
from federatedml.optim.convergence import converge_func_factory
from federatedml.param.homo_nn_param import HomoNNParam
from federatedml.transfer_variable.transfer_class.homo_transfer_variable import HomoTransferVariable
from federatedml.util import consts

Logger = LoggerFactory.get_logger()
MODEL_META_NAME = "HomoNNModelMeta"
MODEL_PARAM_NAME = "HomoNNModelParam"

def _build_model_dict(meta, param):
    return {MODEL_META_NAME: meta, MODEL_PARAM_NAME: param}

def _extract_param(model_dict: dict):
    return model_dict.get(MODEL_PARAM_NAME, None)

def _extract_meta(model_dict: dict):
    return model_dict.get(MODEL_META_NAME, None)

class HomoNNBase(ModelBase):

    def __init__(self):
        super().__init__()
        self.model_param = HomoNNParam()
        self.role = None

    def _init_model(self, param):
        super()._init_model(param)

        self.transfer_variable = HomoTransferVariable()
        secure_aggregate = param.secure_aggregate
        self.aggregator = aggregator.with_role(role=self.role,
                                               transfer_variable=self.transfer_variable,
                                               enable_secure_aggregate=secure_aggregate)
        self.max_iter = param.max_iter
        self.aggregator_iter = 0

    def _iter_suffix(self):
        return self.aggregator_iter,

class HomoNNArbiter(HomoNNBase):

    def __init__(self):
        super().__init__()
        self.role = consts.ARBITER

    def _init_model(self, param):
        super(HomoNNArbiter, self)._init_model(param)
        early_stop = self.model_param.early_stop
        self.converge_func = converge_func_factory(early_stop.converge_func, early_stop.eps).is_converge
        self.loss_consumed = early_stop.converge_func != "weight_diff"

    def callback_loss(self, iter_num, loss):
        metric_meta = MetricMeta(name='train',
                                 metric_type=MetricType.LOSS,
                                 extra_metas={
                                     "unit_name": "iters",
                                 })

        self.callback_meta(metric_name='loss', metric_namespace='train', metric_meta=metric_meta)
        self.callback_metric(metric_name='loss',
                             metric_namespace='train',
                             metric_data=[Metric(iter_num, loss)])

    def _check_monitored_status(self):
        loss = self.aggregator.aggregate_loss(suffix=self._iter_suffix())
        Logger.info(f"loss at iter {self.aggregator_iter}: {loss}")
        self.callback_loss(self.aggregator_iter, loss)
        if self.loss_consumed:
            converge_args = (loss,) if self.loss_consumed else (self.aggregator.model,)
            return self.aggregator.send_converge_status(self.converge_func,
                                                        converge_args=converge_args,
                                                        suffix=self._iter_suffix())

    def fit(self, data_inst):
        while self.aggregator_iter < self.max_iter:
            self.aggregator.aggregate_and_broadcast(suffix=self._iter_suffix())

            if self._check_monitored_status():
                Logger.info(f"early stop at iter {self.aggregator_iter}")
                break
            self.aggregator_iter += 1
        else:
            Logger.warn(f"reach max iter: {self.aggregator_iter}, not converged")

    def save_model(self):
        return self.aggregator.model

class HomoNNClient(HomoNNBase):

    def __init__(self):
        super().__init__()

    def _init_model(self, param):
        super(HomoNNClient, self)._init_model(param)
        self.batch_size = param.batch_size
        self.aggregate_every_n_epoch = 1
        self.nn_define = param.nn_define
        self.config_type = param.config_type
        self.optimizer = param.optimizer
        self.loss = param.loss
        self.metrics = param.metrics
        self.data_converter = nn_model.get_data_converter(self.config_type)
        self.model_builder = nn_model.get_nn_builder(config_type=self.config_type)

    def _check_monitored_status(self, data, epoch_degree):
        metrics = self.nn_model.evaluate(data)
        Logger.info(f"metrics at iter {self.aggregator_iter}: {metrics}")
        loss = metrics["loss"]
        self.aggregator.send_loss(loss=loss,
                                  degree=epoch_degree,
                                  suffix=self._iter_suffix())
        return self.aggregator.get_converge_status(suffix=self._iter_suffix())

    def __build_nn_model(self, input_shape):
        self.nn_model = self.model_builder(input_shape=input_shape,
                                           nn_define=self.nn_define,
                                           optimizer=self.optimizer,
                                           loss=self.loss,
                                           metrics=self.metrics)

    def fit(self, data_inst, *args):

        data = self.data_converter.convert(data_inst, batch_size=self.batch_size)
        self.__build_nn_model(data.get_shape()[0])

        epoch_degree = float(len(data))

        while self.aggregator_iter < self.max_iter:
            Logger.info(f"start {self.aggregator_iter}_th aggregation")

            # train
            self.nn_model.train(data, aggregate_every_n_epoch=self.aggregate_every_n_epoch)

            # send model for aggregate, then set aggregated model to local
            modify_func: typing.Callable = functools.partial(self.aggregator.aggregate_then_get,
                                                             degree=epoch_degree * self.aggregate_every_n_epoch,
                                                             suffix=self._iter_suffix())
            self.nn_model.modify(modify_func)

            # calc loss and check convergence
            if self._check_monitored_status(data, epoch_degree):
                Logger.info(f"early stop at iter {self.aggregator_iter}")
                break

            Logger.info(f"role {self.role} finish {self.aggregator_iter}_th aggregation")
            self.aggregator_iter += 1
        else:
            Logger.warn(f"reach max iter: {self.aggregator_iter}, not converged")

    def export_model(self):
        return _build_model_dict(meta=self._get_meta(), param=self._get_param())

    def _get_meta(self):
        from federatedml.protobuf.generated import nn_model_meta_pb2
        meta_pb = nn_model_meta_pb2.NNModelMeta()
        meta_pb.params.CopyFrom(self.model_param.generate_pb())
        meta_pb.aggregate_iter = self.aggregator_iter
        return meta_pb

    def _get_param(self):
        from federatedml.protobuf.generated import nn_model_param_pb2
        param_pb = nn_model_param_pb2.NNModelParam()
        param_pb.saved_model_bytes = self.nn_model.export_model()
        return param_pb

    def predict(self, data_inst):
        data = self.data_converter.convert(data_inst, batch_size=self.batch_size)
        predict = self.nn_model.predict(data)
        num_output_units = predict.shape[1]

        if num_output_units == 1:
            kv = map(lambda x: (x[0], (0 if x[1][0] < 0.5 else 1, x[1][0])), zip(data.get_keys(), predict))
            pred_tbl = session.parallelize(kv, include_key=True)
            return data_inst.join(pred_tbl, lambda d, pred: [d.label, pred[0], pred[1], {"label": pred[0]}])
        else:
            kv = map(lambda x: (x[0], (x[1].argmax(), [float(e) for e in x[1]])), zip(data.get_keys(), predict))
            pred_tbl = session.parallelize(kv, include_key=True)
            return data_inst.join(pred_tbl,
                                  lambda d, pred: [d.label, pred[0],
                                                   pred[1][pred[0]] / sum(pred[1]),
                                                   {"raw_predict": pred[1]}])

    def _load_model(self, model_dict):
        model_dict = list(model_dict["model"].values())[0]
        model_obj = _extract_param(model_dict)
        meta_obj = _extract_meta(model_dict)
        self.model_param.restore_from_pb(meta_obj.params)
        self._init_model(self.model_param)
        self.aggregator_iter = meta_obj.aggregate_iter
        self.nn_model = restore_nn_model(self.config_type, model_obj.saved_model_bytes)

class HomoNNHost(HomoNNClient):

    def __init__(self):
        super().__init__()
        self.role = consts.HOST

class HomoNNGuest(HomoNNClient):

    def __init__(self):
        super().__init__()
        self.role = consts.GUEST
<commit_msg>fix ISSUE#672<commit_after>#
#  Copyright 2019 The FATE Authors. All Rights Reserved.
#
#  Licensed under the Apache License, Version 2.0 (the "License");
#  you may not use this file except in compliance with the License.
#  You may obtain a copy of the License at
#
#      http://www.apache.org/licenses/LICENSE-2.0
#
#  Unless required by applicable law or agreed to in writing, software
#  distributed under the License is distributed on an "AS IS" BASIS,
#  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
#  See the License for the specific language governing permissions and
#  limitations under the License.
#
import functools
import typing

from arch.api import session
from arch.api.utils.log_utils import LoggerFactory
from fate_flow.entity.metric import MetricType, MetricMeta, Metric
from federatedml.framework.homo.procedure import aggregator
from federatedml.model_base import ModelBase
from federatedml.nn.homo_nn import nn_model
from federatedml.nn.homo_nn.nn_model import restore_nn_model
from federatedml.optim.convergence import converge_func_factory
from federatedml.param.homo_nn_param import HomoNNParam
from federatedml.transfer_variable.transfer_class.homo_transfer_variable import HomoTransferVariable
from federatedml.util import consts

Logger = LoggerFactory.get_logger()
MODEL_META_NAME = "HomoNNModelMeta"
MODEL_PARAM_NAME = "HomoNNModelParam"

def _build_model_dict(meta, param):
    return {MODEL_META_NAME: meta, MODEL_PARAM_NAME: param}

def _extract_param(model_dict: dict):
    return model_dict.get(MODEL_PARAM_NAME, None)

def _extract_meta(model_dict: dict):
    return model_dict.get(MODEL_META_NAME, None)

class HomoNNBase(ModelBase):

    def __init__(self):
        super().__init__()
        self.model_param = HomoNNParam()
        self.role = None

    def _init_model(self, param):
        super()._init_model(param)

        self.transfer_variable = HomoTransferVariable()
        secure_aggregate = param.secure_aggregate
        self.aggregator = aggregator.with_role(role=self.role,
                                               transfer_variable=self.transfer_variable,
                                               enable_secure_aggregate=secure_aggregate)
        self.max_iter = param.max_iter
        self.aggregator_iter = 0

    def _iter_suffix(self):
        return self.aggregator_iter,

class HomoNNArbiter(HomoNNBase):

    def __init__(self):
        super().__init__()
        self.role = consts.ARBITER

    def _init_model(self, param):
        super(HomoNNArbiter, self)._init_model(param)
        early_stop = self.model_param.early_stop
        self.converge_func = converge_func_factory(early_stop.converge_func, early_stop.eps).is_converge
        self.loss_consumed = early_stop.converge_func != "weight_diff"

    def callback_loss(self, iter_num, loss):
        metric_meta = MetricMeta(name='train',
                                 metric_type=MetricType.LOSS,
                                 extra_metas={
                                     "unit_name": "iters",
                                 })

        self.callback_meta(metric_name='loss', metric_namespace='train', metric_meta=metric_meta)
        self.callback_metric(metric_name='loss',
                             metric_namespace='train',
                             metric_data=[Metric(iter_num, loss)])

    def _check_monitored_status(self):
        loss = self.aggregator.aggregate_loss(suffix=self._iter_suffix())
        Logger.info(f"loss at iter {self.aggregator_iter}: {loss}")
        self.callback_loss(self.aggregator_iter, loss)
        if self.loss_consumed:
            converge_args = (loss,) if self.loss_consumed else (self.aggregator.model,)
            return self.aggregator.send_converge_status(self.converge_func,
                                                        converge_args=converge_args,
                                                        suffix=self._iter_suffix())

    def fit(self, data_inst):
        while self.aggregator_iter < self.max_iter:
            self.aggregator.aggregate_and_broadcast(suffix=self._iter_suffix())

            if self._check_monitored_status():
                Logger.info(f"early stop at iter {self.aggregator_iter}")
                break
            self.aggregator_iter += 1
        else:
            Logger.warn(f"reach max iter: {self.aggregator_iter}, not converged")

    def save_model(self):
        return self.aggregator.model

class HomoNNClient(HomoNNBase):

    def __init__(self):
        super().__init__()

    def _init_model(self, param):
        super(HomoNNClient, self)._init_model(param)
        self.batch_size = param.batch_size
        self.aggregate_every_n_epoch = 1
        self.nn_define = param.nn_define
        self.config_type = param.config_type
        self.optimizer = param.optimizer
        self.loss = param.loss
        self.metrics = param.metrics
        self.data_converter = nn_model.get_data_converter(self.config_type)
        self.model_builder = nn_model.get_nn_builder(config_type=self.config_type)

    def _check_monitored_status(self, data, epoch_degree):
        metrics = self.nn_model.evaluate(data)
        Logger.info(f"metrics at iter {self.aggregator_iter}: {metrics}")
        loss = metrics["loss"]
        self.aggregator.send_loss(loss=loss,
                                  degree=epoch_degree,
                                  suffix=self._iter_suffix())
        return self.aggregator.get_converge_status(suffix=self._iter_suffix())

    def __build_nn_model(self, input_shape):
        self.nn_model = self.model_builder(input_shape=input_shape,
                                           nn_define=self.nn_define,
                                           optimizer=self.optimizer,
                                           loss=self.loss,
                                           metrics=self.metrics)

    def fit(self, data_inst, *args):

        data = self.data_converter.convert(data_inst, batch_size=self.batch_size)
        self.__build_nn_model(data.get_shape()[0])

        epoch_degree = float(len(data))

        while self.aggregator_iter < self.max_iter:
            Logger.info(f"start {self.aggregator_iter}_th aggregation")

            # train
            self.nn_model.train(data, aggregate_every_n_epoch=self.aggregate_every_n_epoch)

            # send model for aggregate, then set aggregated model to local
            modify_func: typing.Callable = functools.partial(self.aggregator.aggregate_then_get,
                                                             degree=epoch_degree * self.aggregate_every_n_epoch,
                                                             suffix=self._iter_suffix())
            self.nn_model.modify(modify_func)

            # calc loss and check convergence
            if self._check_monitored_status(data, epoch_degree):
                Logger.info(f"early stop at iter {self.aggregator_iter}")
                break

            Logger.info(f"role {self.role} finish {self.aggregator_iter}_th aggregation")
            self.aggregator_iter += 1
        else:
            Logger.warn(f"reach max iter: {self.aggregator_iter}, not converged")

    def export_model(self):
        return _build_model_dict(meta=self._get_meta(), param=self._get_param())

    def _get_meta(self):
        from federatedml.protobuf.generated import nn_model_meta_pb2
        meta_pb = nn_model_meta_pb2.NNModelMeta()
        meta_pb.params.CopyFrom(self.model_param.generate_pb())
        meta_pb.aggregate_iter = self.aggregator_iter
        return meta_pb

    def _get_param(self):
        from federatedml.protobuf.generated import nn_model_param_pb2
        param_pb = nn_model_param_pb2.NNModelParam()
        param_pb.saved_model_bytes = self.nn_model.export_model()
        return param_pb

    def predict(self, data_inst):
        data = self.data_converter.convert(data_inst, batch_size=self.batch_size)
        predict = self.nn_model.predict(data)
        num_output_units = predict.shape[1]

        if num_output_units == 1:
            kv = [(x[0], (0 if x[1][0] < 0.5 else 1, x[1][0])) for x in zip(data.get_keys(), predict)]
            pred_tbl = session.parallelize(kv, include_key=True)
            return data_inst.join(pred_tbl, lambda d, pred: [d.label, pred[0], pred[1], {"label": pred[0]}])
        else:
            kv = [(x[0], (x[1].argmax(), [float(e) for e in x[1]])) for x in zip(data.get_keys(), predict)]
            pred_tbl = session.parallelize(kv, include_key=True)
            return data_inst.join(pred_tbl,
                                  lambda d, pred: [d.label, pred[0],
                                                   pred[1][pred[0]] / sum(pred[1]),
                                                   {"raw_predict": pred[1]}])

    def _load_model(self, model_dict):
        model_dict = list(model_dict["model"].values())[0]
        model_obj = _extract_param(model_dict)
        meta_obj = _extract_meta(model_dict)
        self.model_param.restore_from_pb(meta_obj.params)
        self._init_model(self.model_param)
        self.aggregator_iter = meta_obj.aggregate_iter
        self.nn_model = restore_nn_model(self.config_type, model_obj.saved_model_bytes)

class HomoNNHost(HomoNNClient):

    def __init__(self):
        super().__init__()
        self.role = consts.HOST

class HomoNNGuest(HomoNNClient):

    def __init__(self):
        super().__init__()
        self.role = consts.GUEST
<|endoftext|>


<commit_before>from datetime import datetime
from django.db import models
from django.template.defaultfilters import filesizeformat
from django.utils.translation import ugettext_lazy as _

from feincms import settings
from feincms.translations import TranslatedObjectMixin, Translation,\
    TranslatedObjectManager

class Category(models.Model):
    title = models.CharField(_('title'), max_length=200)
    parent = models.ForeignKey('self', blank=True, null=True,
        related_name='children', limit_choices_to={'parent__isnull': True},
        verbose_name=_('parent'))

    class Meta:
        ordering = ['parent__title', 'title']
        verbose_name = _('category')
        verbose_name_plural = _('categories')

    def __unicode__(self):
        if self.parent_id:
            return u'%s - %s' % (self.parent.title, self.title)

        return self.title

class MediaFile(models.Model, TranslatedObjectMixin):
    from django.core.files.storage import FileSystemStorage
    fs = FileSystemStorage(location=settings.FEINCMS_MEDIALIBRARY_PATH,
                           base_url=settings.FEINCMS_MEDIALIBRARY_URL)

    file = models.FileField(_('file'), upload_to=settings.FEINCMS_MEDIALIBRARY_FILES, storage=fs)
    created = models.DateTimeField(_('created'), default=datetime.now)
    copyright = models.CharField(_('copyright'), max_length=200, blank=True)

    categories = models.ManyToManyField(Category, verbose_name=_('categories'))

    class Meta:
        verbose_name = _('media file')
        verbose_name_plural = _('media files')

    objects = TranslatedObjectManager()

    @classmethod
    def reconfigure(cls, upload_to=None, storage=None):
        f = cls._meta.get_field('file')
        # Ugh. Copied relevant parts from django/db/models/fields/files.py
        # FileField.__init__ (around line 225)
        if storage:
            f.storage = storage
        if upload_to:
            f.upload_to = upload_to
            if callable(upload_to):
                f.generate_filename = upload_to

    @property
    def file_name(self):
        return self.file.name

    @property
    def file_url(self):
        return self.file.url

    def get_absolute_url(self):
        return self.file_url

    @property
    def file_type(self):
        import re

        FILE_TYPES = (
            ('image', re.compile(r'.(jpg|jpeg|gif|png)$', re.IGNORECASE)),
            ('pdf',   re.compile(r'.pdf$', re.IGNORECASE)),
        )

        filename = self.file_name
        for ftype, pat in FILE_TYPES:
            if pat.search(filename):
                return ftype
        return 'other'

class MediaFileTranslation(Translation(MediaFile)):
    caption = models.CharField(_('caption'), max_length=200)

    class Meta:
        verbose_name = _('media file translation')
        verbose_name_plural = _('media file translations')

    def __unicode__(self):
        return u'%s (%s / %s)' % (
            self.caption,
            self.parent.file.name[21:], # only show filename
            filesizeformat(self.parent.file.size),
            )

<commit_msg>Trim superfluous properties.<commit_after>from datetime import datetime
from django.db import models
from django.template.defaultfilters import filesizeformat
from django.utils.translation import ugettext_lazy as _

from feincms import settings
from feincms.translations import TranslatedObjectMixin, Translation,\
    TranslatedObjectManager

class Category(models.Model):
    title = models.CharField(_('title'), max_length=200)
    parent = models.ForeignKey('self', blank=True, null=True,
        related_name='children', limit_choices_to={'parent__isnull': True},
        verbose_name=_('parent'))

    class Meta:
        ordering = ['parent__title', 'title']
        verbose_name = _('category')
        verbose_name_plural = _('categories')

    def __unicode__(self):
        if self.parent_id:
            return u'%s - %s' % (self.parent.title, self.title)

        return self.title

class MediaFile(models.Model, TranslatedObjectMixin):
    from django.core.files.storage import FileSystemStorage
    fs = FileSystemStorage(location=settings.FEINCMS_MEDIALIBRARY_PATH,
                           base_url=settings.FEINCMS_MEDIALIBRARY_URL)

    file = models.FileField(_('file'), upload_to=settings.FEINCMS_MEDIALIBRARY_FILES, storage=fs)
    created = models.DateTimeField(_('created'), default=datetime.now)
    copyright = models.CharField(_('copyright'), max_length=200, blank=True)

    categories = models.ManyToManyField(Category, verbose_name=_('categories'))

    class Meta:
        verbose_name = _('media file')
        verbose_name_plural = _('media files')

    objects = TranslatedObjectManager()

    @classmethod
    def reconfigure(cls, upload_to=None, storage=None):
        f = cls._meta.get_field('file')
        # Ugh. Copied relevant parts from django/db/models/fields/files.py
        # FileField.__init__ (around line 225)
        if storage:
            f.storage = storage
        if upload_to:
            f.upload_to = upload_to
            if callable(upload_to):
                f.generate_filename = upload_to

    def get_absolute_url(self):
        return self.file_url

    @property
    def file_type(self):
        import re

        FILE_TYPES = (
            ('image', re.compile(r'.(jpg|jpeg|gif|png)$', re.IGNORECASE)),
            ('pdf',   re.compile(r'.pdf$', re.IGNORECASE)),
        )

        filename = self.file_name
        for ftype, pat in FILE_TYPES:
            if pat.search(filename):
                return ftype
        return 'other'

class MediaFileTranslation(Translation(MediaFile)):
    caption = models.CharField(_('caption'), max_length=200)

    class Meta:
        verbose_name = _('media file translation')
        verbose_name_plural = _('media file translations')

    def __unicode__(self):
        return u'%s (%s / %s)' % (
            self.caption,
            self.parent.file.name[21:], # only show filename
            filesizeformat(self.parent.file.size),
            )

<|endoftext|>