<commit_before>from __future__ import print_function, absolute_import, unicode_literals

from prompt_toolkit import prompt
from prompt_toolkit.history import FileHistory
from prompt_toolkit.auto_suggest import AutoSuggestFromHistory
from prompt_toolkit.key_binding.defaults import load_key_bindings_for_prompt
from prompt_toolkit.keys import Keys

from kubeshell.style import StyleFactory
from kubeshell.completer import KubectlCompleter
from kubeshell.lexer import KubectlLexer
from kubeshell.toolbar import Toolbar

import os
import click
import sys
import subprocess
import yaml
import logging
logger = logging.getLogger(__name__)

inline_help = True
registry = load_key_bindings_for_prompt()
completer = KubectlCompleter()

class KubeConfig(object):

    clustername = user = ""
    namespace = "default"
    current_context_index = 0
    current_context_name = ""

    @staticmethod
    def parse_kubeconfig():
        if not os.path.exists(os.path.expanduser("~/.kube/config")):
            return ("", "", "")

        with open(os.path.expanduser("~/.kube/config"), "r") as fd:
            docs = yaml.load_all(fd)
            for doc in docs:
                current_context = doc.get("current-context", "")
                contexts = doc.get("contexts")
                if contexts:
                    for index, context in enumerate(contexts):
                        if context['name'] == current_context:
                            KubeConfig.current_context_index = index
                            KubeConfig.current_context_name = context['name']
                            if 'cluster' in context['context']:
                                KubeConfig.clustername = context['context']['cluster']
                            if 'namespace' in context['context']:
                                KubeConfig.namespace = context['context']['namespace']
                            if 'user' in context['context']:
                                KubeConfig.user = context['context']['user']
                            return (KubeConfig.clustername, KubeConfig.user, KubeConfig.namespace)
        return ("", "", "")

    @staticmethod
    def switch_to_next_cluster():
        if not os.path.exists(os.path.expanduser("~/.kube/config")):
            return

        with open(os.path.expanduser("~/.kube/config"), "r") as fd:
            docs = yaml.load_all(fd)
            for doc in docs:
                contexts = doc.get("contexts")
                if contexts:
                    KubeConfig.current_context_index = (KubeConfig.current_context_index+1) % len(contexts)
                    cluster_name = contexts[KubeConfig.current_context_index]['name']
                    kubectl_config_use_context = "kubectl config use-context " + cluster_name
                    cmd_process = subprocess.Popen(kubectl_config_use_context, shell=True, stdout=subprocess.PIPE)
                    cmd_process.wait()
        return

    @staticmethod
    def switch_to_next_namespace(current_namespace):
        namespace_resources = completer.get_resources("namespace")
        namespaces = sorted(res[0] for res in namespace_resources)
        index = (namespaces.index(current_namespace) + 1) % len(namespaces)
        next_namespace = namespaces[index]
        fmt = "kubectl config set-context {} --namespace={}"
        kubectl_config_set_namespace = fmt.format(KubeConfig.current_context_name, next_namespace)
        cmd_process = subprocess.Popen(kubectl_config_set_namespace, shell=True, stdout=subprocess.PIPE)
        cmd_process.wait()

class Kubeshell(object):

    clustername = user = ""
    namespace = "default"

    def __init__(self, refresh_resources=True):
        shell_dir = os.path.expanduser("~/.kube/shell/")
        self.history = FileHistory(os.path.join(shell_dir, "history"))
        if not os.path.exists(shell_dir):
            os.makedirs(shell_dir)
        self.toolbar = Toolbar(self.get_cluster_name, self.get_namespace, self.get_user, self.get_inline_help)

    @registry.add_binding(Keys.F4)
    def _(event):
        try:
            KubeConfig.switch_to_next_cluster()
            Kubeshell.clustername, Kubeshell.user, Kubeshell.namespace = KubeConfig.parse_kubeconfig()
        except Exception as e:
            logger.warning("failed switching clusters", exc_info=1)

    @registry.add_binding(Keys.F5)
    def _(event):
        try:
            KubeConfig.switch_to_next_namespace(Kubeshell.namespace)
            Kubeshell.clustername, Kubeshell.user, Kubeshell.namespace = KubeConfig.parse_kubeconfig()
        except Exception as e:
            logger.warning("failed namespace switching", exc_info=1)

    @registry.add_binding(Keys.F9)
    def _(event):
        global inline_help
        inline_help = not inline_help
        completer.set_inline_help(inline_help)

    @registry.add_binding(Keys.F10)
    def _(event):
        sys.exit()

    def get_cluster_name(self):
        return Kubeshell.clustername

    def get_namespace(self):
        return Kubeshell.namespace

    def get_user(self):
        return Kubeshell.user

    def get_inline_help(self):
        return inline_help

    def run_cli(self):

        def get_title():
            return "kube-shell"

        logger.info("running kube-shell event loop")
        if not os.path.exists(os.path.expanduser("~/.kube/config")):
            click.secho('Kube-shell uses ~/.kube/config for server side completion. Could not find ~/.kube/config. '
                    'Server side completion functionality may not work.', fg='red', blink=True, bold=True)
        while True:
            global inline_help
            try:
                Kubeshell.clustername, Kubeshell.user, Kubeshell.namespace = KubeConfig.parse_kubeconfig()
            except:
                logger.error("unable to parse ~/.kube/config %s", exc_info=1)
            completer.set_namespace(self.namespace)

            try:
                user_input = prompt('kube-shell> ',
                            history=self.history,
                            auto_suggest=AutoSuggestFromHistory(),
                            style=StyleFactory("vim").style,
                            lexer=KubectlLexer,
                            get_title=get_title,
                            enable_history_search=False,
                            get_bottom_toolbar_tokens=self.toolbar.handler,
                            vi_mode=True,
                            key_bindings_registry=registry,
                            completer=completer)
            except (EOFError, KeyboardInterrupt):
                sys.exit()

            if user_input == "clear":
                click.clear()
            elif user_input == "exit":
                sys.exit()

            # if execute shell command then strip "!"
            if user_input.startswith("!"):
                user_input = user_input[1:]

            if user_input:
                if '-o' in user_input and 'json' in user_input:
                    user_input += ' | pygmentize -l json'
                p = subprocess.Popen(user_input, shell=True)
                p.communicate()
<commit_msg>Fix for F5 shortcut bug<commit_after>from __future__ import print_function, absolute_import, unicode_literals

from prompt_toolkit import prompt
from prompt_toolkit.history import FileHistory
from prompt_toolkit.auto_suggest import AutoSuggestFromHistory
from prompt_toolkit.key_binding.defaults import load_key_bindings_for_prompt
from prompt_toolkit.keys import Keys

from kubeshell.style import StyleFactory
from kubeshell.completer import KubectlCompleter
from kubeshell.lexer import KubectlLexer
from kubeshell.toolbar import Toolbar
from kubeshell.client import KubernetesClient

import os
import click
import sys
import subprocess
import yaml
import logging
logger = logging.getLogger(__name__)

inline_help = True
registry = load_key_bindings_for_prompt()
completer = KubectlCompleter()
client = KubernetesClient()

class KubeConfig(object):

    clustername = user = ""
    namespace = "default"
    current_context_index = 0
    current_context_name = ""

    @staticmethod
    def parse_kubeconfig():
        if not os.path.exists(os.path.expanduser("~/.kube/config")):
            return ("", "", "")

        with open(os.path.expanduser("~/.kube/config"), "r") as fd:
            docs = yaml.load_all(fd)
            for doc in docs:
                current_context = doc.get("current-context", "")
                contexts = doc.get("contexts")
                if contexts:
                    for index, context in enumerate(contexts):
                        if context['name'] == current_context:
                            KubeConfig.current_context_index = index
                            KubeConfig.current_context_name = context['name']
                            if 'cluster' in context['context']:
                                KubeConfig.clustername = context['context']['cluster']
                            if 'namespace' in context['context']:
                                KubeConfig.namespace = context['context']['namespace']
                            if 'user' in context['context']:
                                KubeConfig.user = context['context']['user']
                            return (KubeConfig.clustername, KubeConfig.user, KubeConfig.namespace)
        return ("", "", "")

    @staticmethod
    def switch_to_next_cluster():
        if not os.path.exists(os.path.expanduser("~/.kube/config")):
            return

        with open(os.path.expanduser("~/.kube/config"), "r") as fd:
            docs = yaml.load_all(fd)
            for doc in docs:
                contexts = doc.get("contexts")
                if contexts:
                    KubeConfig.current_context_index = (KubeConfig.current_context_index+1) % len(contexts)
                    cluster_name = contexts[KubeConfig.current_context_index]['name']
                    kubectl_config_use_context = "kubectl config use-context " + cluster_name
                    cmd_process = subprocess.Popen(kubectl_config_use_context, shell=True, stdout=subprocess.PIPE)
                    cmd_process.wait()
        return

    @staticmethod
    def switch_to_next_namespace(current_namespace):
        namespace_resources = client.get_resource("namespace")
        namespaces = sorted(res[0] for res in namespace_resources)
        index = (namespaces.index(current_namespace) + 1) % len(namespaces)
        next_namespace = namespaces[index]
        fmt = "kubectl config set-context {} --namespace={}"
        kubectl_config_set_namespace = fmt.format(KubeConfig.current_context_name, next_namespace)
        cmd_process = subprocess.Popen(kubectl_config_set_namespace, shell=True, stdout=subprocess.PIPE)
        cmd_process.wait()

class Kubeshell(object):

    clustername = user = ""
    namespace = "default"

    def __init__(self, refresh_resources=True):
        shell_dir = os.path.expanduser("~/.kube/shell/")
        self.history = FileHistory(os.path.join(shell_dir, "history"))
        if not os.path.exists(shell_dir):
            os.makedirs(shell_dir)
        self.toolbar = Toolbar(self.get_cluster_name, self.get_namespace, self.get_user, self.get_inline_help)

    @registry.add_binding(Keys.F4)
    def _(event):
        try:
            KubeConfig.switch_to_next_cluster()
            Kubeshell.clustername, Kubeshell.user, Kubeshell.namespace = KubeConfig.parse_kubeconfig()
        except Exception as e:
            logger.warning("failed switching clusters", exc_info=1)

    @registry.add_binding(Keys.F5)
    def _(event):
        try:
            KubeConfig.switch_to_next_namespace(Kubeshell.namespace)
            Kubeshell.clustername, Kubeshell.user, Kubeshell.namespace = KubeConfig.parse_kubeconfig()
        except Exception as e:
            logger.warning("failed namespace switching", exc_info=1)

    @registry.add_binding(Keys.F9)
    def _(event):
        global inline_help
        inline_help = not inline_help
        completer.set_inline_help(inline_help)

    @registry.add_binding(Keys.F10)
    def _(event):
        sys.exit()

    def get_cluster_name(self):
        return Kubeshell.clustername

    def get_namespace(self):
        return Kubeshell.namespace

    def get_user(self):
        return Kubeshell.user

    def get_inline_help(self):
        return inline_help

    def run_cli(self):

        def get_title():
            return "kube-shell"

        logger.info("running kube-shell event loop")
        if not os.path.exists(os.path.expanduser("~/.kube/config")):
            click.secho('Kube-shell uses ~/.kube/config for server side completion. Could not find ~/.kube/config. '
                    'Server side completion functionality may not work.', fg='red', blink=True, bold=True)
        while True:
            global inline_help
            try:
                Kubeshell.clustername, Kubeshell.user, Kubeshell.namespace = KubeConfig.parse_kubeconfig()
            except:
                logger.error("unable to parse ~/.kube/config %s", exc_info=1)
            completer.set_namespace(self.namespace)

            try:
                user_input = prompt('kube-shell> ',
                            history=self.history,
                            auto_suggest=AutoSuggestFromHistory(),
                            style=StyleFactory("vim").style,
                            lexer=KubectlLexer,
                            get_title=get_title,
                            enable_history_search=False,
                            get_bottom_toolbar_tokens=self.toolbar.handler,
                            vi_mode=True,
                            key_bindings_registry=registry,
                            completer=completer)
            except (EOFError, KeyboardInterrupt):
                sys.exit()

            if user_input == "clear":
                click.clear()
            elif user_input == "exit":
                sys.exit()

            # if execute shell command then strip "!"
            if user_input.startswith("!"):
                user_input = user_input[1:]

            if user_input:
                if '-o' in user_input and 'json' in user_input:
                    user_input += ' | pygmentize -l json'
                p = subprocess.Popen(user_input, shell=True)
                p.communicate()
<|endoftext|>


<commit_before>"""All functionality related to interactions with AIB online interface."""

import csv
import datetime
import logging
import re
import tempfile
import time

import dateutil.parser as dparser
import mechanicalsoup

def _to_date(text):
    # AIB CSV export format: DD/MM/YYYY
    return dparser.parse(text, dayfirst=True, yearfirst=False)

def _to_value(text):
    tmp = text.strip().replace(',', '')
    if tmp[-3:] == ' DR':
        return '-' + tmp[:-3]
    return tmp

def _csv2account(csv_data, acc):
    transactions = list(csv_data)
    if not transactions:
        return None
    if 'Masked Card Number' in transactions[0]:
        acc['type'] = 'credit'
    else:
        acc['type'] = 'checking'
        acc['balance'] = transactions[-1].get('Balance', 0)
    operations = []
    for transaction in transactions:
        operation = {}
        operation['timestamp'] = _to_date(transaction['Posted Transactions Date'])
        # The mysterious story of 'Description' field in CSV exports continues!
        # Now the columns differ between CC and current account, on top of the
        # implemented bugs :(
        if acc['type'] == 'credit':
            desc = transaction['Description']
            if len(desc) > 18 and desc[18] == ' ':
                desc = desc[:18] + desc[19:]
        else:
            descriptions = []
            for i in [1, 2, 3]:
                descriptions.append(transaction['Description%s' % i].strip())
            desc = ' '.join(filter(bool, descriptions))
        operation['description'] = desc
        operation['debit'] = _to_value(transaction['Debit Amount'])
        operation['credit'] = _to_value(transaction['Credit Amount'])
        operations.append(operation)
    acc['operations'] = operations
    return acc

class CleansingFormatter(logging.Formatter):
    """Logging formatter that scrubs monetary values out."""

    def __init__(self, fmt=None, datefmt=None):
        self.amount_re = re.compile(r'(?:\d+,)*\d+\.\d+(?: DR)?')
        self.date_re = re.compile(r'\d\d/\d\d/\d\d')
        self.description_re = re.compile('(?!dd/mm/yy).+	')
        logging.Formatter.__init__(self, fmt, datefmt)

    def format(self, record):
        tmp = record.msg
        tmp = self.amount_re.sub('X.XX', tmp)
        tmp = self.date_re.sub('dd/mm/yy', tmp)
        tmp = self.description_re.sub('dummy description	', tmp)
        record.msg = tmp
        return logging.Formatter.format(self, record)

class Aib:
    """Automated browser interacting with AIB online interface."""

    def __init__(self, logindata, chatter):
        self.logindata = logindata
        if chatter['debug']:
            # make a directory for debugging output
            debugdir = tempfile.mkdtemp(prefix='aib2ofx_')
            print('WARNING: putting *sensitive* debug data in %s' % debugdir)
            self.logger = logging.getLogger("mechanize")
            logfile = logging.FileHandler(debugdir + '/mechanize.log', 'w')
            formatter = CleansingFormatter('%(asctime)s\n%(message)s')
            logfile.setFormatter(formatter)
            self.logger.addHandler(logfile)
            self.logger.setLevel(logging.DEBUG)
        else:
            logging.disable(logging.DEBUG)
            self.logger = logging.getLogger(None)
        self.browser = mechanicalsoup.StatefulBrowser()
        self.login_done = False
        self.data = {}

    def extract_value(self, varname):
        """Greps the page content for something that looks like a JS assignment, returns value."""
        response = str(self.browser.get_current_page())
        # Luckily the JS code isn't minified.
        regex = re.compile("%s = '([^']+)';" % varname)
        mangled_value = regex.search(response).group(1)
        value = (
            mangled_value.replace('\\/', '/')
            .replace('\\-', '-')
            .encode('latin1')
            .decode('unicode-escape')
        )
        # I feel dirty now. >_<
        return value

    def login(self):
        """Go through the login process."""
        brw = self.browser

        # Entry page.
        self.logger.debug('Requesting first login page.')
        brw.open('https://onlinebanking.aib.ie/inet/roi/login.htm')
        brw.select_form(selector='#loginCiamForm')
        self.logger.debug('Clicking large CONTINUE button on the entry page.')
        # Note: response code will be 401, as we haven't authorized yet.
        response = brw.submit_selected()
        assert response.status_code == 401

        # Redirect page.
        # This redirect is pure javascript, so we need to extract the target URL by hand.
        url = self.extract_value('window.location')
        # We're also saving one value that would be saved in session storage in a modern browser.
        encoded_post_params = self.extract_value('encodedPostParams')
        self.logger.debug('Bouncing through the interstitial.')
        response = brw.open(url, headers={'Referer': response.url})

        # Actual login form.
        brw.select_form()
        brw['pf.username'] = self.logindata['regNumber']
        brw['pf.pass'] = self.logindata['pin']
        self.logger.debug('Submitting login form.')
        brw.submit_selected()

        # Wait for 2FA on phone
        tfa_done = False
        while not tfa_done:
            brw.select_form('#finalizeForm')
            response = brw.submit_selected(update_state=False)
            if response.content == b'approved':
                tfa_done = True
            elif response.content != b'in_progress':
                raise RuntimeError(
                    'unexpected answer during 2FA auth: %s' % response.content
                )
            time.sleep(1)

        # Forward to normal interface.
        brw.select_form('#finalizeForm')
        response = brw.submit_selected()
        # This form is empty after page loads, fields are added by JS.
        form = brw.select_form(nr=0)
        form.new_control('hidden', 'state', self.extract_value('state.value'))
        form.new_control('hidden', 'nonce', self.extract_value('encodedNonce'))
        form.new_control('hidden', 'postParams', encoded_post_params)
        response = brw.submit_selected()
        assert response.status_code == 200
        brw.select_form(nr=0)
        response = brw.submit_selected()
        assert response.status_code == 200

        # mark login as done
        if brw.get_current_page().find(string='My Accounts'):
            self.login_done = True

    def get_data(self):
        """Download data for all accounts."""
        if not self.login_done:
            self.login()

        brw = self.browser
        self.data = {}
        # parse totals
        main_page = brw.get_current_page()
        for account_line in main_page.find_all(
            'button', attrs={'class': 'account-button'}
        ):
            account_name = account_line.find('div', attrs={'class': 'account-name'})
            account_amount = account_line.find('span', attrs={'class': 'a-amount'})
            if not account_name:
                continue

            # Skip pension and saving accounts
            if account_line.find(text=re.compile('SAVINGS')):
                continue

            account = {}
            account['accountId'] = account_name.get_text(strip=True)
            account['available'] = _to_value(account_amount.get_text(strip=True))
            account['currency'] = 'EUR'
            account['bankId'] = 'AIB'
            account['reportDate'] = datetime.datetime.now()

            self.data[account['accountId']] = account

        # parse transactions
        #
        # Note: As of January 2017 there are *two* places that produce CSV data
        # for an account - 'Historical transactions' and 'Recent transactions'.
        # The latter covers shorter period of time, so we use the former. Both
        # exports differ in the type and amount of fields they export. :(
        self.logger.debug('Switching to transaction listing.')
        brw.select_form('#historicalstatement_form_id')
        brw.submit_selected()

        brw.select_form('#hForm')
        accounts_on_page = {
            o.text: o.get('value')
            for o in brw.get_current_form().form.find_all('option')
        }

        for account in list(self.data):
            if account not in accounts_on_page.keys():
                self.logger.debug(
                    'skipping account %s which is absent on historical'
                    'transactions page',
                    account,
                )
                del self.data[account]
                continue

            # get account's page
            self.logger.debug('Requesting transactions for %s.', account)
            brw.select_form('#hForm')
            brw['dsAccountIndex'] = accounts_on_page[account]
            brw.submit_selected()

            # click the export button
            form = brw.select_form('#historicalTransactionsCommand')
            # Some accounts (eg. freshly opened ones) have export facility
            # disabled. Skip them.
            if form.form.find(attrs={'name': 'export'}).get('value') == 'false':
                self.logger.debug(
                    'skipping account %s which has its "Export" button' 'disabled',
                    account,
                )
                del self.data[account]
                continue
            brw.submit_selected()

            # confirm the export request
            brw.select_form('#historicalTransactionsCommand')
            response = brw.submit_selected(update_state=False)
            csv_data = csv.DictReader(
                response.iter_lines(decode_unicode=True), skipinitialspace=True
            )
            self.data[account] = _csv2account(csv_data, self.data[account])

            # go back to the list of accounts
            brw.select_form('#historicaltransactions_form_id')
            brw.submit_selected()

    def getdata(self):
        """Returns data acquired from online interface."""
        return self.data

    def bye(self):
        """Logs user out of bank's online interface."""
        self.logger.debug('Logging out.')
        brw = self.browser
        brw.select_form('#formLogout')
        brw.submit_selected()
        if not brw.get_current_page().find(string='Logged Out'):
            raise Exception('Logout failed!')
<commit_msg>Stop using deprecated apis of mechanicalsoup.<commit_after>"""All functionality related to interactions with AIB online interface."""

import csv
import datetime
import logging
import re
import tempfile
import time

import dateutil.parser as dparser
import mechanicalsoup

def _to_date(text):
    # AIB CSV export format: DD/MM/YYYY
    return dparser.parse(text, dayfirst=True, yearfirst=False)

def _to_value(text):
    tmp = text.strip().replace(',', '')
    if tmp[-3:] == ' DR':
        return '-' + tmp[:-3]
    return tmp

def _csv2account(csv_data, acc):
    transactions = list(csv_data)
    if not transactions:
        return None
    if 'Masked Card Number' in transactions[0]:
        acc['type'] = 'credit'
    else:
        acc['type'] = 'checking'
        acc['balance'] = transactions[-1].get('Balance', 0)
    operations = []
    for transaction in transactions:
        operation = {}
        operation['timestamp'] = _to_date(transaction['Posted Transactions Date'])
        # The mysterious story of 'Description' field in CSV exports continues!
        # Now the columns differ between CC and current account, on top of the
        # implemented bugs :(
        if acc['type'] == 'credit':
            desc = transaction['Description']
            if len(desc) > 18 and desc[18] == ' ':
                desc = desc[:18] + desc[19:]
        else:
            descriptions = []
            for i in [1, 2, 3]:
                descriptions.append(transaction['Description%s' % i].strip())
            desc = ' '.join(filter(bool, descriptions))
        operation['description'] = desc
        operation['debit'] = _to_value(transaction['Debit Amount'])
        operation['credit'] = _to_value(transaction['Credit Amount'])
        operations.append(operation)
    acc['operations'] = operations
    return acc

class CleansingFormatter(logging.Formatter):
    """Logging formatter that scrubs monetary values out."""

    def __init__(self, fmt=None, datefmt=None):
        self.amount_re = re.compile(r'(?:\d+,)*\d+\.\d+(?: DR)?')
        self.date_re = re.compile(r'\d\d/\d\d/\d\d')
        self.description_re = re.compile('(?!dd/mm/yy).+	')
        logging.Formatter.__init__(self, fmt, datefmt)

    def format(self, record):
        tmp = record.msg
        tmp = self.amount_re.sub('X.XX', tmp)
        tmp = self.date_re.sub('dd/mm/yy', tmp)
        tmp = self.description_re.sub('dummy description	', tmp)
        record.msg = tmp
        return logging.Formatter.format(self, record)

class Aib:
    """Automated browser interacting with AIB online interface."""

    def __init__(self, logindata, chatter):
        self.logindata = logindata
        if chatter['debug']:
            # make a directory for debugging output
            debugdir = tempfile.mkdtemp(prefix='aib2ofx_')
            print('WARNING: putting *sensitive* debug data in %s' % debugdir)
            self.logger = logging.getLogger("mechanize")
            logfile = logging.FileHandler(debugdir + '/mechanize.log', 'w')
            formatter = CleansingFormatter('%(asctime)s\n%(message)s')
            logfile.setFormatter(formatter)
            self.logger.addHandler(logfile)
            self.logger.setLevel(logging.DEBUG)
        else:
            logging.disable(logging.DEBUG)
            self.logger = logging.getLogger(None)
        self.browser = mechanicalsoup.StatefulBrowser()
        self.login_done = False
        self.data = {}

    def extract_value(self, varname):
        """Greps the page content for something that looks like a JS assignment, returns value."""
        response = str(self.browser.page)
        # Luckily the JS code isn't minified.
        regex = re.compile("%s = '([^']+)';" % varname)
        mangled_value = regex.search(response).group(1)
        value = (
            mangled_value.replace('\\/', '/')
            .replace('\\-', '-')
            .encode('latin1')
            .decode('unicode-escape')
        )
        # I feel dirty now. >_<
        return value

    def login(self):
        """Go through the login process."""
        brw = self.browser

        # Entry page.
        self.logger.debug('Requesting first login page.')
        brw.open('https://onlinebanking.aib.ie/inet/roi/login.htm')
        brw.select_form(selector='#loginCiamForm')
        self.logger.debug('Clicking large CONTINUE button on the entry page.')
        # Note: response code will be 401, as we haven't authorized yet.
        response = brw.submit_selected()
        assert response.status_code == 401

        # Redirect page.
        # This redirect is pure javascript, so we need to extract the target URL by hand.
        url = self.extract_value('window.location')
        # We're also saving one value that would be saved in session storage in a modern browser.
        encoded_post_params = self.extract_value('encodedPostParams')
        self.logger.debug('Bouncing through the interstitial.')
        response = brw.open(url, headers={'Referer': response.url})

        # Actual login form.
        brw.select_form()
        brw['pf.username'] = self.logindata['regNumber']
        brw['pf.pass'] = self.logindata['pin']
        self.logger.debug('Submitting login form.')
        brw.submit_selected()

        # Wait for 2FA on phone
        tfa_done = False
        while not tfa_done:
            brw.select_form('#finalizeForm')
            response = brw.submit_selected(update_state=False)
            if response.content == b'approved':
                tfa_done = True
            elif response.content != b'in_progress':
                raise RuntimeError(
                    'unexpected answer during 2FA auth: %s' % response.content
                )
            time.sleep(1)

        # Forward to normal interface.
        brw.select_form('#finalizeForm')
        response = brw.submit_selected()
        # This form is empty after page loads, fields are added by JS.
        form = brw.select_form(nr=0)
        form.new_control('hidden', 'state', self.extract_value('state.value'))
        form.new_control('hidden', 'nonce', self.extract_value('encodedNonce'))
        form.new_control('hidden', 'postParams', encoded_post_params)
        response = brw.submit_selected()
        assert response.status_code == 200
        brw.select_form(nr=0)
        response = brw.submit_selected()
        assert response.status_code == 200

        # mark login as done
        if brw.page.find(string='My Accounts'):
            self.login_done = True

    def get_data(self):
        """Download data for all accounts."""
        if not self.login_done:
            self.login()

        brw = self.browser
        self.data = {}
        # parse totals
        main_page = brw.page
        for account_line in main_page.find_all(
            'button', attrs={'class': 'account-button'}
        ):
            account_name = account_line.find('div', attrs={'class': 'account-name'})
            account_amount = account_line.find('span', attrs={'class': 'a-amount'})
            if not account_name:
                continue

            # Skip pension and saving accounts
            if account_line.find(text=re.compile('SAVINGS')):
                continue

            account = {}
            account['accountId'] = account_name.get_text(strip=True)
            account['available'] = _to_value(account_amount.get_text(strip=True))
            account['currency'] = 'EUR'
            account['bankId'] = 'AIB'
            account['reportDate'] = datetime.datetime.now()

            self.data[account['accountId']] = account

        # parse transactions
        #
        # Note: As of January 2017 there are *two* places that produce CSV data
        # for an account - 'Historical transactions' and 'Recent transactions'.
        # The latter covers shorter period of time, so we use the former. Both
        # exports differ in the type and amount of fields they export. :(
        self.logger.debug('Switching to transaction listing.')
        brw.select_form('#historicalstatement_form_id')
        brw.submit_selected()

        brw.select_form('#hForm')
        accounts_on_page = {
            o.text: o.get('value') for o in brw.form.form.find_all('option')
        }

        for account in list(self.data):
            if account not in accounts_on_page.keys():
                self.logger.debug(
                    'skipping account %s which is absent on historical'
                    'transactions page',
                    account,
                )
                del self.data[account]
                continue

            # get account's page
            self.logger.debug('Requesting transactions for %s.', account)
            brw.select_form('#hForm')
            brw['dsAccountIndex'] = accounts_on_page[account]
            brw.submit_selected()

            # click the export button
            form = brw.select_form('#historicalTransactionsCommand')
            # Some accounts (eg. freshly opened ones) have export facility
            # disabled. Skip them.
            if form.form.find(attrs={'name': 'export'}).get('value') == 'false':
                self.logger.debug(
                    'skipping account %s which has its "Export" button' 'disabled',
                    account,
                )
                del self.data[account]
                continue
            brw.submit_selected()

            # confirm the export request
            brw.select_form('#historicalTransactionsCommand')
            response = brw.submit_selected(update_state=False)
            csv_data = csv.DictReader(
                response.iter_lines(decode_unicode=True), skipinitialspace=True
            )
            self.data[account] = _csv2account(csv_data, self.data[account])

            # go back to the list of accounts
            brw.select_form('#historicaltransactions_form_id')
            brw.submit_selected()

    def getdata(self):
        """Returns data acquired from online interface."""
        return self.data

    def bye(self):
        """Logs user out of bank's online interface."""
        self.logger.debug('Logging out.')
        brw = self.browser
        brw.select_form('#formLogout')
        brw.submit_selected()
        if not brw.page.find(string='Logged Out'):
            raise Exception('Logout failed!')
<|endoftext|>


<commit_before>import os
import re

import django.db.models.options as options
from django.db import models
from django.core.exceptions import ObjectDoesNotExist
from django.core.paginator import Paginator, EmptyPage, PageNotAnInteger

from wagtail.wagtailcore.fields import RichTextField
from wagtail.wagtailcore.models import Page
from wagtail.wagtailsearch import index

from modelcluster.fields import ParentalKey
from modelcluster.tags import ClusterTaggableManager
from taggit.models import TaggedItemBase, Tag
from bs4 import BeautifulSoup

from core.utilities import validate_only_one_instance
from core.panels import *
from core.snippets import *
from core.forms import *

options.DEFAULT_NAMES = options.DEFAULT_NAMES + ('description',)

class IndexPage(models.Model):

    """
    Abstract Index Page class. Declare a couple of abstract methods that should be implemented by
    any class implementing this 'interface'.
    """
    # Just one instance allowed

    def clean(self):
        validate_only_one_instance(self)

    def children(self):
        raise NotImplementedError("Class %s doesn't implement aMethod()" % (self.__class__.__name__))

    def get_context(self, request):
        raise NotImplementedError("Class %s doesn't implement aMethod()" % (self.__class__.__name__))

    class Meta:
        abstract = True

class HomePage(Page, IndexPage):

    """
    HomePage class, inheriting from wagtailcore.Page straight away
    """
    subpage_types = [
        'core.WagtailPage',
        'core.CompanyIndex',
        'core.WagtailCompanyPage',
        'core.WagtailSitePage',
        'core.SubmitFormPage'
    ]
    feed_image = models.ForeignKey(
        'wagtailimages.Image',
        null=True,
        blank=True,
        on_delete=models.SET_NULL,
        related_name='+'
    )
    search_fields = ()

    body = RichTextField(blank=True)

    @property
    def og_image(self):
        # Returns image and image type of feed_image, if exists
        image = {'image': None, 'type': None}
        if self.feed_image:
            image['image'] = self.feed_image
        name, extension = os.path.splitext(image['image'].file.url)
        image['type'] = extension[1:]
        return image

    def children(self):
        return self.get_children().live()

    def get_context(self, request):
        # Get pages
        pages = WagtailSitePage.objects.live().descendant_of(self).order_by('-is_featured', '-latest_revision_created_at')

        # Filter by tag
        tag = request.GET.get('tag')
        if tag:
            pages = pages.filter(tags__name__iexact=tag)

        # Pagination
        page = request.GET.get('page')
        paginator = Paginator(pages, 12)  # Show 12 pages per page
        try:
            pages = paginator.page(page)
        except PageNotAnInteger:
            pages = paginator.page(1)
        except EmptyPage:
            pages = paginator.page(paginator.num_pages)

        # Update template context
        context = super(HomePage, self).get_context(request)
        context['pages'] = pages
        context['tag'] = tag
        # Only tags used by live pages
        context['tags'] = Tag.objects.filter(
            core_pagetag_items__isnull=False,
            core_pagetag_items__content_object__live=True
        ).distinct().order_by('name')

        return context

    def __unicode__(self):
        return self.title

    class Meta:
        verbose_name = "Home page"
        description = "Where the good stuff happens!"

HomePage.content_panels = HOME_PAGE_CONTENT_PANELS
HomePage.promote_panels = WAGTAIL_PAGE_PROMOTE_PANELS

class CompanyIndex(Page, IndexPage):

    """
    HomePage class, inheriting from wagtailcore.Page straight away
    """
    subpage_types = ['core.WagtailCompanyPage']
    search_fields = ()
    body = RichTextField(null=True, blank=True)
    show_map = models.BooleanField(default=False, help_text='Show map of companies around the world.')

    def children(self):
        return self.get_children().live()

    def get_context(self, request):
        # Get pages
        pages = WagtailCompanyPage.objects.live().descendant_of(self).distinct().order_by('-numchild', '-latest_revision_created_at')

        # Filter by tag
        tag = request.GET.get('tag')
        if tag:
            pages = pages.filter(tags__name__iexact=tag)

        # Pagination
        page = request.GET.get('page')
        paginator = Paginator(pages, 12)  # Show 12 pages per page
        try:
            pages = paginator.page(page)
        except PageNotAnInteger:
            pages = paginator.page(1)
        except EmptyPage:
            pages = paginator.page(paginator.num_pages)

        # Update template context
        context = super(CompanyIndex, self).get_context(request)
        context['pages'] = pages
        context['tag'] = tag
        return context

    def __unicode__(self):
        return self.title

    class Meta:
        verbose_name = "Company Pages"
        description = "Companies index"

CompanyIndex.content_panels = WAGTAIL_COMPANY_INDEX_PAGE_CONTENT_PANELS

class PageTag(TaggedItemBase):
    content_object = ParentalKey('core.WagtailPage', related_name='tagged_items')

# Main core Page model. All main content pages inherit from this class.
class WagtailPage(Page):

    """
    Our main custom Page class. All content pages should inherit from this one.
    """
    feed_image = models.ForeignKey(
        'wagtailimages.Image',
        null=True,
        blank=True,
        on_delete=models.SET_NULL,
        related_name='+'
    )
    body = RichTextField(blank=True)
    tags = ClusterTaggableManager(through=PageTag, blank=True)
    search_fields = ()

    @property
    def parent(self):
        try:
            return self.get_ancestors().reverse()[0]
        except:
            return None

    @property
    def child(self):
        for related_object in self._meta.get_all_related_objects():
            if not issubclass(related_object.model, self.__class__):
                continue
            try:
                return getattr(self, related_object.get_accessor_name())
            except ObjectDoesNotExist:
                pass

    @property
    def body_text(self):
        return BeautifulSoup(self.body, "html5lib").get_text()

    @property
    def body_excerpt(self):
        """
        Return body text replacing end of lines (. ? ! chars) with a blank space
        """
        return re.sub(r'([\.?!])([a-zA-Z])', r'\1 \2', self.body_text)

    @property
    def og_image(self):
        # Returns image and image type of feed_image or image as fallback, if exists
        image = {'image': None, 'type': None}
        if self.feed_image:
            image['image'] = self.feed_image
        name, extension = os.path.splitext(image['image'].file.url)
        image['type'] = extension[1:]
        return image

WagtailPage.content_panels = WAGTAIL_PAGE_CONTENT_PANELS
WagtailPage.promote_panels = WAGTAIL_PAGE_PROMOTE_PANELS

class WagtailCompanyPage(WagtailPage):
    """
    Company page listing a bunch of site pages
    """
    parent_types = ['core.HomePage']
    subpage_types = ['core.WagtailSitePage']

    company_url = models.URLField(
        blank=True,
        null=True,
        help_text='Paste the URL of your site, something like "https://www.springload.co.nz"',
    )
    github_url = models.URLField(null=True, blank=True)
    twitter_url = models.URLField(null=True, blank=True)
    location = models.CharField(max_length=128, blank=True, null=True)
    show_map = models.BooleanField(default=True, help_text='Show company in the map of companies around the world.')
    coords = models.CharField(max_length=255, blank=True, null=True)

    logo = models.ForeignKey(
        'wagtailimages.Image',
        null=True,
        blank=True,
        on_delete=models.SET_NULL,
        related_name='+'
    )

    search_fields = Page.search_fields + (  # Inherit search_fields from Page
        index.SearchField('company_url', boost=1),
        index.SearchField('body_text', boost=1)
    )

    @property
    def lat(self):
        if self.coords:
            return self.coords.split(",")[0].strip()
        else:
            return None

    @property
    def lon(self):
        if self.coords:
            return self.coords.split(",")[1].strip()
        else:
            return None

    @property
    def twitter_handler(self):
        if self.twitter_url:
            return "@%s" % self.twitter_url.split("/")[-1]
        else:
            return None

    @property
    def github_user(self):
        if self.github_url:
            return self.github_url.split("/")[-1]
        else:
            return None

    @property
    def children_count(self):
        return self.children().count()

    @property
    def og_image(self):
        # Returns image and image type of logo or feed_image as fallback, if exists
        image = {'image': None, 'type': None}
        if self.logo:
            image['image'] = self.logo
        elif self.feed_image:
            image['image'] = self.feed_image
        name, extension = os.path.splitext(image['image'].file.url)
        image['type'] = extension[1:]
        return image

    def children(self):
        return WagtailSitePage.objects.live().descendant_of(self).order_by('-first_published_at')

    def get_context(self, request):
        # Get pages
        pages = self.children()
        # Pagination
        page = request.GET.get('page')
        paginator = Paginator(pages, 12)  # Show 12 pages per page
        try:
            pages = paginator.page(page)
        except PageNotAnInteger:
            pages = paginator.page(1)
        except EmptyPage:
            pages = paginator.page(paginator.num_pages)

        # Update template context
        context = super(WagtailCompanyPage, self).get_context(request)
        context['pages'] = pages
        return context

    def __unicode__(self):
        return self.title

    class Meta:
        verbose_name = "Company page"
        description = "Page for companies developing Wagtail"

WagtailCompanyPage.content_panels = WAGTAIL_COMPANY_PAGE_CONTENT_PANELS

class WagtailSitePage(WagtailPage):
    """
    Site page
    """
    parent_types = ['core.WagtailCompanyPage', 'core.HomePage']
    subpage_types = []
    is_featured = models.BooleanField(
        "Featured",
        default=False,
        blank=False,
        help_text='''If enabled, this site will appear on top of the sites list of the homepage.'''
    )
    image_desktop = models.ForeignKey(
        'wagtailimages.Image',
        null=True,
        blank=True,
        on_delete=models.SET_NULL,
        related_name='+'
    )
    image_tablet = models.ForeignKey(
        'wagtailimages.Image',
        null=True,
        blank=True,
        on_delete=models.SET_NULL,
        related_name='+'
    )
    image_phone = models.ForeignKey(
        'wagtailimages.Image',
        null=True,
        blank=True,
        on_delete=models.SET_NULL,
        related_name='+'
    )
    site_url = models.URLField(
        blank=False,
        null=True,
        help_text='Paste the URL of your site, something like "http://www.springload.co.nz"',
    )

    search_fields = Page.search_fields + (  # Inherit search_fields from Page
        index.SearchField('site_url'),
        index.SearchField('body_text')
    )

    @property
    def og_image(self):
        # Returns image and image type of feed_image, if exists
        image = {'image': None, 'type': None}
        if self.feed_image:
            image['image'] = self.feed_image
        elif self.image_desktop:
            image['image'] = self.image_desktop
        name, extension = os.path.splitext(image['image'].file.url)
        image['type'] = extension[1:]
        return image

    def __unicode__(self):
        return self.title

    class Meta:
        verbose_name = "Site page"
        description = "Page to show case an existing site based on Wagtail"

WagtailSitePage.content_panels = WAGTAIL_SITE_PAGE_CONTENT_PANELS
WagtailSitePage.promote_panels = WAGTAIL_SITE_PAGE_PROMOTE_PANELS
<commit_msg>Fix incorrect social handles for urls with trailing slash<commit_after>import os
import re

import django.db.models.options as options
from django.db import models
from django.core.exceptions import ObjectDoesNotExist
from django.core.paginator import Paginator, EmptyPage, PageNotAnInteger

from wagtail.wagtailcore.fields import RichTextField
from wagtail.wagtailcore.models import Page
from wagtail.wagtailsearch import index

from modelcluster.fields import ParentalKey
from modelcluster.tags import ClusterTaggableManager
from taggit.models import TaggedItemBase, Tag
from bs4 import BeautifulSoup

from core.utilities import validate_only_one_instance
from core.panels import *
from core.snippets import *
from core.forms import *

options.DEFAULT_NAMES = options.DEFAULT_NAMES + ('description',)

class IndexPage(models.Model):

    """
    Abstract Index Page class. Declare a couple of abstract methods that should be implemented by
    any class implementing this 'interface'.
    """
    # Just one instance allowed

    def clean(self):
        validate_only_one_instance(self)

    def children(self):
        raise NotImplementedError("Class %s doesn't implement aMethod()" % (self.__class__.__name__))

    def get_context(self, request):
        raise NotImplementedError("Class %s doesn't implement aMethod()" % (self.__class__.__name__))

    class Meta:
        abstract = True

class HomePage(Page, IndexPage):

    """
    HomePage class, inheriting from wagtailcore.Page straight away
    """
    subpage_types = [
        'core.WagtailPage',
        'core.CompanyIndex',
        'core.WagtailCompanyPage',
        'core.WagtailSitePage',
        'core.SubmitFormPage'
    ]
    feed_image = models.ForeignKey(
        'wagtailimages.Image',
        null=True,
        blank=True,
        on_delete=models.SET_NULL,
        related_name='+'
    )
    search_fields = ()

    body = RichTextField(blank=True)

    @property
    def og_image(self):
        # Returns image and image type of feed_image, if exists
        image = {'image': None, 'type': None}
        if self.feed_image:
            image['image'] = self.feed_image
        name, extension = os.path.splitext(image['image'].file.url)
        image['type'] = extension[1:]
        return image

    def children(self):
        return self.get_children().live()

    def get_context(self, request):
        # Get pages
        pages = WagtailSitePage.objects.live().descendant_of(self).order_by('-is_featured', '-latest_revision_created_at')

        # Filter by tag
        tag = request.GET.get('tag')
        if tag:
            pages = pages.filter(tags__name__iexact=tag)

        # Pagination
        page = request.GET.get('page')
        paginator = Paginator(pages, 12)  # Show 12 pages per page
        try:
            pages = paginator.page(page)
        except PageNotAnInteger:
            pages = paginator.page(1)
        except EmptyPage:
            pages = paginator.page(paginator.num_pages)

        # Update template context
        context = super(HomePage, self).get_context(request)
        context['pages'] = pages
        context['tag'] = tag
        # Only tags used by live pages
        context['tags'] = Tag.objects.filter(
            core_pagetag_items__isnull=False,
            core_pagetag_items__content_object__live=True
        ).distinct().order_by('name')

        return context

    def __unicode__(self):
        return self.title

    class Meta:
        verbose_name = "Home page"
        description = "Where the good stuff happens!"

HomePage.content_panels = HOME_PAGE_CONTENT_PANELS
HomePage.promote_panels = WAGTAIL_PAGE_PROMOTE_PANELS

class CompanyIndex(Page, IndexPage):

    """
    HomePage class, inheriting from wagtailcore.Page straight away
    """
    subpage_types = ['core.WagtailCompanyPage']
    search_fields = ()
    body = RichTextField(null=True, blank=True)
    show_map = models.BooleanField(default=False, help_text='Show map of companies around the world.')

    def children(self):
        return self.get_children().live()

    def get_context(self, request):
        # Get pages
        pages = WagtailCompanyPage.objects.live().descendant_of(self).distinct().order_by('-numchild', '-latest_revision_created_at')

        # Filter by tag
        tag = request.GET.get('tag')
        if tag:
            pages = pages.filter(tags__name__iexact=tag)

        # Pagination
        page = request.GET.get('page')
        paginator = Paginator(pages, 12)  # Show 12 pages per page
        try:
            pages = paginator.page(page)
        except PageNotAnInteger:
            pages = paginator.page(1)
        except EmptyPage:
            pages = paginator.page(paginator.num_pages)

        # Update template context
        context = super(CompanyIndex, self).get_context(request)
        context['pages'] = pages
        context['tag'] = tag
        return context

    def __unicode__(self):
        return self.title

    class Meta:
        verbose_name = "Company Pages"
        description = "Companies index"

CompanyIndex.content_panels = WAGTAIL_COMPANY_INDEX_PAGE_CONTENT_PANELS

class PageTag(TaggedItemBase):
    content_object = ParentalKey('core.WagtailPage', related_name='tagged_items')

# Main core Page model. All main content pages inherit from this class.
class WagtailPage(Page):

    """
    Our main custom Page class. All content pages should inherit from this one.
    """
    feed_image = models.ForeignKey(
        'wagtailimages.Image',
        null=True,
        blank=True,
        on_delete=models.SET_NULL,
        related_name='+'
    )
    body = RichTextField(blank=True)
    tags = ClusterTaggableManager(through=PageTag, blank=True)
    search_fields = ()

    @property
    def parent(self):
        try:
            return self.get_ancestors().reverse()[0]
        except:
            return None

    @property
    def child(self):
        for related_object in self._meta.get_all_related_objects():
            if not issubclass(related_object.model, self.__class__):
                continue
            try:
                return getattr(self, related_object.get_accessor_name())
            except ObjectDoesNotExist:
                pass

    @property
    def body_text(self):
        return BeautifulSoup(self.body, "html5lib").get_text()

    @property
    def body_excerpt(self):
        """
        Return body text replacing end of lines (. ? ! chars) with a blank space
        """
        return re.sub(r'([\.?!])([a-zA-Z])', r'\1 \2', self.body_text)

    @property
    def og_image(self):
        # Returns image and image type of feed_image or image as fallback, if exists
        image = {'image': None, 'type': None}
        if self.feed_image:
            image['image'] = self.feed_image
        name, extension = os.path.splitext(image['image'].file.url)
        image['type'] = extension[1:]
        return image

WagtailPage.content_panels = WAGTAIL_PAGE_CONTENT_PANELS
WagtailPage.promote_panels = WAGTAIL_PAGE_PROMOTE_PANELS

class WagtailCompanyPage(WagtailPage):
    """
    Company page listing a bunch of site pages
    """
    parent_types = ['core.HomePage']
    subpage_types = ['core.WagtailSitePage']

    company_url = models.URLField(
        blank=True,
        null=True,
        help_text='Paste the URL of your site, something like "https://www.springload.co.nz"',
    )
    github_url = models.URLField(null=True, blank=True)
    twitter_url = models.URLField(null=True, blank=True)
    location = models.CharField(max_length=128, blank=True, null=True)
    show_map = models.BooleanField(default=True, help_text='Show company in the map of companies around the world.')
    coords = models.CharField(max_length=255, blank=True, null=True)

    logo = models.ForeignKey(
        'wagtailimages.Image',
        null=True,
        blank=True,
        on_delete=models.SET_NULL,
        related_name='+'
    )

    search_fields = Page.search_fields + (  # Inherit search_fields from Page
        index.SearchField('company_url', boost=1),
        index.SearchField('body_text', boost=1)
    )

    @property
    def lat(self):
        if self.coords:
            return self.coords.split(",")[0].strip()
        else:
            return None

    @property
    def lon(self):
        if self.coords:
            return self.coords.split(",")[1].strip()
        else:
            return None

    @property
    def twitter_handler(self):
        if self.twitter_url:
            return "@%s" % self.twitter_url.strip('/ ').split("/")[-1]
        else:
            return None

    @property
    def github_user(self):
        if self.github_url:
            return self.github_url.strip('/ ').split("/")[-1]
        else:
            return None

    @property
    def children_count(self):
        return self.children().count()

    @property
    def og_image(self):
        # Returns image and image type of logo or feed_image as fallback, if exists
        image = {'image': None, 'type': None}
        if self.logo:
            image['image'] = self.logo
        elif self.feed_image:
            image['image'] = self.feed_image
        name, extension = os.path.splitext(image['image'].file.url)
        image['type'] = extension[1:]
        return image

    def children(self):
        return WagtailSitePage.objects.live().descendant_of(self).order_by('-first_published_at')

    def get_context(self, request):
        # Get pages
        pages = self.children()
        # Pagination
        page = request.GET.get('page')
        paginator = Paginator(pages, 12)  # Show 12 pages per page
        try:
            pages = paginator.page(page)
        except PageNotAnInteger:
            pages = paginator.page(1)
        except EmptyPage:
            pages = paginator.page(paginator.num_pages)

        # Update template context
        context = super(WagtailCompanyPage, self).get_context(request)
        context['pages'] = pages
        return context

    def __unicode__(self):
        return self.title

    class Meta:
        verbose_name = "Company page"
        description = "Page for companies developing Wagtail"

WagtailCompanyPage.content_panels = WAGTAIL_COMPANY_PAGE_CONTENT_PANELS

class WagtailSitePage(WagtailPage):
    """
    Site page
    """
    parent_types = ['core.WagtailCompanyPage', 'core.HomePage']
    subpage_types = []
    is_featured = models.BooleanField(
        "Featured",
        default=False,
        blank=False,
        help_text='''If enabled, this site will appear on top of the sites list of the homepage.'''
    )
    image_desktop = models.ForeignKey(
        'wagtailimages.Image',
        null=True,
        blank=True,
        on_delete=models.SET_NULL,
        related_name='+'
    )
    image_tablet = models.ForeignKey(
        'wagtailimages.Image',
        null=True,
        blank=True,
        on_delete=models.SET_NULL,
        related_name='+'
    )
    image_phone = models.ForeignKey(
        'wagtailimages.Image',
        null=True,
        blank=True,
        on_delete=models.SET_NULL,
        related_name='+'
    )
    site_url = models.URLField(
        blank=False,
        null=True,
        help_text='Paste the URL of your site, something like "http://www.springload.co.nz"',
    )

    search_fields = Page.search_fields + (  # Inherit search_fields from Page
        index.SearchField('site_url'),
        index.SearchField('body_text')
    )

    @property
    def og_image(self):
        # Returns image and image type of feed_image, if exists
        image = {'image': None, 'type': None}
        if self.feed_image:
            image['image'] = self.feed_image
        elif self.image_desktop:
            image['image'] = self.image_desktop
        name, extension = os.path.splitext(image['image'].file.url)
        image['type'] = extension[1:]
        return image

    def __unicode__(self):
        return self.title

    class Meta:
        verbose_name = "Site page"
        description = "Page to show case an existing site based on Wagtail"

WagtailSitePage.content_panels = WAGTAIL_SITE_PAGE_CONTENT_PANELS
WagtailSitePage.promote_panels = WAGTAIL_SITE_PAGE_PROMOTE_PANELS
<|endoftext|>


<commit_before># coding=utf-8
# Copyright 2018 The Google AI Language Team Authors.
#
# 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.
"""Model utility functions."""

from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

from tensor2tensor.layers import common_layers

import tensorflow as tf

__all__ = [
    "GNMTAttentionMultiCell",
    "gnmt_residual_fn",
    "create_rnn_cell",
    "create_gnmt_rnn_cell",
    "build_unidirectional_rnn",
    "build_bidirectional_rnn",
    "make_sequences_compatible",
    "get_embeddings",
    "build_logits",
    "get_global_step",
]

def _single_cell(unit_type,
                 num_units,
                 forget_bias,
                 dropout,
                 mode,
                 residual_connection=False,
                 residual_fn=None,
                 trainable=True):
  """Create an instance of a single RNN cell."""
  # dropout (= 1 - keep_prob) is set to 0 during eval and infer
  dropout = dropout if mode == tf.estimator.ModeKeys.TRAIN else 0.0

  # Cell Type
  if unit_type == "lstm":
    single_cell = tf.contrib.rnn.LSTMCell(
        num_units,
        forget_bias=forget_bias,
        trainable=trainable)
  elif unit_type == "gru":
    single_cell = tf.contrib.rnn.GRUCell(
        num_units,
        trainable=trainable)
  elif unit_type == "layer_norm_lstm":
    single_cell = tf.contrib.rnn.LayerNormBasicLSTMCell(
        num_units,
        forget_bias=forget_bias,
        layer_norm=True,
        trainable=trainable)
  elif unit_type == "nas":
    single_cell = tf.contrib.rnn.NASCell(
        num_units,
        trainable=trainable)
  else:
    raise ValueError("Unknown unit type %s!" % unit_type)

  # Dropout (= 1 - keep_prob).
  if dropout > 0.0:
    single_cell = tf.contrib.rnn.DropoutWrapper(
        cell=single_cell, input_keep_prob=(1.0 - dropout))

  # Residual.
  if residual_connection:
    single_cell = tf.contrib.rnn.ResidualWrapper(
        single_cell, residual_fn=residual_fn)

  return single_cell

def _cell_list(unit_type,
               num_units,
               num_layers,
               num_residual_layers,
               forget_bias,
               dropout,
               mode,
               single_cell_fn=None,
               residual_fn=None,
               trainable=True):
  """Create a list of RNN cells."""
  if single_cell_fn is None:
    single_cell_fn = _single_cell

  cell_list = []
  for i in range(num_layers):
    single_cell = single_cell_fn(
        unit_type=unit_type,
        num_units=num_units,
        forget_bias=forget_bias,
        dropout=dropout,
        mode=mode,
        residual_connection=(i >= num_layers - num_residual_layers),
        residual_fn=residual_fn,
        trainable=trainable)
    cell_list.append(single_cell)

  return cell_list

class GNMTAttentionMultiCell(tf.nn.rnn_cell.MultiRNNCell):
  """A MultiCell with GNMT attention style."""

  def __init__(self, attention_cell, cells, use_new_attention=False):
    """Creates a GNMTAttentionMultiCell.

    Args:
      attention_cell: An instance of AttentionWrapper.
      cells: A list of RNNCell wrapped with AttentionInputWrapper.
      use_new_attention: Whether to use the attention generated from current
        step bottom layer's output. Default is False.
    """
    cells = [attention_cell] + cells
    self.use_new_attention = use_new_attention
    super(GNMTAttentionMultiCell, self).__init__(cells, state_is_tuple=True)

  def __call__(self, inputs, state, scope=None):
    """Run the cell with bottom layer's attention copied to all upper layers."""
    if not tf.contrib.framework.nest.is_sequence(state):
      raise ValueError(
          "Expected state to be a tuple of length %d, but received: %s"
          % (len(self.state_size), state))

    scope = "multi_rnn_cell" if scope is None else scope
    with tf.variable_scope(scope):
      new_states = []

      with tf.variable_scope("cell_0_attention"):
        attention_cell = self._cells[0]
        attention_state = state[0]
        cur_inp, new_attention_state = attention_cell(inputs, attention_state)
        new_states.append(new_attention_state)

      for i in range(1, len(self._cells)):
        with tf.variable_scope("cell_%d" % i):

          cell = self._cells[i]
          cur_state = state[i]

          if self.use_new_attention:
            cur_inp = tf.concat([cur_inp, new_attention_state.attention], -1)
          else:
            cur_inp = tf.concat([cur_inp, attention_state.attention], -1)

          cur_inp, new_state = cell(cur_inp, cur_state)
          new_states.append(new_state)

    return cur_inp, tuple(new_states)

def gnmt_residual_fn(inputs, outputs):
  """Residual function that handles different inputs and outputs inner dims.

  Args:
    inputs: A potentially nested structure of <tensor> [..., input_dim] that
      represents cell inputs.
    outputs: A potentially nested structure of <tensor> [..., output_dim] that
      represents cell outputs. Must have the same structure and number of
      dimensions as inputs and output_dim >= input_dim must hold.

  Returns:
    outputs + actual_inputs where actual_inputs are a nested structure of slices
    of the inputs along the last dimension up to the output_dim.
  """
  def split_input(inp, out):
    out_dim = out.get_shape().as_list()[-1]
    inp_dim = inp.get_shape().as_list()[-1]
    return tf.split(inp, [out_dim, inp_dim - out_dim], axis=-1)
  actual_inputs, _ = tf.contrib.framework.nest.map_structure(
      split_input, inputs, outputs)
  def assert_shape_match(inp, out):
    inp.get_shape().assert_is_compatible_with(out.get_shape())
  tf.contrib.framework.nest.assert_same_structure(actual_inputs, outputs)
  tf.contrib.framework.nest.map_structure(
      assert_shape_match, actual_inputs, outputs)
  return tf.contrib.framework.nest.map_structure(
      lambda inp, out: inp + out, actual_inputs, outputs)

def create_rnn_cell(unit_type,
                    num_units,
                    num_layers,
                    num_residual_layers,
                    forget_bias,
                    dropout,
                    mode,
                    attention_mechanism=None,
                    attention_num_heads=1,
                    attention_layer_size=None,
                    output_attention=False,
                    single_cell_fn=None,
                    trainable=True):
  """Returns an instance of an RNN cell.

  Args:
    unit_type: A string that specifies the type of the recurrent unit. Must be
      one of {"lstm", "gru", "lstm_norm", "nas"}.
    num_units: An integer for the numner of units per layer.
    num_layers: An integer for the number of recurrent layers.
    num_residual_layers: An integer for the number of residual layers.
    forget_bias: A float for the forget bias in LSTM cells.
    dropout: A float for the recurrent dropout rate.
    mode: TRAIN | EVAL | PREDICT
    attention_mechanism: An instance of tf.contrib.seq2seq.AttentionMechanism.
    attention_num_heads: An integer for the number of attention heads.
    attention_layer_size: Optional integer for the size of the attention layer.
    output_attention: A boolean indicating whether RNN cell outputs attention.
    single_cell_fn: A function for building a single RNN cell.
    trainable: A boolean indicating whether the cell weights are trainable.

  Returns:
    An RNNCell instance.
  """
  cell_list = _cell_list(
      unit_type=unit_type,
      num_units=num_units,
      num_layers=num_layers,
      forget_bias=forget_bias,
      dropout=dropout,
      mode=mode,
      num_residual_layers=num_residual_layers,
      single_cell_fn=single_cell_fn,
      trainable=trainable)

  if len(cell_list) == 1:  # Single layer.
    cell = cell_list[0]
  else:  # Multiple layers.
    cell = tf.contrib.rnn.MultiRNNCell(cell_list)

  # Wrap with attention, if necessary.
  if attention_mechanism is not None:
    cell = tf.contrib.seq2seq.AttentionWrapper(
        cell, [attention_mechanism] * attention_num_heads,
        attention_layer_size=[attention_layer_size] * attention_num_heads,
        alignment_history=False,
        output_attention=output_attention,
        name="attention")

  return cell

def create_gnmt_rnn_cell(unit_type,
                         num_units,
                         num_layers,
                         num_residual_layers,
                         forget_bias,
                         dropout,
                         mode,
                         attention_mechanism,
                         attention_num_heads=1,
                         attention_layer_size=None,
                         output_attention=False,
                         single_cell_fn=None):
  """Returns an instance of an GNMT-style RNN cell.

  Args:
    unit_type: A string that specifies the type of the recurrent unit. Must be
      one of {"lstm", "gru", "lstm_norm", "nas"}.
    num_units: An integer for the numner of units per layer.
    num_layers: An integer for the number of recurrent layers.
    num_residual_layers: An integer for the number of residual layers.
    forget_bias: A float for the forget bias in LSTM cells.
    dropout: A float for the recurrent dropout rate.
    mode: TRAIN | EVAL | PREDICT
    attention_mechanism: An instance of tf.contrib.seq2seq.AttentionMechanism.
    attention_num_heads: An integer for the number of attention heads.
    attention_layer_size: Optional integer for the size of the attention layer.
    output_attention: A boolean indicating whether RNN cell outputs attention.
    single_cell_fn: A function for building a single RNN cell.

  Returns:
    An RNNCell instance.
  """
  cell_list = _cell_list(
      unit_type=unit_type,
      num_units=num_units,
      num_layers=num_layers,
      forget_bias=forget_bias,
      dropout=dropout,
      mode=mode,
      num_residual_layers=num_residual_layers,
      single_cell_fn=single_cell_fn,
      residual_fn=gnmt_residual_fn)

  if attention_num_heads > 1:
    attention_mechanism = [attention_mechanism] * attention_num_heads
    attention_layer_size = [attention_layer_size] * attention_num_heads

  # Only wrap the bottom layer with the attention mechanism.
  attention_cell = cell_list.pop(0)
  attention_cell = tf.contrib.seq2seq.AttentionWrapper(
      attention_cell, attention_mechanism,
      attention_layer_size=attention_layer_size,
      alignment_history=False,
      output_attention=output_attention,
      name="gnmt_attention")

  cell = GNMTAttentionMultiCell(
      attention_cell, cell_list, use_new_attention=True)

  return cell

def build_unidirectional_rnn(sequences, length,
                             num_layers, num_residual_layers, num_units,
                             unit_type, forget_bias, dropout, mode,
                             trainable=True):
  """Builds sequences encoded with a unidirectional RNN."""
  cell = create_rnn_cell(
      unit_type=unit_type,
      num_units=num_units,
      num_layers=num_layers,
      num_residual_layers=num_residual_layers,
      forget_bias=forget_bias,
      dropout=dropout,
      mode=mode,
      trainable=trainable)

  # Encode.
  outputs, final_state = tf.nn.dynamic_rnn(
      cell, sequences,
      sequence_length=length,
      dtype=tf.float32,
      time_major=False)

  return outputs, final_state

def build_bidirectional_rnn(sequences, length,
                            num_layers, num_residual_layers, num_units,
                            unit_type, forget_bias, dropout, mode,
                            trainable=True):
  """Builds sequences encoded with a bidirectional RNN."""
  # Create RNN cells.
  cells = {
      direction: create_rnn_cell(
          unit_type=unit_type,
          num_units=num_units,
          num_layers=num_layers,
          num_residual_layers=num_residual_layers,
          forget_bias=forget_bias,
          dropout=dropout,
          mode=mode,
          trainable=trainable)
      for direction in ["fw", "bw"]}

  # Encode.
  outputs, state = tf.nn.bidirectional_dynamic_rnn(
      cells["fw"], cells["bw"], sequences,
      sequence_length=length,
      dtype=tf.float32,
      time_major=False)

  # Concatenate bidirectional outputs.
  outputs = tf.concat(outputs, axis=2)

  return outputs, state

def make_sequences_compatible(seq1, seq2, op_type="max_pool"):
  """Preprocesses sequences to make them compatible for computing a loss.

  Args:
    seq1: <type> [batch_size, seq_len, ...].
    seq2: <type> [batch_size, seq_len, ...].
    op_type: String for the operation type. Must be one of
      {"avg_pool", "max_pool", "truncate"}.

  Returns:
    new_seq1, new_seq2: <type> [batch_size, ...] for preprocessed sequences.

  Raises:
    ValueError: if op_type is unknown.
  """
  if op_type not in {"avg_pool", "max_pool", "truncate"}:
    raise ValueError("Unknown op_type: %s" % op_type)

  if op_type == "avg_pool":
    new_seq1 = tf.reduce_mean(seq1, axis=1)
    new_seq2 = tf.reduce_mean(seq2, axis=1)
  elif op_type == "max_pool":
    new_seq1 = tf.reduce_max(seq1, axis=1)
    new_seq2 = tf.reduce_max(seq2, axis=1)
  elif op_type == "truncate":
    max_length = tf.minimum(tf.shape(seq1)[1], tf.shape(seq2)[1])
    new_seq1 = seq1[:, :max_length, :]
    new_seq2 = seq2[:, :max_length, :]

  return  new_seq1, new_seq2

def get_embeddings(modality, outer_scope, inner_scope="shared"):
  """Returns embeddings for the given modality stricly forcing reuse.

  Args:
    modality: A tensor2tensro.utils.modality.Modality object.
    outer_scope: A variable scope object used as the outer scope.
    inner_scope: A string used as the inner-most variable scope name.

  Returns:
    embeddings: <float32> [vocab_size, emb_size].
  """
  with tf.variable_scope(outer_scope, reuse=True):
    with tf.variable_scope(inner_scope, reuse=True):
      embeddings = modality._get_weights()  # pylint: disable=protected-access
  return embeddings

def build_logits(sequences, embeddings, vocab_size):
  """Builds logits for sequences and a given modality."""
  sequences_shape = common_layers.shape_list(sequences)
  sequences = tf.reshape(sequences, [-1, sequences_shape[-1]])
  # Compute logits.
  logits = tf.matmul(sequences, embeddings, transpose_b=True)
  logits_shape = sequences_shape[:-1] + [1, vocab_size]
  return tf.reshape(logits, logits_shape)

def get_global_step(hparams):
  """Returns the global optimization step."""
  step = tf.to_float(tf.train.get_or_create_global_step())
  multiplier = hparams.optimizer_multistep_accumulate_steps
  if not multiplier:
    return step
  return step / tf.to_float(multiplier)
<commit_msg>Fix typo<commit_after># coding=utf-8
# Copyright 2018 The Google AI Language Team Authors.
#
# 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.
"""Model utility functions."""

from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

from tensor2tensor.layers import common_layers

import tensorflow as tf

__all__ = [
    "GNMTAttentionMultiCell",
    "gnmt_residual_fn",
    "create_rnn_cell",
    "create_gnmt_rnn_cell",
    "build_unidirectional_rnn",
    "build_bidirectional_rnn",
    "make_sequences_compatible",
    "get_embeddings",
    "build_logits",
    "get_global_step",
]

def _single_cell(unit_type,
                 num_units,
                 forget_bias,
                 dropout,
                 mode,
                 residual_connection=False,
                 residual_fn=None,
                 trainable=True):
  """Create an instance of a single RNN cell."""
  # dropout (= 1 - keep_prob) is set to 0 during eval and infer
  dropout = dropout if mode == tf.estimator.ModeKeys.TRAIN else 0.0

  # Cell Type
  if unit_type == "lstm":
    single_cell = tf.contrib.rnn.LSTMCell(
        num_units,
        forget_bias=forget_bias,
        trainable=trainable)
  elif unit_type == "gru":
    single_cell = tf.contrib.rnn.GRUCell(
        num_units,
        trainable=trainable)
  elif unit_type == "layer_norm_lstm":
    single_cell = tf.contrib.rnn.LayerNormBasicLSTMCell(
        num_units,
        forget_bias=forget_bias,
        layer_norm=True,
        trainable=trainable)
  elif unit_type == "nas":
    single_cell = tf.contrib.rnn.NASCell(
        num_units,
        trainable=trainable)
  else:
    raise ValueError("Unknown unit type %s!" % unit_type)

  # Dropout (= 1 - keep_prob).
  if dropout > 0.0:
    single_cell = tf.contrib.rnn.DropoutWrapper(
        cell=single_cell, input_keep_prob=(1.0 - dropout))

  # Residual.
  if residual_connection:
    single_cell = tf.contrib.rnn.ResidualWrapper(
        single_cell, residual_fn=residual_fn)

  return single_cell

def _cell_list(unit_type,
               num_units,
               num_layers,
               num_residual_layers,
               forget_bias,
               dropout,
               mode,
               single_cell_fn=None,
               residual_fn=None,
               trainable=True):
  """Create a list of RNN cells."""
  if single_cell_fn is None:
    single_cell_fn = _single_cell

  cell_list = []
  for i in range(num_layers):
    single_cell = single_cell_fn(
        unit_type=unit_type,
        num_units=num_units,
        forget_bias=forget_bias,
        dropout=dropout,
        mode=mode,
        residual_connection=(i >= num_layers - num_residual_layers),
        residual_fn=residual_fn,
        trainable=trainable)
    cell_list.append(single_cell)

  return cell_list

class GNMTAttentionMultiCell(tf.nn.rnn_cell.MultiRNNCell):
  """A MultiCell with GNMT attention style."""

  def __init__(self, attention_cell, cells, use_new_attention=False):
    """Creates a GNMTAttentionMultiCell.

    Args:
      attention_cell: An instance of AttentionWrapper.
      cells: A list of RNNCell wrapped with AttentionInputWrapper.
      use_new_attention: Whether to use the attention generated from current
        step bottom layer's output. Default is False.
    """
    cells = [attention_cell] + cells
    self.use_new_attention = use_new_attention
    super(GNMTAttentionMultiCell, self).__init__(cells, state_is_tuple=True)

  def __call__(self, inputs, state, scope=None):
    """Run the cell with bottom layer's attention copied to all upper layers."""
    if not tf.contrib.framework.nest.is_sequence(state):
      raise ValueError(
          "Expected state to be a tuple of length %d, but received: %s"
          % (len(self.state_size), state))

    scope = "multi_rnn_cell" if scope is None else scope
    with tf.variable_scope(scope):
      new_states = []

      with tf.variable_scope("cell_0_attention"):
        attention_cell = self._cells[0]
        attention_state = state[0]
        cur_inp, new_attention_state = attention_cell(inputs, attention_state)
        new_states.append(new_attention_state)

      for i in range(1, len(self._cells)):
        with tf.variable_scope("cell_%d" % i):

          cell = self._cells[i]
          cur_state = state[i]

          if self.use_new_attention:
            cur_inp = tf.concat([cur_inp, new_attention_state.attention], -1)
          else:
            cur_inp = tf.concat([cur_inp, attention_state.attention], -1)

          cur_inp, new_state = cell(cur_inp, cur_state)
          new_states.append(new_state)

    return cur_inp, tuple(new_states)

def gnmt_residual_fn(inputs, outputs):
  """Residual function that handles different inputs and outputs inner dims.

  Args:
    inputs: A potentially nested structure of <tensor> [..., input_dim] that
      represents cell inputs.
    outputs: A potentially nested structure of <tensor> [..., output_dim] that
      represents cell outputs. Must have the same structure and number of
      dimensions as inputs and output_dim >= input_dim must hold.

  Returns:
    outputs + actual_inputs where actual_inputs are a nested structure of slices
    of the inputs along the last dimension up to the output_dim.
  """
  def split_input(inp, out):
    out_dim = out.get_shape().as_list()[-1]
    inp_dim = inp.get_shape().as_list()[-1]
    return tf.split(inp, [out_dim, inp_dim - out_dim], axis=-1)
  actual_inputs, _ = tf.contrib.framework.nest.map_structure(
      split_input, inputs, outputs)
  def assert_shape_match(inp, out):
    inp.get_shape().assert_is_compatible_with(out.get_shape())
  tf.contrib.framework.nest.assert_same_structure(actual_inputs, outputs)
  tf.contrib.framework.nest.map_structure(
      assert_shape_match, actual_inputs, outputs)
  return tf.contrib.framework.nest.map_structure(
      lambda inp, out: inp + out, actual_inputs, outputs)

def create_rnn_cell(unit_type,
                    num_units,
                    num_layers,
                    num_residual_layers,
                    forget_bias,
                    dropout,
                    mode,
                    attention_mechanism=None,
                    attention_num_heads=1,
                    attention_layer_size=None,
                    output_attention=False,
                    single_cell_fn=None,
                    trainable=True):
  """Returns an instance of an RNN cell.

  Args:
    unit_type: A string that specifies the type of the recurrent unit. Must be
      one of {"lstm", "gru", "lstm_norm", "nas"}.
    num_units: An integer for the numner of units per layer.
    num_layers: An integer for the number of recurrent layers.
    num_residual_layers: An integer for the number of residual layers.
    forget_bias: A float for the forget bias in LSTM cells.
    dropout: A float for the recurrent dropout rate.
    mode: TRAIN | EVAL | PREDICT
    attention_mechanism: An instance of tf.contrib.seq2seq.AttentionMechanism.
    attention_num_heads: An integer for the number of attention heads.
    attention_layer_size: Optional integer for the size of the attention layer.
    output_attention: A boolean indicating whether RNN cell outputs attention.
    single_cell_fn: A function for building a single RNN cell.
    trainable: A boolean indicating whether the cell weights are trainable.

  Returns:
    An RNNCell instance.
  """
  cell_list = _cell_list(
      unit_type=unit_type,
      num_units=num_units,
      num_layers=num_layers,
      forget_bias=forget_bias,
      dropout=dropout,
      mode=mode,
      num_residual_layers=num_residual_layers,
      single_cell_fn=single_cell_fn,
      trainable=trainable)

  if len(cell_list) == 1:  # Single layer.
    cell = cell_list[0]
  else:  # Multiple layers.
    cell = tf.contrib.rnn.MultiRNNCell(cell_list)

  # Wrap with attention, if necessary.
  if attention_mechanism is not None:
    cell = tf.contrib.seq2seq.AttentionWrapper(
        cell, [attention_mechanism] * attention_num_heads,
        attention_layer_size=[attention_layer_size] * attention_num_heads,
        alignment_history=False,
        output_attention=output_attention,
        name="attention")

  return cell

def create_gnmt_rnn_cell(unit_type,
                         num_units,
                         num_layers,
                         num_residual_layers,
                         forget_bias,
                         dropout,
                         mode,
                         attention_mechanism,
                         attention_num_heads=1,
                         attention_layer_size=None,
                         output_attention=False,
                         single_cell_fn=None):
  """Returns an instance of an GNMT-style RNN cell.

  Args:
    unit_type: A string that specifies the type of the recurrent unit. Must be
      one of {"lstm", "gru", "lstm_norm", "nas"}.
    num_units: An integer for the numner of units per layer.
    num_layers: An integer for the number of recurrent layers.
    num_residual_layers: An integer for the number of residual layers.
    forget_bias: A float for the forget bias in LSTM cells.
    dropout: A float for the recurrent dropout rate.
    mode: TRAIN | EVAL | PREDICT
    attention_mechanism: An instance of tf.contrib.seq2seq.AttentionMechanism.
    attention_num_heads: An integer for the number of attention heads.
    attention_layer_size: Optional integer for the size of the attention layer.
    output_attention: A boolean indicating whether RNN cell outputs attention.
    single_cell_fn: A function for building a single RNN cell.

  Returns:
    An RNNCell instance.
  """
  cell_list = _cell_list(
      unit_type=unit_type,
      num_units=num_units,
      num_layers=num_layers,
      forget_bias=forget_bias,
      dropout=dropout,
      mode=mode,
      num_residual_layers=num_residual_layers,
      single_cell_fn=single_cell_fn,
      residual_fn=gnmt_residual_fn)

  if attention_num_heads > 1:
    attention_mechanism = [attention_mechanism] * attention_num_heads
    attention_layer_size = [attention_layer_size] * attention_num_heads

  # Only wrap the bottom layer with the attention mechanism.
  attention_cell = cell_list.pop(0)
  attention_cell = tf.contrib.seq2seq.AttentionWrapper(
      attention_cell, attention_mechanism,
      attention_layer_size=attention_layer_size,
      alignment_history=False,
      output_attention=output_attention,
      name="gnmt_attention")

  cell = GNMTAttentionMultiCell(
      attention_cell, cell_list, use_new_attention=True)

  return cell

def build_unidirectional_rnn(sequences, length,
                             num_layers, num_residual_layers, num_units,
                             unit_type, forget_bias, dropout, mode,
                             trainable=True):
  """Builds sequences encoded with a unidirectional RNN."""
  cell = create_rnn_cell(
      unit_type=unit_type,
      num_units=num_units,
      num_layers=num_layers,
      num_residual_layers=num_residual_layers,
      forget_bias=forget_bias,
      dropout=dropout,
      mode=mode,
      trainable=trainable)

  # Encode.
  outputs, final_state = tf.nn.dynamic_rnn(
      cell, sequences,
      sequence_length=length,
      dtype=tf.float32,
      time_major=False)

  return outputs, final_state

def build_bidirectional_rnn(sequences, length,
                            num_layers, num_residual_layers, num_units,
                            unit_type, forget_bias, dropout, mode,
                            trainable=True):
  """Builds sequences encoded with a bidirectional RNN."""
  # Create RNN cells.
  cells = {
      direction: create_rnn_cell(
          unit_type=unit_type,
          num_units=num_units,
          num_layers=num_layers,
          num_residual_layers=num_residual_layers,
          forget_bias=forget_bias,
          dropout=dropout,
          mode=mode,
          trainable=trainable)
      for direction in ["fw", "bw"]}

  # Encode.
  outputs, state = tf.nn.bidirectional_dynamic_rnn(
      cells["fw"], cells["bw"], sequences,
      sequence_length=length,
      dtype=tf.float32,
      time_major=False)

  # Concatenate bidirectional outputs.
  outputs = tf.concat(outputs, axis=2)

  return outputs, state

def make_sequences_compatible(seq1, seq2, op_type="max_pool"):
  """Preprocesses sequences to make them compatible for computing a loss.

  Args:
    seq1: <type> [batch_size, seq_len, ...].
    seq2: <type> [batch_size, seq_len, ...].
    op_type: String for the operation type. Must be one of
      {"avg_pool", "max_pool", "truncate"}.

  Returns:
    new_seq1, new_seq2: <type> [batch_size, ...] for preprocessed sequences.

  Raises:
    ValueError: if op_type is unknown.
  """
  if op_type not in {"avg_pool", "max_pool", "truncate"}:
    raise ValueError("Unknown op_type: %s" % op_type)

  if op_type == "avg_pool":
    new_seq1 = tf.reduce_mean(seq1, axis=1)
    new_seq2 = tf.reduce_mean(seq2, axis=1)
  elif op_type == "max_pool":
    new_seq1 = tf.reduce_max(seq1, axis=1)
    new_seq2 = tf.reduce_max(seq2, axis=1)
  elif op_type == "truncate":
    max_length = tf.minimum(tf.shape(seq1)[1], tf.shape(seq2)[1])
    new_seq1 = seq1[:, :max_length, :]
    new_seq2 = seq2[:, :max_length, :]

  return  new_seq1, new_seq2

def get_embeddings(modality, outer_scope, inner_scope="shared"):
  """Returns embeddings for the given modality stricly forcing reuse.

  Args:
    modality: A tensor2tensor.utils.modality.Modality object.
    outer_scope: A variable scope object used as the outer scope.
    inner_scope: A string used as the inner-most variable scope name.

  Returns:
    embeddings: <float32> [vocab_size, emb_size].
  """
  with tf.variable_scope(outer_scope, reuse=True):
    with tf.variable_scope(inner_scope, reuse=True):
      embeddings = modality._get_weights()  # pylint: disable=protected-access
  return embeddings

def build_logits(sequences, embeddings, vocab_size):
  """Builds logits for sequences and a given modality."""
  sequences_shape = common_layers.shape_list(sequences)
  sequences = tf.reshape(sequences, [-1, sequences_shape[-1]])
  # Compute logits.
  logits = tf.matmul(sequences, embeddings, transpose_b=True)
  logits_shape = sequences_shape[:-1] + [1, vocab_size]
  return tf.reshape(logits, logits_shape)

def get_global_step(hparams):
  """Returns the global optimization step."""
  step = tf.to_float(tf.train.get_or_create_global_step())
  multiplier = hparams.optimizer_multistep_accumulate_steps
  if not multiplier:
    return step
  return step / tf.to_float(multiplier)
<|endoftext|>