<commit_before>class ConsoleView:

    def __init__(self, model):
        self.model = model
        self.pos_w    = max(1, len(str(len(self.model.player_list)))) if self.model.player_list else 0
        self.player_w = max([ len(p.name) for p in self.model.player_list ]) if self.model.player_list else 0
        self.wlr_w    = max(2, max([ len("%.2f" % round(p.wlr, 2)) for p in self.model.player_list ])) if self.model.player_list else 0
        self.ppg_w    = max(2, max([ len("%.2f" % round(p.ppg, 2)) for p in self.model.player_list ])) if self.model.player_list else 0
        self.gp_w     = max(2, max([ len(str(p.gp))   for p in self.model.player_list ])) if self.model.player_list else 0
        self.w_w      = max(1, max([ len(str(p.w))    for p in self.model.player_list ])) if self.model.player_list else 0
        self.l_w      = max(1, max([ len(str(p.l))    for p in self.model.player_list ])) if self.model.player_list else 0
        self.ot_w     = max(2, max([ len(str(p.otw))  for p in self.model.player_list ])) if self.model.player_list else 0
        self.gf_w     = max(2, max([ len(str(p.gf))   for p in self.model.player_list ])) if self.model.player_list else 0
        self.ga_w     = max(2, max([ len(str(p.ga))   for p in self.model.player_list ])) if self.model.player_list else 0
        self.diff_w   = max(4, max([ len(str(p.diff)) for p in self.model.player_list ])) if self.model.player_list else 0
        self.total_w  = self.pos_w + self.player_w + self.wlr_w + self.ppg_w + self.gp_w + self.w_w + self.l_w + self.ot_w + self.gf_w + self.ga_w + self.diff_w + 32

    def _get_filled_line(self, c):
        return self.format_string_line % ( c * self.pos_w
                                         , c * self.player_w
                                         , c * self.wlr_w
                                         , c * self.ppg_w
                                         , c * self.gp_w
                                         , c * self.w_w
                                         , c * self.l_w
                                         , c * self.ot_w
                                         , c * self.gf_w
                                         , c * self.ga_w
                                         , c * self.diff_w
                                         )

    def _get_header(self):
        rows = []
        rows.append(self._get_filled_line("-"))
        rows.append("|" + self.model.current_season.replace("_", " ").upper().center(self.total_w) + "|")
        rows.append(self._get_filled_line("-"))
        rows.append(self.format_string_row % ( "#"
                                             , "Player"
                                             , "WLR"
                                             , "PPG"
                                             , "GP"
                                             , "W"
                                             , "L"
                                             , "OT"
                                             , "GF"
                                             , "GA"
                                             , "DIFF"
                                             ))
        rows.append(self._get_filled_line("-"))
        return rows

    def _get_rows(self):
        self.format_string_row = (("| %%-%ds " + "| %%-%ds " * 10 + "|")) % ( self.pos_w
                                                                            , self.player_w
                                                                            , self.wlr_w
                                                                            , self.ppg_w
                                                                            , self.gp_w
                                                                            , self.w_w
                                                                            , self.l_w
                                                                            , self.ot_w
                                                                            , self.gf_w
                                                                            , self.ga_w
                                                                            , self.diff_w
                                                                            )
        self.format_string_line = ("|" + self.format_string_row[1:-1].replace(" ", "-").replace("|", "-") + "|")

        rows = []
        rows.extend(self._get_header())

        position = 0
        for player in self.model.player_list:
            position += 1
            rows.append(self.format_string_row % ( position
                                                 , player.name
                                                 , "%.2f" % player.wlr
                                                 , "%.2f" % player.ppg
                                                 , player.gp
                                                 , player.w
                                                 , player.l
                                                 , player.otw
                                                 , player.gf
                                                 , player.ga
                                                 , player.diff
                                                 ))

        rows.append(self._get_filled_line("-"))
        return rows

    def print_stats(self):
        for row in self._get_rows():
            print row

    def print_suggestions(self, player_name, excluded_players):
        excluded_players.append(player_name)
        least_played_player = self.model.get_least_played_player(player_name, excluded_players)
        player_least_played_games = self.model.get_player_least_played_games(excluded_players)
        leader = self.model.get_leader(excluded_players)
        player = self.model.get_player(player_name)

        print "Least played player            - %s, %s" % (least_played_player.name, "away" if least_played_player > player else "home")
        print "Player with least played games - %s, %s" % (player_least_played_games.name, "away" if player_least_played_games > player else "home")
        print "Season leader                  - %s, away" % leader.name

    def print_game_addition(self, home, away, home_score, away_score, overtime):
        overtime_print = ", OT" if overtime else ""
        print "Successfully added game \"%s @ %s: %d-%d%s\"" % ( away
                                                               , home
                                                               , away_score
                                                               , home_score
                                                               , overtime_print
                                                               )

class HtmlView():

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

    def _print_stats_table(self):
        print "
"
        print "
%s
" % self.model.current_season.replace("_", " ").upper()
        print "
"
        print "#	"
        print "Player	"
        print "WLR	"
        print "PPG	"
        print "GP	"
        print "W	"
        print "L	"
        print "OT	"
        print "GF	"
        print "GA	"
        print "DIFF	"
        print "
"

        position = 0
        for player in self.model.player_list:
            position += 1
            print "
"
            print "%d	"   % position
            print "%s	"   % player.name
            print "%.2f	" % player.wlr
            print "%.2f	" % player.ppg
            print "%d	"   % player.gp
            print "%d	"   % player.w
            print "%d	"   % player.l
            print "%d	"   % player.otw
            print "%d	"   % player.gf
            print "%d	"   % player.ga
            print "%d	"   % player.diff
            print "
"
        print "
"

    def _print_champions_table(self):
        print "
"
        print "
CHAMPIONS
"
        print "
"
        print "Season	"
        print "Player	"
        print "WLR	"
        print "PPG	"
        print "GP	"
        print "W	"
        print "L	"
        print "OT	"
        print "GF	"
        print "GA	"
        print "DIFF	"
        print "
"

        for season in [ season for season in self.model.seasons if season != self.model.current_season ]:
            player = self.model.get_leader(season=season)
            print "
"
            print "%s	"   % season.replace("_", " ").upper()
            print "%s	"   % player.name
            print "%.2f	" % player.wlr
            print "%.2f	" % player.ppg
            print "%d	"   % player.gp
            print "%d	"   % player.w
            print "%d	"   % player.l
            print "%d	"   % player.otw
            print "%d	"   % player.gf
            print "%d	"   % player.ga
            print "%d	"   % player.diff
            print "
"
        print "
"

    def _print_match_count_table(self):
        print "
"
        print "
GAMES PLAYED
"
        print "
"
        print "Player	"
        print "Player	"
        print "Games played	"
        print "
"

        printed = {}
        for player in self.model.player_list:
            printed[player.name] = []
            for away_player, count in player.match_counts.iteritems():
                if away_player in printed and not player.name in printed[away_player]:
                    print "
"
                    print "%s	" % player.name
                    print "%s	" % away_player
                    print "%d	" % count
                    print "
"
                    printed[player.name].append(away_player)
        print "
"

    def _print_suggestions_table(self):
        print "
"
        print "
MATCH SUGGESTIONS
"
        print "
"
        print "Player	"
        print "LPP	"
        print "LPG	"
        print "Leader	"
        print "
"

        printed = {}
        for player_name, player in self.model.player_dict.iteritems():
            least_played_player = self.model.get_least_played_player(player_name)
            player_least_played_games = self.model.get_player_least_played_games([player_name])
            leader = self.model.get_leader([player_name])
            print "
"
            print "%s	" % player.name
            print "%s	" % least_played_player.name
            print "%s	" % player_least_played_games.name
            print "%s	" % leader.name
            print "
"
        print "
"

    def print_html(self):
        print "<!DOCTYPE html>"
        print ""
        print "<title>NHL STEGEN 3.0</title><link rel=\"stylesheet\" href=\"styles.css\">"
        print ""
        self._print_stats_table()
        print "
"
        self._print_match_count_table()
        print "
"
        self._print_suggestions_table()
        print "
"
        self._print_champions_table()
        print ""
        print ""
<commit_msg>Better output<commit_after>class ConsoleView:

    def __init__(self, model):
        self.model = model
        self.pos_w    = max(1, len(str(len(self.model.player_list)))) if self.model.player_list else 0
        self.player_w = max([ len(p.name) for p in self.model.player_list ]) if self.model.player_list else 0
        self.wlr_w    = max(2, max([ len("%.2f" % round(p.wlr, 2)) for p in self.model.player_list ])) if self.model.player_list else 0
        self.ppg_w    = max(2, max([ len("%.2f" % round(p.ppg, 2)) for p in self.model.player_list ])) if self.model.player_list else 0
        self.gp_w     = max(2, max([ len(str(p.gp))   for p in self.model.player_list ])) if self.model.player_list else 0
        self.w_w      = max(1, max([ len(str(p.w))    for p in self.model.player_list ])) if self.model.player_list else 0
        self.l_w      = max(1, max([ len(str(p.l))    for p in self.model.player_list ])) if self.model.player_list else 0
        self.ot_w     = max(2, max([ len(str(p.otw))  for p in self.model.player_list ])) if self.model.player_list else 0
        self.gf_w     = max(2, max([ len(str(p.gf))   for p in self.model.player_list ])) if self.model.player_list else 0
        self.ga_w     = max(2, max([ len(str(p.ga))   for p in self.model.player_list ])) if self.model.player_list else 0
        self.diff_w   = max(4, max([ len(str(p.diff)) for p in self.model.player_list ])) if self.model.player_list else 0
        self.total_w  = self.pos_w + self.player_w + self.wlr_w + self.ppg_w + self.gp_w + self.w_w + self.l_w + self.ot_w + self.gf_w + self.ga_w + self.diff_w + 32

    def _get_filled_line(self, c):
        return self.format_string_line % ( c * self.pos_w
                                         , c * self.player_w
                                         , c * self.wlr_w
                                         , c * self.ppg_w
                                         , c * self.gp_w
                                         , c * self.w_w
                                         , c * self.l_w
                                         , c * self.ot_w
                                         , c * self.gf_w
                                         , c * self.ga_w
                                         , c * self.diff_w
                                         )

    def _get_header(self):
        rows = []
        rows.append(self._get_filled_line("-"))
        rows.append("|" + self.model.current_season.replace("_", " ").upper().center(self.total_w) + "|")
        rows.append(self._get_filled_line("-"))
        rows.append(self.format_string_row % ( "#"
                                             , "Player"
                                             , "WLR"
                                             , "PPG"
                                             , "GP"
                                             , "W"
                                             , "L"
                                             , "OT"
                                             , "GF"
                                             , "GA"
                                             , "DIFF"
                                             ))
        rows.append(self._get_filled_line("-"))
        return rows

    def _get_rows(self):
        self.format_string_row = (("| %%-%ds " + "| %%-%ds " * 10 + "|")) % ( self.pos_w
                                                                            , self.player_w
                                                                            , self.wlr_w
                                                                            , self.ppg_w
                                                                            , self.gp_w
                                                                            , self.w_w
                                                                            , self.l_w
                                                                            , self.ot_w
                                                                            , self.gf_w
                                                                            , self.ga_w
                                                                            , self.diff_w
                                                                            )
        self.format_string_line = ("|" + self.format_string_row[1:-1].replace(" ", "-").replace("|", "-") + "|")

        rows = []
        rows.extend(self._get_header())

        position = 0
        for player in self.model.player_list:
            position += 1
            rows.append(self.format_string_row % ( position
                                                 , player.name
                                                 , "%.2f" % player.wlr
                                                 , "%.2f" % player.ppg
                                                 , player.gp
                                                 , player.w - player.otw
                                                 , player.l
                                                 , player.otw
                                                 , player.gf
                                                 , player.ga
                                                 , player.diff
                                                 ))

        rows.append(self._get_filled_line("-"))
        return rows

    def print_stats(self):
        for row in self._get_rows():
            print row

    def print_suggestions(self, player_name, excluded_players):
        excluded_players.append(player_name)
        least_played_player = self.model.get_least_played_player(player_name, excluded_players)
        player_least_played_games = self.model.get_player_least_played_games(excluded_players)
        leader = self.model.get_leader(excluded_players)
        player = self.model.get_player(player_name)

        print "Least played player            - %s, %s" % (least_played_player.name, "away" if least_played_player > player else "home")
        print "Player with least played games - %s, %s" % (player_least_played_games.name, "away" if player_least_played_games > player else "home")
        print "Season leader                  - %s, away" % leader.name

    def print_game_addition(self, home, away, home_score, away_score, overtime):
        overtime_print = ", OT" if overtime else ""
        print "Successfully added game \"%s @ %s: %d-%d%s\"" % ( away
                                                               , home
                                                               , away_score
                                                               , home_score
                                                               , overtime_print
                                                               )

class HtmlView():

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

    def _print_stats_table(self):
        print "
"
        print "
%s
" % self.model.current_season.replace("_", " ").upper()
        print "
"
        print "#	"
        print "Player	"
        print "WLR	"
        print "PPG	"
        print "GP	"
        print "W	"
        print "L	"
        print "OT	"
        print "GF	"
        print "GA	"
        print "DIFF	"
        print "
"

        position = 0
        for player in self.model.player_list:
            position += 1
            print "
"
            print "%d	"   % position
            print "%s	"   % player.name
            print "%.2f	" % player.wlr
            print "%.2f	" % player.ppg
            print "%d	"   % player.gp
            print "%d	"   % player.w - player.otw
            print "%d	"   % player.l
            print "%d	"   % player.otw
            print "%d	"   % player.gf
            print "%d	"   % player.ga
            print "%d	"   % player.diff
            print "
"
        print "
"

    def _print_champions_table(self):
        print "
"
        print "
CHAMPIONS
"
        print "
"
        print "Season	"
        print "Player	"
        print "WLR	"
        print "PPG	"
        print "GP	"
        print "W	"
        print "L	"
        print "OT	"
        print "GF	"
        print "GA	"
        print "DIFF	"
        print "
"

        for season in [ season for season in self.model.seasons if season != self.model.current_season ]:
            player = self.model.get_leader(season=season)
            print "
"
            print "%s	"   % season.replace("_", " ").upper()
            print "%s	"   % player.name
            print "%.2f	" % player.wlr
            print "%.2f	" % player.ppg
            print "%d	"   % player.gp
            print "%d	"   % player.w - player.otw
            print "%d	"   % player.l
            print "%d	"   % player.otw
            print "%d	"   % player.gf
            print "%d	"   % player.ga
            print "%d	"   % player.diff
            print "
"
        print "
"

    def _print_match_count_table(self):
        print "
"
        print "
GAMES PLAYED
"
        print "
"
        print "Player	"
        print "Player	"
        print "Games played	"
        print "
"

        printed = {}
        for player in self.model.player_list:
            printed[player.name] = []
            for away_player, count in player.match_counts.iteritems():
                if away_player in printed and not player.name in printed[away_player]:
                    print "
"
                    print "%s	" % player.name
                    print "%s	" % away_player
                    print "%d	" % count
                    print "
"
                    printed[player.name].append(away_player)
        print "
"

    def _print_suggestions_table(self):
        print "
"
        print "
MATCH SUGGESTIONS
"
        print "
"
        print "Player	"
        print "LPP	"
        print "LPG	"
        print "Leader	"
        print "
"

        printed = {}
        for player_name, player in self.model.player_dict.iteritems():
            least_played_player = self.model.get_least_played_player(player_name)
            player_least_played_games = self.model.get_player_least_played_games([player_name])
            leader = self.model.get_leader([player_name])
            print "
"
            print "%s	" % player.name
            print "%s	" % least_played_player.name
            print "%s	" % player_least_played_games.name
            print "%s	" % leader.name
            print "
"
        print "
"

    def print_html(self):
        print "<!DOCTYPE html>"
        print ""
        print "<title>NHL STEGEN 3.0</title><link rel=\"stylesheet\" href=\"styles.css\">"
        print ""
        self._print_stats_table()
        print "
"
        self._print_match_count_table()
        print "
"
        self._print_suggestions_table()
        print "
"
        self._print_champions_table()
        print ""
        print ""
<|endoftext|>


<commit_before><commit_msg>Fix tests<commit_after><|endoftext|>


<commit_before><commit_msg>Ensure test example follows PEP 8 guidelines<commit_after><|endoftext|>


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

# TODO: Rekeying http://www.cisco.com/web/about/security/intelligence/securing-voip.html#12
# TODO: Standard configuration directory
# TODO: xdg

'''
A simple ffmpeg-based SRTP stream setup tool.

Doesn't have any ambitions of dealing with more complex matters such as
conferencing, NAT traversal, codec negotiation, etc—sane defaults and sane
environments only supported.
'''

# TODO: https://github.com/webrtc/samples
#       https://webrtchacks.com/chrome-extension/

from argparse import ArgumentParser
#from tempfile import NamedTemporaryFile
from base64 import b64encode
from socket import (socket, SOCK_STREAM, IPPROTO_TCP,
                    getaddrinfo, AI_ADDRCONFIG, AI_PASSIVE, AI_NUMERICHOST)
import subprocess
import logging
import json
import ssl

_DEFAULT_CODEC = 'opus'
_DEFAULT_CIPHERS = '!eNULL:!aNULL:kDHE+aRSA+HIGH'

def ssl_context_for(purpose, ca_certs, own_cert):
    ssl_context = ssl.create_default_context(purpose, cafile=ca_certs)
    #ssl_context = ssl.SSLContext(ssl.PROTOCOL_TLSv1_2)
    ssl_context.load_cert_chain(own_cert)
    if ca_certs is None:
        ssl_context.load_default_certs(purpose)
    else:
        ssl_context.load_verify_locations(cafile=ca_certs)
    ssl_context.verify_mode = ssl.CERT_REQUIRED
    ##ssl_context.verify_flags |= ssl.VERIFY_X509_STRICT
    # Since we use only DH KEX later, we have to provide DH params. They aren't
    # automatically generated. There are no compiled in ones. If you don't do
    # this, you get weird "No shared cipher" errors at the client hello.
    if purpose == ssl.Purpose.CLIENT_AUTH:
        ssl_context.load_dh_params(args.dh_params)
    # Enforce encryption and authentication.
    # Enforce perfect forward secrecy—only provided by Diffie-Hellman ephemeral
    # so far.
    # Enforce RSA-based authentication because of better failure modes.
    # Enforce 'high'—higher security suites.
    # See http://security.stackexchange.com/questions/5096/rsa-vs-dsa-for-ssh-authentication-keys/46781#46781.
    # TODO: Figure out how to enforce *generically*, better hash suites, and
    # not have outdated, slow, and known weaker ciphers like 3DES.
    ssl_context.set_ciphers(_DEFAULT_CIPHERS)
    ssl_context.set_alpn_protocols(['simplevoip/0'])

    return ssl_context

class NullFramedJSONSocket:
    '"socket"'

    def __init__(self, socket):
        self.socket = socket
        # Not the most efficient, but who cares here? It's a bloody control
        # channel, with small payoads.
        self.buffer = bytearray()

    def load(self):
        while self.buffer.rfind(b'\0') == -1:
            chunk = self.socket.recv(128)
            self.buffer += chunk
            if len(chunk) == 0:
                break
        body, _, self.buffer = self.buffer.partition(b'\0')
        return json.loads(body.decode())

    def dump(self, payload):
        self.socket.sendall(self._frame_json(payload))

    @staticmethod
    def _frame_json(payload):
        return json.dumps(payload).encode() + b'\0'

def ffmpeg(*args, **kwargs):
    return  subprocess.Popen(['ffmpeg', '-loglevel', 'warning',
                                        '-nostdin'] +
                             list(args),
                              **kwargs)

def ffmpeg_in(speaker, sdp):
    #with NamedTemporaryFile(mode='w') as fp:
    #    fp.write(sdp)
    #    fp.seek(0)
    popen = ffmpeg('-f', 'sdp',
                   '-i', 'pipe:',
                   '-f', 'alsa',
                   speaker,
                   stdin=subprocess.PIPE,
                   universal_newlines=True)
    # Debug code
    #popen.terminate()
    #popen.wait()
    #popen = subprocess.Popen(['ffplay', '-f', 'sdp',
    #                                    '-i', 'pipe:'] +
    #                         (['showmode', '1'] if args.log_level == logging.DEBUG else []),
    #                         stdin=subprocess.PIPE,
    #                         universal_newlines=True)
    # Not .communicate(), which tries to read stdout, and does a wait().
    with popen.stdin:
        popen.stdin.write(sdp)
    return popen

def ffmpeg_out(microphone, address):
    return ffmpeg('-f', 'alsa',
                  '-i', microphone,
                  '-f', 'rtp',
                  '-c:a', _DEFAULT_CODEC,
                  'rtp://{}:{}'.format(*address))

def server(sk, address, ssl_context):
    sk.bind(address)
    sk.listen(1)
    logging.debug('Listening for one connection at a time')
    while True:
        connection, address = sk.accept()
        logging.debug('TCP connectionf from %s', address)
        ssl_socket = ssl_context.wrap_socket(connection, server_side=True)
        logging.debug('TLS handshake')
        with ssl_socket:
            json_socket = NullFramedJSONSocket(ssl_socket)
            payload = {}
            if args.public_address:
                payload['public_address'] = args.public_address
            payload['audio_sdp'] = audio_sdp(*address)
            # TODO: Wait for application-layer message rather? For now, since
            #       the only thing we do is talk, nothing else, no point in
            #       waiting.
            json_socket.dump(payload)
            logging.debug('Sent %s', payload)
            response = json_socket.load()
            logging.debug('Got %s', response)

            in_ffmpeg = ffmpeg_in(args.speaker, response['audio_sdp'])
            logging.debug('ffmpeg listening')
            json_socket.dump({'clear_to_send': True})
            logging.debug('Sent CTS')
            while not json_socket.load().get('clear_to_send', False):
                pass
            logging.debug('Got CTS')

            out_ffmpeg = ffmpeg_out(args.microphone,
                                    response.get('public_address', address))
            logging.debug('ffmpeg sending')

            # TODO: Subprocess polling. Dirty trick with pipes?
            with in_ffmpeg, out_ffmpeg:
                pass
            logging.debug('Shutdown')

def client(sk, address, ssl_context):
    # TODO?
    ssl_context.check_hostname = True
    ssl_socket = ssl_context.wrap_socket(sk, server_hostname=args.name)
    ssl_socket.connect(address)

    with ssl_socket:
        json_socket = NullFramedJSONSocket(ssl_socket)
        payload = {}
        if args.public_address:
            payload['public_address'] = args.public_address
        payload['audio_sdp'] = audio_sdp(args.host, args.port)
        json_socket.dump(payload)
        logging.debug('Sent %s', payload)
        response = json_socket.load()
        logging.debug('Got %s', response)

        in_ffmpeg = ffmpeg_in(args.speaker, response['audio_sdp'])
        logging.debug('ffmpeg listening')
        json_socket.dump({'clear_to_send': True})
        logging.debug('Sent CTS')
        while not json_socket.load().get('clear_to_send', False):
            pass
        logging.debug('Got CTS')

        out_ffmpeg = ffmpeg_out(args.microphone,
                                response.get('public_address', address))
        logging.debug('ffmpeg sending')

        # TODO: Subprocess polling. Dirty trick with pipes?
        with in_ffmpeg, out_ffmpeg:
            pass
        logging.debug('Shutdown')

def audio_sdp(host, port):
    # FIXME: Why does it say c=… 127.0.0.1? We're not originating from
    # localhost!
    ffmpeg = subprocess.Popen(['ffmpeg',
                               '-loglevel', 'warning',
                               '-f', 'alsa',
                               '-i', args.microphone,
                               '-f', 'rtp',
                               '-t', '0',
                               '-c:a', _DEFAULT_CODEC,
                               'rtp://{}:{}'.format(host, port)],
                              universal_newlines=True,
                              stdout=subprocess.PIPE,
                              stderr=subprocess.PIPE)
    ffmpeg_stdout, ffmpeg_stderr = ffmpeg.communicate()
    ffmpeg.wait()
    if ffmpeg.returncode != 0:
        raise subprocess.SubprocessError(ffmpeg_stderr)
    return ffmpeg_stdout

# TODO: Use
def gen_srtp_params(direction):
    # Doesn't seem like ffmpeg supports RFC 5764 (DTLS-SRTP), despite
    # supporting some of the ciphers, so we have to do the key negotiation
    # ourselves, so we just exchange the master key and master salt over a
    # TCP/TLS channel.
    return ['-srtp_{:direction}_params'.format(direction=direction),
            b64encode(ssl.RAND_BYTES(30))]

def argument_parser():
    ap = ArgumentParser()

    ap.add_argument('-q', action='store_const', const=logging.WARNING, dest='log_level')
    ap.add_argument('-D', action='store_const', const=logging.DEBUG, dest='log_level')

    ap.add_argument('-n', '--numeric', action='store_true', default=False)

    ap.add_argument('--cert', required=True)
    ap.add_argument('--certs')

    ap.add_argument('--microphone', default='default', help='arecord -L')
    ap.add_argument('--speaker', default='default', help='aplay -L')
    ap.add_argument('--webcam', default='/dev/video0', help='v4l2-ctl --list-devices')

    ap.add_argument('-p', '--public-address', nargs=2)

    sub_aps = ap.add_subparsers()

    client_ap = sub_aps.add_parser('client', aliases=['c'])
    # The extent to which we help you work around NAT
    client_ap.add_argument('name')
    client_ap.add_argument('host')
    client_ap.add_argument('port', type=int)
    client_ap.set_defaults(listen=False)

    server_ap = sub_aps.add_parser('server', aliases=['s'])
    server_ap.add_argument('--dh-params', required=True)
    # The extent to which we help you work around NAT
    server_ap.add_argument('host')
    server_ap.add_argument('port', type=int)
    server_ap.set_defaults(listen=True)

    init_ap = sub_aps.add_parser('init')
    init_ap.add_argument('--dh-params', required=True)
    init_ap.set_defaults(init=True)

    return ap

if __name__ == '__main__':
    ap = argument_parser()
    args = ap.parse_args()

    if args.log_level is not None:
        logging.basicConfig(level=args.log_level)

    if getattr(args, 'init', False):
        subprocess.check_call(['openssl', 'req', '-new',
                                                 '-x509',
                                                 '-days', '365',
                                                 '-nodes',
                                                 '-out', args.cert,
                                                 '-keyout', args.cert])
        subprocess.check_call(['openssl', 'dhparam', '-out', args.dh_params,
                                                     '2048'])
        raise SystemExit()

    flags = 0
    if args.listen:
        flags |= AI_ADDRCONFIG | AI_PASSIVE
    if args.numeric:
        flags |= AI_NMERICHOST
    family, type_, proto, _, address = \
            getaddrinfo(args.host, args.port,
                        type=SOCK_STREAM, proto=IPPROTO_TCP,
                        flags=flags)[0]
    sk = socket(family, type_, proto)

    if args.listen:
        mainloop = server
        purpose = ssl.Purpose.CLIENT_AUTH
    else:
        mainloop = client
        purpose = ssl.Purpose.SERVER_AUTH
    ssl_context = ssl_context_for(purpose,
                                  ca_certs=args.certs, own_cert=args.cert)
    with sk:
        mainloop(sk, address, ssl_context)
<commit_msg>SRTP support.<commit_after>#! /usr/bin/env python3

# TODO: Rekeying http://www.cisco.com/web/about/security/intelligence/securing-voip.html#12
# TODO: Standard configuration directory
# TODO: xdg

'''
A simple ffmpeg-based SRTP stream setup tool.

Doesn't have any ambitions of dealing with more complex matters such as
conferencing, NAT traversal, codec negotiation, etc—sane defaults and sane
environments only supported.
'''

# TODO: https://github.com/webrtc/samples
#       https://webrtchacks.com/chrome-extension/

from argparse import ArgumentParser
#from tempfile import NamedTemporaryFile
from base64 import b64encode
from socket import (socket, SOCK_STREAM, IPPROTO_TCP,
                    getaddrinfo, AI_ADDRCONFIG, AI_PASSIVE, AI_NUMERICHOST)
import subprocess
import logging
import json
import ssl

_DEFAULT_CODEC = 'opus'
_DEFAULT_TLS_CIPHERS = '!eNULL:!aNULL:kDHE+aRSA+HIGH'
_DEFAULT_SRTP_CIPHER = 'AES_CM_128_HMAC_SHA1_80'

def ssl_context_for(purpose, ca_certs, own_cert):
    ssl_context = ssl.create_default_context(purpose, cafile=ca_certs)
    #ssl_context = ssl.SSLContext(ssl.PROTOCOL_TLSv1_2)
    ssl_context.load_cert_chain(own_cert)
    if ca_certs is None:
        ssl_context.load_default_certs(purpose)
    else:
        ssl_context.load_verify_locations(cafile=ca_certs)
    ssl_context.verify_mode = ssl.CERT_REQUIRED
    ##ssl_context.verify_flags |= ssl.VERIFY_X509_STRICT
    # Since we use only DH KEX later, we have to provide DH params. They aren't
    # automatically generated. There are no compiled in ones. If you don't do
    # this, you get weird "No shared cipher" errors at the client hello.
    if purpose == ssl.Purpose.CLIENT_AUTH:
        ssl_context.load_dh_params(args.dh_params)
    # Enforce encryption and authentication.
    # Enforce perfect forward secrecy—only provided by Diffie-Hellman ephemeral
    # so far.
    # Enforce RSA-based authentication because of better failure modes.
    # Enforce 'high'—higher security suites.
    # See http://security.stackexchange.com/questions/5096/rsa-vs-dsa-for-ssh-authentication-keys/46781#46781.
    # TODO: Figure out how to enforce *generically*, better hash suites, and
    # not have outdated, slow, and known weaker ciphers like 3DES.
    ssl_context.set_ciphers(_DEFAULT_TLS_CIPHERS)
    ssl_context.set_alpn_protocols(['simplevoip/0'])

    return ssl_context

class NullFramedJSONSocket:
    '"socket"'

    def __init__(self, socket):
        self.socket = socket
        # Not the most efficient, but who cares here? It's a bloody control
        # channel, with small payoads.
        self.buffer = bytearray()

    def load(self):
        while self.buffer.rfind(b'\0') == -1:
            chunk = self.socket.recv(128)
            self.buffer += chunk
            if len(chunk) == 0:
                break
        body, _, self.buffer = self.buffer.partition(b'\0')
        return json.loads(body.decode())

    def dump(self, payload):
        self.socket.sendall(self._frame_json(payload))

    @staticmethod
    def _frame_json(payload):
        return json.dumps(payload).encode() + b'\0'

def ffmpeg(*args, **kwargs):
    return  subprocess.Popen(['ffmpeg', '-loglevel', 'warning',
                                        '-nostdin'] +
                             list(args),
                              **kwargs)

def ffmpeg_in(speaker, sdp):
    #with NamedTemporaryFile(mode='w') as fp:
    #    fp.write(sdp)
    #    fp.seek(0)
    popen = ffmpeg('-f', 'sdp',
                   '-i', 'pipe:',
                   '-f', 'alsa',
                   speaker,
                   stdin=subprocess.PIPE,
                   universal_newlines=True)
    # Debug code
    #popen.terminate()
    #popen.wait()
    #popen = subprocess.Popen(['ffplay', '-f', 'sdp',
    #                                    '-i', 'pipe:'] +
    #                         (['showmode', '1'] if args.log_level == logging.DEBUG else []),
    #                         stdin=subprocess.PIPE,
    #                         universal_newlines=True)
    # Not .communicate(), which tries to read stdout, and does a wait().
    with popen.stdin:
        popen.stdin.write(sdp)
    return popen

def ffmpeg_out(microphone, address, srtp_params):
    return ffmpeg('-f', 'alsa',
                  '-i', microphone,
                  '-f', 'rtp',
                  '-c:a', _DEFAULT_CODEC,
                  *srtp_params,
                  'srtp://{}:{}'.format(*address))

def server(sk, address, ssl_context):
    sk.bind(address)
    sk.listen(1)
    logging.debug('Listening for one connection at a time')
    while True:
        connection, address = sk.accept()
        logging.debug('TCP connectionf from %s', address)
        ssl_socket = ssl_context.wrap_socket(connection, server_side=True)
        logging.debug('TLS handshake')
        with ssl_socket:
            json_socket = NullFramedJSONSocket(ssl_socket)
            payload = {}
            if args.public_address:
                payload['public_address'] = args.public_address
            srtp_key = ssl.RAND_bytes(30)
            srtp_params = gen_srtp_params(srtp_key)
            payload['audio_sdp'] = audio_sdp(*address, srtp_params=srtp_params)
            # TODO: Wait for application-layer message rather? For now, since
            #       the only thing we do is talk, nothing else, no point in
            #       waiting.
            json_socket.dump(payload)
            logging.debug('Sent %s', payload)
            response = json_socket.load()
            logging.debug('Got %s', response)

            inbound_media = ffmpeg_in(args.speaker, response['audio_sdp'])
            logging.debug('ffmpeg listening')
            json_socket.dump({'clear_to_send': True})
            logging.debug('Sent CTS')
            while not json_socket.load().get('clear_to_send', False):
                pass
            logging.debug('Got CTS')

            outbound_media = ffmpeg_out(args.microphone,
                                        response.get('public_address', address),
                                        srtp_params=srtp_params)
            logging.debug('ffmpeg sending')

            # TODO: Subprocess polling. Dirty trick with pipes?
            with inbound_media, outbound_media:
                pass
            logging.debug('Shutdown')

def client(sk, address, ssl_context):
    # TODO?
    ssl_context.check_hostname = True
    ssl_socket = ssl_context.wrap_socket(sk, server_hostname=args.name)
    ssl_socket.connect(address)

    with ssl_socket:
        json_socket = NullFramedJSONSocket(ssl_socket)
        payload = {}
        if args.public_address:
            payload['public_address'] = args.public_address
        srtp_key = ssl.RAND_bytes(30)
        srtp_params = gen_srtp_params(srtp_key)
        payload['audio_sdp'] = audio_sdp(*address, srtp_params=srtp_params)
        json_socket.dump(payload)
        logging.debug('Sent %s', payload)
        response = json_socket.load()
        logging.debug('Got %s', response)

        inbound_media = ffmpeg_in(args.speaker, response['audio_sdp'])
        logging.debug('ffmpeg listening')
        json_socket.dump({'clear_to_send': True})
        logging.debug('Sent CTS')
        while not json_socket.load().get('clear_to_send', False):
            pass
        logging.debug('Got CTS')

        outbound_media = ffmpeg_out(args.microphone,
                                    response.get('public_address', address),
                                    srtp_params=srtp_params)
        logging.debug('ffmpeg sending')

        # TODO: Subprocess polling. Dirty trick with pipes?
        with inbound_media, outbound_media:
            pass
        logging.debug('Shutdown')

def audio_sdp(host, port, srtp_params):
    # FIXME: Why does it say c=… 127.0.0.1? We're not originating from
    # localhost!
    popen = ffmpeg('-f', 'alsa',
                   '-i', args.microphone,
                   '-f', 'rtp',
                   '-t', '0',
                   '-c:a', _DEFAULT_CODEC,
                   *srtp_params,
                   'srtp://{}:{}'.format(host, port),
                   universal_newlines=True,
                   stdout=subprocess.PIPE,
                   stderr=subprocess.PIPE)
    ffmpeg_stdout, ffmpeg_stderr = popen.communicate()
    popen.wait()
    if popen.returncode != 0:
        raise subprocess.SubprocessError(ffmpeg_stderr)
    return ffmpeg_stdout

# TODO: Use
def gen_srtp_params(srtp_key):
    # Doesn't seem like ffmpeg supports RFC 5764 (DTLS-SRTP), despite
    # supporting some of the ciphers, so we have to do the key negotiation
    # ourselves, so we just exchange the master key and master salt over a
    # TCP/TLS channel.
    return ['-srtp_out_suite', _DEFAULT_SRTP_CIPHER,
            '-srtp_out_params', b64encode(srtp_key)]

def argument_parser():
    ap = ArgumentParser()

    ap.add_argument('-q', action='store_const', const=logging.WARNING, dest='log_level')
    ap.add_argument('-D', action='store_const', const=logging.DEBUG, dest='log_level')

    ap.add_argument('-n', '--numeric', action='store_true', default=False)

    ap.add_argument('--cert', required=True)
    ap.add_argument('--certs')

    ap.add_argument('--microphone', default='default', help='arecord -L')
    ap.add_argument('--speaker', default='default', help='aplay -L')
    ap.add_argument('--webcam', default='/dev/video0', help='v4l2-ctl --list-devices')

    ap.add_argument('-p', '--public-address', nargs=2)

    sub_aps = ap.add_subparsers()

    client_ap = sub_aps.add_parser('client', aliases=['c'])
    # The extent to which we help you work around NAT
    client_ap.add_argument('name')
    client_ap.add_argument('host')
    client_ap.add_argument('port', type=int)
    client_ap.set_defaults(listen=False)

    server_ap = sub_aps.add_parser('server', aliases=['s'])
    server_ap.add_argument('--dh-params', required=True)
    # The extent to which we help you work around NAT
    server_ap.add_argument('host')
    server_ap.add_argument('port', type=int)
    server_ap.set_defaults(listen=True)

    init_ap = sub_aps.add_parser('init')
    init_ap.add_argument('--dh-params', required=True)
    init_ap.set_defaults(init=True)

    return ap

if __name__ == '__main__':
    ap = argument_parser()
    args = ap.parse_args()

    if args.log_level is not None:
        logging.basicConfig(level=args.log_level)

    if getattr(args, 'init', False):
        subprocess.check_call(['openssl', 'req', '-new',
                                                 '-x509',
                                                 '-days', '365',
                                                 '-nodes',
                                                 '-out', args.cert,
                                                 '-keyout', args.cert])
        subprocess.check_call(['openssl', 'dhparam', '-out', args.dh_params,
                                                     '2048'])
        raise SystemExit()

    flags = 0
    if args.listen:
        flags |= AI_ADDRCONFIG | AI_PASSIVE
    if args.numeric:
        flags |= AI_NMERICHOST
    family, type_, proto, _, address = \
            getaddrinfo(args.host, args.port,
                        type=SOCK_STREAM, proto=IPPROTO_TCP,
                        flags=flags)[0]
    sk = socket(family, type_, proto)

    if args.listen:
        mainloop = server
        purpose = ssl.Purpose.CLIENT_AUTH
    else:
        mainloop = client
        purpose = ssl.Purpose.SERVER_AUTH
    ssl_context = ssl_context_for(purpose,
                                  ca_certs=args.certs, own_cert=args.cert)
    with sk:
        mainloop(sk, address, ssl_context)
<|endoftext|>


<commit_before># -*- coding: utf-8 -*-
"""
Created on Wed Oct 19 17:16:40 2016

@author: ee10sjo
"""
import os
import numpy as np
import pandas as pd

folder = 'E:\\NormalFault\\Tltd_1degS_2km_30L'

os.chdir(folder)

fl = os.listdir(folder)

### For testing, selects horizon Hz5.ts

with open(fl[26], 'r') as f:
    data = f.readlines()

def z_invert(data):
    # split lines, remove linebreak marker
    for i in np.arange(0, len(data)):
        data[i] = data[i].strip()

    # split columns deliminated by space
    for i in np.arange(0, len(data)):
        data[i] = data[i].split()

    # multiply z by -1
    for i in np.arange(0,len(data)):
        if data[i][0] == 'VRTX':
            data[i][-1] = str(np.multiply(float(data[i][-1]), -1))

    # replace header keyword to indicate depth format
    data[10] = ['ZNEGATIVE', 'Depth']

    return data

data_inv = z_invert(data)

d = data_inv

##with open("newfile.ts", "w") as f:
for i in d:
    out = "\n".join(d)
    f.writelines(out)
#
#for line in d:
#    d[line] = ['\t'.join(d[line])]

#adds tab to every break
with open("newfile.ts", "w") as f:
    for i in d:
        out = ["%s\t" % item for item in i]
        f.writelines(out)

<commit_msg>Add beginning of trimech manipulation tool.<commit_after># -*- coding: utf-8 -*-
"""
Created on Wed Oct 19 17:16:40 2016

@author: ee10sjo
"""
import os
import numpy as np
import pandas as pd

#folder = 'E:\\NormalFault\\Tltd_1degS_2km_30L'
folder = 'Y:\\ee10sjo\\Oldfields_PhD\\4_Applications\\MOVE\\Tltd_1degS_2km_30L'

os.chdir(folder)

fl = os.listdir(folder)

### For testing, selects horizon Hz5.ts

with open(fl[26], 'r') as f:
    data = f.readlines()

def z_invert(data):
    # split lines, remove linebreak marker
    for i in np.arange(0, len(data)):
        data[i] = data[i].strip()

    # split columns deliminated by space
    for i in np.arange(0, len(data)):
        data[i] = data[i].split()

    # multiply z by -1
    for i in np.arange(0,len(data)):
        # select only vertices (does not edit individual triangles)
        if data[i][0] == 'VRTX':
            # Multiple each vertex z by -1
            data[i][-1] = str(np.multiply(float(data[i][-1]), -1))

            data[i][-1] = str((float(data[i][-1])+2000))

    # replace header keyword to indicate depth format
    data[10] = ['ZNEGATIVE', 'Depth']

    return data

data_inv = z_invert(data)

d = data_inv

with open(fl[26], 'w') as f:
    for i in d:
        out = "\n".join(str(d))
        f.writelines(out)

#for line in d:
#    d[line] = ['\t'.join(d[line])]

#adds tab to every break
with open("newfile.ts", "w") as f:
    for i in d:
        out = ["%s\t" % item for item in i]
        f.writelines(out)

<|endoftext|>


<commit_before>import json
import logging
import os
import random
import time
import urllib2

import cv2

import cache
import data_augmentation
import downloader

logger = logging.getLogger(__name__)

def _parse_url_file(file_data, separator=" ", get_utf=False):
  """ ImageNet generally uses a specific format to store the URLs of images.
  This function parses data in that format.
  Args:
    file_data: The input data to parse.
    separator: Default string used to separate WNID from URL in the file.
    get_utf: If this is True, it will, in addition, return a version of the same
    output that is encoded in UTF-8, which is more efficient if we're dumping to
    a JSON file, since we need to encode it in this format anyway.
  Returns:
    A list of tuples, where each tuple contains a wnid of an image, and the URL
    of that image. """
  lines = file_data.split("\n")[:-1]
  # Split ids and urls.
  mappings = []
  utf_mappings = []
  for line in lines:
    if not line.startswith("n"):
      # Bad line.
      continue

    wnid, url = line.split(separator, 1)
    try:
      unicode_url = url.decode("utf8")
    except UnicodeDecodeError:
      # We got a URL that's not valid unicode. (It does happen.)
      logger.warning("Skipping invalid URL: %s", url)
      continue

    mappings.append([wnid, url])
    if get_utf:
      utf_mappings.append([wnid, unicode_url])

  if get_utf:
    return mappings, utf_mappings
  return mappings

class ImageGetter(object):
  """ Gets random sets of images for use in training and testing. """
  def __init__(self, synset_location, batch_size, test_percentage=0.1,
               download_words=False):
    """
    Args:
      synset_location: Where to save synsets. Will be created if it doesn't
      exist.
      batch_size: The size of each batch to load.
      test_percentage: The percentage of the total images that will be used for
      testing.
      download_words: Whether to download the words for each synset as well as
      just the numbers. """
    self.__cache = cache.DiskCache("image_cache", 50000000000,
                                   download_words=download_words)
    self.__mem_buffer = cache.MemoryBuffer(224, batch_size, channels=3)
    self.__download_manager = downloader.DownloadManager(200,
        self.__cache, self.__mem_buffer)
    self.__batch_size = batch_size

    self.__synset_location = synset_location

    self._synsets = {}
    self._populate_synsets()

    # Calculate and store sizes for each synset.
    self.__synset_sizes = {}
    self.__total_images = 0
    for synset, urls in self._synsets.iteritems():
      self.__synset_sizes[synset] = len(urls)
      self.__total_images += len(urls)

    logger.info("Have %d total images in database." % (self.__total_images))

    self.__choose_test_images(test_percentage)

  def _populate_synsets(self):
    """ Populates the synset dictionary. """
    if not os.path.exists(self.__synset_location):
      os.mkdir(self.__synset_location)

    loaded = self.__load_synsets()
    self.__get_image_list(loaded)

  def __choose_test_images(self, test_percentage):
    """ Chooses the images to use in the testing set.
    Args:
      test_percentage: The percentage of the total number of images that are
      testing images. """
    required_images = int(self.__total_images * test_percentage)

    self.__already_picked = set([])
    self.__test_images = set([])
    for i in range(0, required_images):
      wnid, _, index = self.__pick_random_image()
      synset, _ = wnid.split("_")
      self.__test_images.add((synset, index))

  def __get_image_list(self, loaded):
    """ Downloads a comprehensive list of all images available.
    Args:
      loaded: Set of synsets that were already loaded successfully. """
    logger.info("Downloading list of synsets...")
    url = "http://www.image-net.org/api/text/imagenet.synset.obtain_synset_list"
    response = urllib2.urlopen(url, timeout=1000)
    synsets = set(response.read().split("\n")[:-1])

    # Remove any that were already loaded.
    synsets = synsets - loaded
    logger.info("Downloading %d synsets." % (len(synsets)))
    logger.debug("Downloading synsets: %s", synsets)

    # Get the urls for each synset.
    base_url = \
        "http://www.image-net.org/api/text/" \
        "imagenet.synset.geturls.getmapping?wnid=%s"
    for synset in synsets:
      if not synset:
        continue

      logger.info("Downloading urls for synset: %s", synset)
      response = urllib2.urlopen(base_url % (synset), timeout=1000)

      # Parse the image data.
      mappings, utf_mappings = _parse_url_file(response.read(), get_utf=True)

      self._synsets[synset] = mappings
      # Save it for later.
      self.__save_synset(synset, utf_mappings)

  def condense_loaded_synsets(self):
    """ Using is-a relationships obtained from ImageNet, it combines smaller
    synsets until only ones that have a sufficient amount of data remain. """
    logger.info("Condensing synsets...")

    # First, download the list of is-a relationships.
    url = "http://www.image-net.org/archive/wordnet.is_a.txt"
    response = urllib2.urlopen(url)
    lines = response.read().split("\n")[:-1]

    # Convert to actual mappings. Here, the keys are a synset, and the values
    # are what that synset is.
    mappings = {}
    for line in lines:
      if not line:
        # Bad line.
        continue
      sup, sub = line.split()
      mappings[sub] = sup

    # Now, go through our synsets and combine any that are too small.
    while True:
      sizes_to_remove = []
      for synset, size in self.__synset_sizes.iteritems():
        merged_at_least_one = False
        if size < self.MINIMUM_SYNSET_SIZE:
          if synset in mappings:
            # Combine with superset.
            superset = mappings[synset]
            if superset not in self.__synsets:
              logger.warning("Superset %s is not loaded!" % (superset))
            else:
              logger.info("Merging %s with %s." % (synset, superset))

              # It looks like synsets don't by default include sub-synsets.
              self.__synsets[superset].extend(self.__synsets[synset])
              self.__synset_sizes[superset] += self.__synset_sizes[synset]
              logger.debug("New superset size: %d." % \
                          (self.__synset_sizes[superset]))

          # Delete it since it's too small.
          logger.info("Deleting %s with size %d." % (synset, size))
          self.__synsets.pop(synset)
          # We can't pop items from the dictionary while we're iterating.
          sizes_to_remove.append(synset)

          merged_at_least_one = True

      for synset in sizes_to_remove:
        self.__synset_sizes.pop(synset)

      if not merged_at_least_one:
        # We're done here.
        logger.info("Finished merging synsets.")
        break

  def __save_synset(self, synset, data):
    """ Saves a synset to a file.
    Args:
      synset: The name of the synset.
      data: The data that will be saved in the file. """
    if not self.__synset_location:
      # We're not saving synsets.
      return

    logger.info("Saving synset: %s" % (synset))

    synset_path = os.path.join(self.__synset_location, "%s.json" % (synset))
    synset_file = open(synset_path, "w")
    json.dump(data, synset_file)
    synset_file.close()

  def __load_synsets(self):
    """ Loads synsets from a file.
    Returns:
      A set of the synsets that were loaded successfully. """
    loaded = set([])
    loaded_text = []
    logger.info("Loading synsets...")

    for path in os.listdir(self.__synset_location):
      if path.endswith(".json"):
        # Load synset from file.
        synset_name = path[:-5]
        logger.debug("Loading synset %s." % (synset_name))
        full_path = os.path.join(self.__synset_location, path)

        synset_file = file(full_path)
        loaded_text.append((synset_name, synset_file.read()))
        synset_file.close()

    logger.info("De-jsoning...")
    for synset_name, text in loaded_text:
      # Storing stuff in the default unicode format takes up a
      # massive amount of memory.
      self._synsets[synset_name] = \
          [[elem[0].encode("utf8"), elem[1].encode("utf8")] \
            for elem in json.loads(text)]

      loaded.add(synset_name)

    return loaded

  def __pick_random_image(self):
    """ Picks a random image from our database.
    Returns:
      wnid and url, and index of the image in it's synset list. """
    # Pick a random image.
    image_number = random.randint(0, self.__total_images - 1)
    # Traverse our map of synset sizes until we find this image.
    total = 0
    for synset, size in self.__synset_sizes.iteritems():
      total += size
      if total > image_number:
        # This is the image we want.
        use_synset = synset
        use_index = size - (total - image_number)
        break

    if (use_synset, use_index) in self.__test_images:
      # This is a test image, don't use it here.
      return self.__pick_random_image()

    wnid, url = self._synsets[use_synset][use_index]

    if wnid in self.__already_picked:
      # This is a duplicate, pick another.
      return self.__pick_random_image()
    self.__already_picked.add(wnid)

    return [wnid, url, use_index]

  def __pick_random_test_image(self):
    """ Picks a random image from the test set.
    Returns:
      wnid and url of the image. """
    synset, index = random.choice(tuple(self.__test_images))
    return self._synsets[synset][index]

  def get_random_batch(self, testing=False):
    """ Loads a random batch of images from the whole dataset.
    Args:
      testing: Whether to use testing images. Defaults to False, in which case
      it will use training images.
    Returns:
      The array of loaded images, and a list of the synsets each image belongs
      to. """
    logger.info("Getting batch.")

    self.__mem_buffer.clear()

    # Keeps track of images that were already used this batch.
    self.__already_picked = set([])

    # Try to download initial images.
    for _ in range(0, self.__batch_size):
      self.__load_random_image(testing=testing)

    # Wait for all the downloads to complete, replacing any ones that fail.
    while True:
      failures = self.__download_manager.get_failures()
      if len(failures):
        logger.debug("Replacing %d failed downloads." % (len(failures)))

      # Remove failed images.
      for synset, name, url in failures:
        # Load a new image to replace it.
        self.__load_random_image(testing=testing)

        # Remove failed image.
        wnid = "%s_%s" % (synset, name)
        logger.info("Removing bad image %s." % (wnid))
        self._synsets[synset].remove([wnid, url])
        self.__synset_sizes[synset] -= 1
        self.__total_images -= 1
        # Update the json file.
        self.__save_synset(synset, self._synsets[synset])

      if not self.__download_manager.update():
        break

      time.sleep(0.2)

    return self.__mem_buffer.get_storage()

  def __load_random_image(self, testing=False):
    """ Loads a random image from either the cache or the internet. If loading
    from the internet, it adds it to the download manager, otherwise, it adds
    them to the memory buffer.
    Args:
      testing: Whether to use testing images. Defaults to False, in which case
      it will use training images. """
    if testing:
      wnid, url = self.__pick_random_test_image()
    else:
      wnid, url, _ = self.__pick_random_image()
    synset, number = wnid.split("_")

    image = self.__get_cached_image(synset, number)
    if image is None:
      # We have to download the image instead.
      self.__download_manager.download_new(synset, number, url)
    else:
      self.__mem_buffer.add(image, number, synset)

  def __get_cached_image(self, synset, image_number):
    """ Checks if an image is in the cache, and returns it.
    Args:
      synset: The synset of the image.
      image_number: The image number in the synset.
    Returns:
      The actual image data, or None if the image is not in the cache. """
    # First check in the cache for the image.
    cached_image = self.__cache.get(synset, image_number)
    if cached_image != None:
      # We had a cached copy, so we're done.
      logger.debug("Found image in cache: %s_%s" % (synset, image_number))

      # Select a patch.
      patches = data_augmentation.extract_patches(cached_image)
      cached_image = patches[random.randint(0, len(patches) - 1)]

      return cached_image

    return None

  def get_synsets(self):
    return self._synsets

class FilteredImageGetter(ImageGetter):
  """ Works like an ImageGetter, but only loads images from a specific
  pre-defined file containing image names and their URLs. """

  def __init__(self, url_file, batch_size, **kwargs):
    """
    Args:
      url_file: Name of the file containing the images to use.
      batch_size: The size of each batch. """
    self.__url_file = url_file

    super(FilteredImageGetter, self).__init__(None, batch_size, **kwargs)

  def _populate_synsets(self):
    """ Populates the synsets dict. """
    # Load data from the url file.
    urls = file(self.__url_file).read()
    images = _parse_url_file(urls, separator="\t")

    # Separate it by synset and load it into the dictionary.
    for wnid, url in images:
      synset, _ = wnid.split("_")

      if synset not in self._synsets:
        self._synsets[synset] = []
      self._synsets[synset].append([wnid, url])
<commit_msg>Refactor image_getter.<commit_after>import json
import logging
import os
import random
import time
import urllib2

import cv2

import cache
import data_augmentation
import downloader

logger = logging.getLogger(__name__)

def _parse_url_file(file_data, separator=" ", get_utf=False):
  """ ImageNet generally uses a specific format to store the URLs of images.
  This function parses data in that format.
  Args:
    file_data: The input data to parse.
    separator: Default string used to separate WNID from URL in the file.
    get_utf: If this is True, it will, in addition, return a version of the same
    output that is encoded in UTF-8, which is more efficient if we're dumping to
    a JSON file, since we need to encode it in this format anyway.
  Returns:
    A list of tuples, where each tuple contains a wnid of an image, and the URL
    of that image. """
  lines = file_data.split("\n")[:-1]
  # Split ids and urls.
  mappings = []
  utf_mappings = []
  for line in lines:
    if not line.startswith("n"):
      # Bad line.
      continue

    wnid, url = line.split(separator, 1)
    try:
      unicode_url = url.decode("utf8")
    except UnicodeDecodeError:
      # We got a URL that's not valid unicode. (It does happen.)
      logger.warning("Skipping invalid URL: %s", url)
      continue

    mappings.append([wnid, url])
    if get_utf:
      utf_mappings.append([wnid, unicode_url])

  if get_utf:
    return mappings, utf_mappings
  return mappings

class ImageGetter(object):
  """ Gets random sets of images for use in training and testing. """

  def __init__(self, synset_location, batch_size, test_percentage=0.1,
               download_words=False):
    """
    Args:
      synset_location: Where to save synsets. Will be created if it doesn't
      exist.
      batch_size: The size of each batch to load.
      test_percentage: The percentage of the total images that will be used for
      testing.
      download_words: Whether to download the words for each synset as well as
      just the numbers. """
    self.__cache = cache.DiskCache("image_cache", 50000000000,
                                   download_words=download_words)
    self.__batch_size = batch_size

    self.__synset_location = synset_location

    self._synsets = {}
    self._populate_synsets()

    # Calculate and store sizes for each synset.
    self.__synset_sizes = {}
    self.__total_images = 0
    for synset, urls in self._synsets.iteritems():
      self.__synset_sizes[synset] = len(urls)
      self.__total_images += len(urls)

    # Make internal datasets for training and testing.
    train, test = self.__split_train_test_images(test_percentage)
    self.__train_set = _Dataset(train, self.__cache, self.__batch_size)
    self.__test_set = _Dataset(test, self.__cache, self.__batch_size)

  def _populate_synsets(self):
    """ Populates the synset dictionary. """
    if not os.path.exists(self.__synset_location):
      os.mkdir(self.__synset_location)

    loaded = self.__load_synsets()
    self.__get_image_list(loaded)

  def __split_train_test_images(self, test_percentage):
    """ Chooses the images to use in the training set and testing set.
    NOTE: For memory reasons, it also removes the images from the synset map.
    Args:
      test_percentage: The percentage of the total number of images that are
      testing images.
    Returns:
      The list of training images, and the list of testing images. """
    train = []
    test = []

    for synset in self._synsets.keys():
      for image in self._synsets[synset]:
        rand = random.random()
        if rand < test_percentage:
          test.append(tuple(image))
        else:
          train.append(tuple(image))

      # Remove the synset.
      self._synsets.pop(synset)

    return train, test

  def __get_image_list(self, loaded):
    """ Downloads a comprehensive list of all images available.
    Args:
      loaded: Set of synsets that were already loaded successfully. """
    logger.info("Downloading list of synsets...")
    url = "http://www.image-net.org/api/text/imagenet.synset.obtain_synset_list"
    response = urllib2.urlopen(url, timeout=1000)
    synsets = set(response.read().split("\n")[:-1])

    # Remove any that were already loaded.
    synsets = synsets - loaded
    logger.info("Downloading %d synsets." % (len(synsets)))
    logger.debug("Downloading synsets: %s", synsets)

    # Get the urls for each synset.
    base_url = \
        "http://www.image-net.org/api/text/" \
        "imagenet.synset.geturls.getmapping?wnid=%s"
    for synset in synsets:
      if not synset:
        continue

      logger.info("Downloading urls for synset: %s", synset)
      response = urllib2.urlopen(base_url % (synset), timeout=1000)

      # Parse the image data.
      mappings, utf_mappings = _parse_url_file(response.read(), get_utf=True)

      self._synsets[synset] = mappings
      # Save it for later.
      self.__save_synset(synset, utf_mappings)

  def condense_loaded_synsets(self):
    """ Using is-a relationships obtained from ImageNet, it combines smaller
    synsets until only ones that have a sufficient amount of data remain. """
    logger.info("Condensing synsets...")

    # First, download the list of is-a relationships.
    url = "http://www.image-net.org/archive/wordnet.is_a.txt"
    response = urllib2.urlopen(url)
    lines = response.read().split("\n")[:-1]

    # Convert to actual mappings. Here, the keys are a synset, and the values
    # are what that synset is.
    mappings = {}
    for line in lines:
      if not line:
        # Bad line.
        continue
      sup, sub = line.split()
      mappings[sub] = sup

    # Now, go through our synsets and combine any that are too small.
    while True:
      sizes_to_remove = []
      for synset, size in self.__synset_sizes.iteritems():
        merged_at_least_one = False
        if size < self.MINIMUM_SYNSET_SIZE:
          if synset in mappings:
            # Combine with superset.
            superset = mappings[synset]
            if superset not in self.__synsets:
              logger.warning("Superset %s is not loaded!" % (superset))
            else:
              logger.info("Merging %s with %s." % (synset, superset))

              # It looks like synsets don't by default include sub-synsets.
              self.__synsets[superset].extend(self.__synsets[synset])
              self.__synset_sizes[superset] += self.__synset_sizes[synset]
              logger.debug("New superset size: %d." % \
                          (self.__synset_sizes[superset]))

          # Delete it since it's too small.
          logger.info("Deleting %s with size %d." % (synset, size))
          self.__synsets.pop(synset)
          # We can't pop items from the dictionary while we're iterating.
          sizes_to_remove.append(synset)

          merged_at_least_one = True

      for synset in sizes_to_remove:
        self.__synset_sizes.pop(synset)

      if not merged_at_least_one:
        # We're done here.
        logger.info("Finished merging synsets.")
        break

  def __save_synset(self, synset, data):
    """ Saves a synset to a file.
    Args:
      synset: The name of the synset.
      data: The data that will be saved in the file. """
    if not self.__synset_location:
      # We're not saving synsets.
      return

    logger.info("Saving synset: %s" % (synset))

    synset_path = os.path.join(self.__synset_location, "%s.json" % (synset))
    synset_file = open(synset_path, "w")
    json.dump(data, synset_file)
    synset_file.close()

  def __load_synsets(self):
    """ Loads synsets from a file.
    Returns:
      A set of the synsets that were loaded successfully. """
    loaded = set([])
    loaded_text = []
    logger.info("Loading synsets...")

    for path in os.listdir(self.__synset_location):
      if path.endswith(".json"):
        # Load synset from file.
        synset_name = path[:-5]
        logger.debug("Loading synset %s." % (synset_name))
        full_path = os.path.join(self.__synset_location, path)

        synset_file = file(full_path)
        loaded_text.append((synset_name, synset_file.read()))
        synset_file.close()

    logger.info("De-jsoning...")
    for synset_name, text in loaded_text:
      # Storing stuff in the default unicode format takes up a
      # massive amount of memory.
      self._synsets[synset_name] = \
          [[elem[0].encode("utf8"), elem[1].encode("utf8")] \
            for elem in json.loads(text)]

      loaded.add(synset_name)

    return loaded

  def _remove_bad_images(self, bad_images):
    """ Removes bad images from the synset files.
    Args:
      bad_images: A list of the images to remove. Should contain tuples of the
      image WNID and URL. """
    for wnid, url in bad_images:
      synset, _ = wnid.split("_")

      # Load the proper file.
      file_name = "%s.json" % (synset)
      file_path = os.path.join(self.__synset_location, file_name)
      file_object = file(file_path)
      synset_data = json.load(file_object)
      file_object.close()

      # Remove the correct entries.
      synset_data.remove([wnid, url])

      # Write back to the file.
      self.__save_synset(synset, synset_data)

  def get_synsets(self):
    return self._synsets

  def get_random_train_batch(self):
    """ Gets a random training batch.
    Returns:
      The array of loaded images, and the list of labels. """
    images, failures = self.__train_set.get_random_batch()

    # Remove failures from the synset files.
    self._remove_bad_images(failures)

    return images

  def get_random_test_batch(self):
    """ Gets a random testing batch.
    Returns:
      The array of loaded images, and the list of labels. """
    images, failures = self.__test_set.get_random_batch()

    # Remove failures from the synset files.
    self._remove_bad_images(failures)

    return images

class _Dataset(object):
  """ Represents a single dataset, with a fixed group of images that it draws
  randomly from. """

  def __init__(self, images, disk_cache, batch_size):
    """
    Args:
      images: The list of images that makes up this dataset. The list should
      contain tuples with each image's WNID and URL.
      disk_cache: The disk cache where we can store downloaded images.
      batch_size: The size of each batch from this dataset. """
    self.__images = dict(zip(range(0, len(images) - 1), images))
    self.__cache = disk_cache

    self.__mem_buffer = cache.MemoryBuffer(224, batch_size, channels=3)
    self.__download_manager = downloader.DownloadManager(200,
        self.__cache, self.__mem_buffer)
    self.__batch_size = batch_size

    logger.info("Have %d total images in database." % (len(self.__images)))

  def __pick_random_image(self):
    """ Picks a random image from our database.
    Returns:
      wnid and url, and key of the image in the images map. """
    # Pick a random image.
    key = random.randint(0, len(self.__images) - 1)
    wnid, url = self.__images[key]

    if wnid in self.__already_picked:
      # This is a duplicate, pick another.
      return self.__pick_random_image()
    self.__already_picked.add(wnid)

    return wnid, url, key

  def get_random_batch(self):
    """ Loads a random batch of images from the whole dataset.
    Returns:
      The array of loaded images, and a list of the synsets each image belongs
      to, as well as a list of all images that failed to download. """
    logger.info("Getting batch.")

    self.__mem_buffer.clear()

    # Keeps track of images that were already used this batch.
    self.__already_picked = set([])

    # Try to download initial images.
    loading = {}
    for _ in range(0, self.__batch_size):
      wnid, key = self.__load_random_image()
      loading[wnid] = key

    # Wait for all the downloads to complete, replacing any ones that fail.
    to_remove = []
    while True:
      failures = self.__download_manager.get_failures()
      if len(failures):
        logger.debug("Replacing %d failed downloads." % (len(failures)))

      # Remove failed images.
      for synset, name, url in failures:
        # Load a new image to replace it.
        wnid, key = self.__load_random_image()
        loading[wnid] = key

        # Remove failed image.
        wnid = "%s_%s" % (synset, name)
        logger.info("Removing bad image %s." % (wnid))
        self.__images.pop(loading[wnid])
        loading.pop(wnid)

        to_remove.append((wnid, url))

      if not self.__download_manager.update():
        break

      time.sleep(0.2)

    return self.__mem_buffer.get_storage(), to_remove

  def __load_random_image(self):
    """ Loads a random image from either the cache or the internet. If loading
    from the internet, it adds it to the download manager, otherwise, it adds
    them to the memory buffer.
    Returns:
      The WNID and key of the loaded image. """
    wnid, url, key = self.__pick_random_image()
    synset, number = wnid.split("_")

    image = self.__get_cached_image(synset, number)
    if image is None:
      # We have to download the image instead.
      self.__download_manager.download_new(synset, number, url)
    else:
      self.__mem_buffer.add(image, number, synset)

    return wnid, key

  def __get_cached_image(self, synset, image_number):
    """ Checks if an image is in the cache, and returns it.
    Args:
      synset: The synset of the image.
      image_number: The image number in the synset.
    Returns:
      The actual image data, or None if the image is not in the cache. """
    # First check in the cache for the image.
    cached_image = self.__cache.get(synset, image_number)
    if cached_image != None:
      # We had a cached copy, so we're done.
      logger.debug("Found image in cache: %s_%s" % (synset, image_number))

      # Select a patch.
      patches = data_augmentation.extract_patches(cached_image)
      cached_image = patches[random.randint(0, len(patches) - 1)]

      return cached_image

    return None

class FilteredImageGetter(ImageGetter):
  """ Works like an ImageGetter, but only loads images from a specific
  pre-defined file containing image names and their URLs. """

  def __init__(self, url_file, batch_size, **kwargs):
    """
    Args:
      url_file: Name of the file containing the images to use.
      batch_size: The size of each batch. """
    self.__url_file = url_file

    super(FilteredImageGetter, self).__init__(None, batch_size, **kwargs)

  def _populate_synsets(self):
    """ Populates the synsets dict. """
    # Load data from the url file.
    urls = file(self.__url_file).read()
    images = _parse_url_file(urls, separator="\t")

    # Separate it by synset and load it into the dictionary.
    for wnid, url in images:
      synset, _ = wnid.split("_")

      if synset not in self._synsets:
        self._synsets[synset] = []
      self._synsets[synset].append([wnid, url])

  def _remove_bad_images(self, *args, **kwargs):
    """ No-op to stop it from trying to access synsets that we don't know about.
    """
    pass
<|endoftext|>


<commit_before>from collections import OrderedDict
import json
import unittest

def vsepr_parse_user_answer(user_input):
    d = OrderedDict(json.loads(user_input))
    d['atoms'] = OrderedDict(sorted(d['atoms'].items()))
    return d

def vsepr_build_correct_answer(geometry, atoms):
    correct_answer = OrderedDict()
    correct_answer['geometry'] = geometry
    correct_answer['atoms'] = OrderedDict(sorted(atoms.items()))
    return correct_answer

def vsepr_grade(user_input, correct_answer, convert_to_perepherial=False):
    """
        Allowed cases:
            c0, a, e
            c0, p
    """
    # print user_input, type(user_input)
    # print correct_answer, type(correct_answer)
    if user_input['geometry'] != correct_answer['geometry']:
        return False

    if user_input['atoms']['c0'] != correct_answer['atoms']['c0']:
        return False

    if convert_to_perepherial:
        # convert user_input from (a,e,e1,e2) to (p)
        # correct_answer must be set in (p) using this flag
        user_input['atoms'] = {'p' + str(i): v for i, v in enumerate(user_input['atoms'].values())}

    for ea_position in ['p', 'a', 'e', 'e1', 'e2']:
        # collecting atoms:
        a_user = [v for k, v in user_input['atoms'].items() if k.startswith(ea_position)]
        a_correct = [v for k, v in correct_answer['atoms'].items() if k.startswith(ea_position)]
        # print a_user, a_correct
        if len(a_user) != len(a_correct):
            return False
        if sorted(a_user) != sorted(a_correct):
            return False

    return True

class Test_Grade(unittest.TestCase):
    ''' test grade function '''

    def test_incorrect_geometry(self):
        correct_answer = vsepr_build_correct_answer(geometry="AX4E0", atoms={"c0": "N", "p0": "H", "p1": "(ep)", "p2": "H", "p3": "H"})
        user_answer = vsepr_parse_user_answer(u'{"geometry":"AX3E0","atoms":{"c0":"B","p0":"F","p1":"B","p2":"F"}}')
        self.assertFalse(vsepr_grade(user_answer, correct_answer))

    def test_incorrect_positions(self):
        correct_answer = vsepr_build_correct_answer(geometry="AX4E0", atoms={"c0": "N", "p0": "H", "p1": "(ep)", "p2": "H", "p3": "H"})
        user_answer = vsepr_parse_user_answer(u'{"geometry":"AX4E0","atoms":{"c0":"B","p0":"F","p1":"B","p2":"F"}}')
        self.assertFalse(vsepr_grade(user_answer, correct_answer))

    def test_correct_answer(self):
        correct_answer = vsepr_build_correct_answer(geometry="AX4E0", atoms={"c0": "N", "p0": "H", "p1": "(ep)", "p2": "H", "p3": "H"})
        user_answer = vsepr_parse_user_answer(u'{"geometry":"AX4E0","atoms":{"c0":"N","p0":"H","p1":"(ep)","p2":"H", "p3":"H"}}')
        self.assertTrue(vsepr_grade(user_answer, correct_answer))

    def test_incorrect_position_order_p(self):
        correct_answer = vsepr_build_correct_answer(geometry="AX4E0", atoms={"c0": "N", "p0": "H", "p1": "(ep)", "p2": "H", "p3": "H"})
        user_answer = vsepr_parse_user_answer(u'{"geometry":"AX4E0","atoms":{"c0":"N","p0":"H","p1":"H","p2":"(ep)", "p3":"H"}}')
        self.assertFalse(vsepr_grade(user_answer, correct_answer))

    def test_correct_position_order_with_ignore_p_order(self):
        correct_answer = vsepr_build_correct_answer(geometry="AX4E0", atoms={"c0": "N", "p0": "H", "p1": "(ep)", "p2": "H", "p3": "H"})
        user_answer = vsepr_parse_user_answer(u'{"geometry":"AX4E0","atoms":{"c0":"N","p0":"H","p1":"H","p2":"(ep)", "p3":"H"}}')
        self.assertTrue(vsepr_grade(user_answer, correct_answer, ignore_p_order=True))

    def test_incorrect_position_order_ae(self):
        correct_answer = vsepr_build_correct_answer(geometry="AX6E0", atoms={"c0": "Br", "a0": "test", "a1": "(ep)", "e0": "H", "e1": "H", "e2": "(ep)", "e3": "(ep)"})
        user_answer = vsepr_parse_user_answer(u'{"geometry":"AX6E0","atoms":{"c0":"Br","a0":"test","a1":"(ep)","e0":"H","e1":"(ep)","e2":"(ep)","e3":"(ep)"}}')
        self.assertFalse(vsepr_grade(user_answer, correct_answer))

    def test_correct_position_order_with_ignore_a_order_not_e(self):
        correct_answer = vsepr_build_correct_answer(geometry="AX6E0", atoms={"c0": "Br", "a0": "(ep)", "a1": "test", "e0": "H", "e1": "H", "e2": "(ep)", "e3": "(ep)"})
        user_answer = vsepr_parse_user_answer(u'{"geometry":"AX6E0","atoms":{"c0":"Br","a0":"test","a1":"(ep)","e0":"H","e1":"H","e2":"(ep)","e3":"(ep)"}}')
        self.assertTrue(vsepr_grade(user_answer, correct_answer, ignore_a_order=True))

    def test_incorrect_position_order_with_ignore_a_order_not_e(self):
        correct_answer = vsepr_build_correct_answer(geometry="AX6E0", atoms={"c0": "Br", "a0": "(ep)", "a1": "test", "e0": "H", "e1": "H", "e2": "H", "e3": "(ep)"})
        user_answer = vsepr_parse_user_answer(u'{"geometry":"AX6E0","atoms":{"c0":"Br","a0":"test","a1":"(ep)","e0":"H","e1":"H","e2":"(ep)","e3":"H"}}')
        self.assertFalse(vsepr_grade(user_answer, correct_answer, ignore_a_order=True))

    def test_correct_position_order_with_ignore_e_order_not_a(self):
        correct_answer = vsepr_build_correct_answer(geometry="AX6E0", atoms={"c0": "Br", "a0": "(ep)", "a1": "test", "e0": "H", "e1": "H", "e2": "H", "e3": "(ep)"})
        user_answer = vsepr_parse_user_answer(u'{"geometry":"AX6E0","atoms":{"c0":"Br","a0":"(ep)","a1":"test","e0":"H","e1":"H","e2":"(ep)","e3":"H"}}')
        self.assertTrue(vsepr_grade(user_answer, correct_answer, ignore_e_order=True))

    def test_incorrect_position_order_with_ignore_e_order__not_a(self):
        correct_answer = vsepr_build_correct_answer(geometry="AX6E0", atoms={"c0": "Br", "a0": "(ep)", "a1": "test", "e0": "H", "e1": "H", "e2": "H", "e3": "(ep)"})
        user_answer = vsepr_parse_user_answer(u'{"geometry":"AX6E0","atoms":{"c0":"Br","a0":"test","a1":"(ep)","e0":"H","e1":"H","e2":"(ep)","e3":"H"}}')
        self.assertFalse(vsepr_grade(user_answer, correct_answer, ignore_e_order=True))

    def test_correct_position_order_with_ignore_ae_order(self):
        correct_answer = vsepr_build_correct_answer(geometry="AX6E0", atoms={"c0": "Br", "a0": "(ep)", "a1": "test", "e0": "H", "e1": "H", "e2": "H", "e3": "(ep)"})
        user_answer = vsepr_parse_user_answer(u'{"geometry":"AX6E0","atoms":{"c0":"Br","a0":"test","a1":"(ep)","e0":"H","e1":"H","e2":"(ep)","e3":"H"}}')
        self.assertTrue(vsepr_grade(user_answer, correct_answer, ignore_e_order=True, ignore_a_order=True))

    def test_incorrect_c0(self):
        correct_answer = vsepr_build_correct_answer(geometry="AX6E0", atoms={"c0": "Br", "a0": "(ep)", "a1": "test", "e0": "H", "e1": "H", "e2": "H", "e3": "(ep)"})
        user_answer = vsepr_parse_user_answer(u'{"geometry":"AX6E0","atoms":{"c0":"H","a0":"test","a1":"(ep)","e0":"H","e1":"H","e2":"(ep)","e3":"H"}}')
        self.assertFalse(vsepr_grade(user_answer, correct_answer, ignore_e_order=True, ignore_a_order=True))

def suite():

    testcases = [Test_Grade]
    suites = []
    for testcase in testcases:
        suites.append(unittest.TestLoader().loadTestsFromTestCase(testcase))
    return unittest.TestSuite(suites)

if __name__ == "__main__":
    unittest.TextTestRunner(verbosity=2).run(suite())
<commit_msg>changed syntax to support e1,e2 in AX6 and add convert_to_p behavior<commit_after>import json
import unittest
import itertools

def vsepr_parse_user_answer(user_input):
    return json.loads(user_input)

def vsepr_build_correct_answer(geometry, atoms):
    return {'geometry': geometry, 'atoms': atoms}

def vsepr_grade(user_input, correct_answer, convert_to_perepherial=False):
    """
        Allowed cases:
            c0, a, e
            c0, p
    """
    # import ipdb; ipdb.set_trace()
    # print user_input, type(user_input)
    # print correct_answer, type(correct_answer)
    if user_input['geometry'] != correct_answer['geometry']:
        return False

    if user_input['atoms']['c0'] != correct_answer['atoms']['c0']:
        return False

    if convert_to_perepherial:
        # convert user_input from (a,e,e1,e2) to (p)
        # correct_answer must be set in (p) using this flag
        c0 = user_input['atoms'].pop('c0')
        user_input['atoms'] = {'p' + str(i): v for i, v in enumerate(user_input['atoms'].values())}
        user_input['atoms']['c0'] = c0

    # special case for AX6
    if 'e10' in correct_answer['atoms']:  # need check e1x, e2x symmetry for AX6..
        a_user = {}
        a_correct = {}
        for ea_position in ['a', 'e1', 'e2']:  # collecting positions:
            a_user[ea_position] = [v for k, v in user_input['atoms'].items() if k.startswith(ea_position)]
            a_correct[ea_position] = [v for k, v in correct_answer['atoms'].items() if k.startswith(ea_position)]

        correct = [sorted(a_correct['a'])] + [sorted(a_correct['e1'])] + [sorted(a_correct['e2'])]
        for permutation in itertools.permutations(['a', 'e1', 'e2']):
            if correct == [sorted(a_user[permutation[0]])] + [sorted(a_user[permutation[1]])] + [sorted(a_user[permutation[2]])]:
                return True
        return False

    else:  # no need to checl e1x,e2x symmetry - convert them to ex
        if 'e10' in user_input['atoms']:  # e1x exists, it is AX6.. case
            e_index = 0
            for k, v in user_input['atoms'].items():
                if len(k) == 3:  # e1x
                    del user_input['atoms'][k]
                    user_input['atoms']['e' + str(e_index)] = v
                    e_index += 1

        # common case
        for ea_position in ['p', 'a', 'e']:
            # collecting atoms:
            a_user = [v for k, v in user_input['atoms'].items() if k.startswith(ea_position)]
            a_correct = [v for k, v in correct_answer['atoms'].items() if k.startswith(ea_position)]
            # print a_user, a_correct
            if len(a_user) != len(a_correct):
                return False
            if sorted(a_user) != sorted(a_correct):
                return False

        return True

class Test_Grade(unittest.TestCase):
    ''' test grade function '''

    def test_incorrect_geometry(self):
        correct_answer = vsepr_build_correct_answer(geometry="AX4E0", atoms={"c0": "N", "p0": "H", "p1": "(ep)", "p2": "H", "p3": "H"})
        user_answer = vsepr_parse_user_answer(u'{"geometry":"AX3E0","atoms":{"c0":"B","p0":"F","p1":"B","p2":"F"}}')
        self.assertFalse(vsepr_grade(user_answer, correct_answer))

    def test_incorrect_positions(self):
        correct_answer = vsepr_build_correct_answer(geometry="AX4E0", atoms={"c0": "N", "p0": "H", "p1": "(ep)", "p2": "H", "p3": "H"})
        user_answer = vsepr_parse_user_answer(u'{"geometry":"AX4E0","atoms":{"c0":"B","p0":"F","p1":"B","p2":"F"}}')
        self.assertFalse(vsepr_grade(user_answer, correct_answer))

    def test_correct_answer(self):
        correct_answer = vsepr_build_correct_answer(geometry="AX4E0", atoms={"c0": "N", "p0": "H", "p1": "(ep)", "p2": "H", "p3": "H"})
        user_answer = vsepr_parse_user_answer(u'{"geometry":"AX4E0","atoms":{"c0":"N","p0":"H","p1":"(ep)","p2":"H", "p3":"H"}}')
        self.assertTrue(vsepr_grade(user_answer, correct_answer))

    def test_incorrect_position_order_p(self):
        correct_answer = vsepr_build_correct_answer(geometry="AX4E0", atoms={"c0": "N", "p0": "H", "p1": "(ep)", "p2": "H", "p3": "H"})
        user_answer = vsepr_parse_user_answer(u'{"geometry":"AX4E0","atoms":{"c0":"N","p0":"H","p1":"H","p2":"(ep)", "p3":"H"}}')
        self.assertFalse(vsepr_grade(user_answer, correct_answer))

    def test_correct_position_order_with_ignore_p_order(self):
        correct_answer = vsepr_build_correct_answer(geometry="AX4E0", atoms={"c0": "N", "p0": "H", "p1": "(ep)", "p2": "H", "p3": "H"})
        user_answer = vsepr_parse_user_answer(u'{"geometry":"AX4E0","atoms":{"c0":"N","p0":"H","p1":"H","p2":"(ep)", "p3":"H"}}')
        self.assertTrue(vsepr_grade(user_answer, correct_answer, ignore_p_order=True))

    def test_incorrect_position_order_ae(self):
        correct_answer = vsepr_build_correct_answer(geometry="AX6E0", atoms={"c0": "Br", "a0": "test", "a1": "(ep)", "e0": "H", "e1": "H", "e2": "(ep)", "e3": "(ep)"})
        user_answer = vsepr_parse_user_answer(u'{"geometry":"AX6E0","atoms":{"c0":"Br","a0":"test","a1":"(ep)","e0":"H","e1":"(ep)","e2":"(ep)","e3":"(ep)"}}')
        self.assertFalse(vsepr_grade(user_answer, correct_answer))

    def test_correct_position_order_with_ignore_a_order_not_e(self):
        correct_answer = vsepr_build_correct_answer(geometry="AX6E0", atoms={"c0": "Br", "a0": "(ep)", "a1": "test", "e0": "H", "e1": "H", "e2": "(ep)", "e3": "(ep)"})
        user_answer = vsepr_parse_user_answer(u'{"geometry":"AX6E0","atoms":{"c0":"Br","a0":"test","a1":"(ep)","e0":"H","e1":"H","e2":"(ep)","e3":"(ep)"}}')
        self.assertTrue(vsepr_grade(user_answer, correct_answer, ignore_a_order=True))

    def test_incorrect_position_order_with_ignore_a_order_not_e(self):
        correct_answer = vsepr_build_correct_answer(geometry="AX6E0", atoms={"c0": "Br", "a0": "(ep)", "a1": "test", "e0": "H", "e1": "H", "e2": "H", "e3": "(ep)"})
        user_answer = vsepr_parse_user_answer(u'{"geometry":"AX6E0","atoms":{"c0":"Br","a0":"test","a1":"(ep)","e0":"H","e1":"H","e2":"(ep)","e3":"H"}}')
        self.assertFalse(vsepr_grade(user_answer, correct_answer, ignore_a_order=True))

    def test_correct_position_order_with_ignore_e_order_not_a(self):
        correct_answer = vsepr_build_correct_answer(geometry="AX6E0", atoms={"c0": "Br", "a0": "(ep)", "a1": "test", "e0": "H", "e1": "H", "e2": "H", "e3": "(ep)"})
        user_answer = vsepr_parse_user_answer(u'{"geometry":"AX6E0","atoms":{"c0":"Br","a0":"(ep)","a1":"test","e0":"H","e1":"H","e2":"(ep)","e3":"H"}}')
        self.assertTrue(vsepr_grade(user_answer, correct_answer, ignore_e_order=True))

    def test_incorrect_position_order_with_ignore_e_order__not_a(self):
        correct_answer = vsepr_build_correct_answer(geometry="AX6E0", atoms={"c0": "Br", "a0": "(ep)", "a1": "test", "e0": "H", "e1": "H", "e2": "H", "e3": "(ep)"})
        user_answer = vsepr_parse_user_answer(u'{"geometry":"AX6E0","atoms":{"c0":"Br","a0":"test","a1":"(ep)","e0":"H","e1":"H","e2":"(ep)","e3":"H"}}')
        self.assertFalse(vsepr_grade(user_answer, correct_answer, ignore_e_order=True))

    def test_correct_position_order_with_ignore_ae_order(self):
        correct_answer = vsepr_build_correct_answer(geometry="AX6E0", atoms={"c0": "Br", "a0": "(ep)", "a1": "test", "e0": "H", "e1": "H", "e2": "H", "e3": "(ep)"})
        user_answer = vsepr_parse_user_answer(u'{"geometry":"AX6E0","atoms":{"c0":"Br","a0":"test","a1":"(ep)","e0":"H","e1":"H","e2":"(ep)","e3":"H"}}')
        self.assertTrue(vsepr_grade(user_answer, correct_answer, ignore_e_order=True, ignore_a_order=True))

    def test_incorrect_c0(self):
        correct_answer = vsepr_build_correct_answer(geometry="AX6E0", atoms={"c0": "Br", "a0": "(ep)", "a1": "test", "e0": "H", "e1": "H", "e2": "H", "e3": "(ep)"})
        user_answer = vsepr_parse_user_answer(u'{"geometry":"AX6E0","atoms":{"c0":"H","a0":"test","a1":"(ep)","e0":"H","e1":"H","e2":"(ep)","e3":"H"}}')
        self.assertFalse(vsepr_grade(user_answer, correct_answer, ignore_e_order=True, ignore_a_order=True))

def suite():

    testcases = [Test_Grade]
    suites = []
    for testcase in testcases:
        suites.append(unittest.TestLoader().loadTestsFromTestCase(testcase))
    return unittest.TestSuite(suites)

if __name__ == "__main__":
    unittest.TextTestRunner(verbosity=2).run(suite())
<|endoftext|>