<commit_before>from selenium import webdriver
from selenium.webdriver.common.keys import Keys
from django.test import LiveServerTestCase
import os
import time
os.environ['DJANGO_LIVE_TEST_SERVER_ADDRESS'] = 'localhost:8082'

from django.test.utils import override_settings
from selenium.webdriver.firefox.webdriver import WebDriver

class MySeleniumTests(LiveServerTestCase):
    fixtures = ['initial_data.json']

    @classmethod
    def setUpClass(cls):
        cls.selenium = WebDriver()
        super(MySeleniumTests, cls).setUpClass()

    @classmethod
    def tearDownClass(cls):
        cls.selenium.quit()
        super(MySeleniumTests, cls).tearDownClass()

    @override_settings(DEBUG=True)
    def test_signup(self):
        self.selenium.get('%s%s' % (self.live_server_url, '/accounts/signup/'))
        self.selenium.find_element_by_name("username").send_keys('bugbugbug')
        self.selenium.find_element_by_name("email").send_keys('bugbugbug@bugbug.com')
        self.selenium.find_element_by_name("password1").send_keys('secret123')
        self.selenium.find_element_by_name("password2").send_keys('secret123')
        self.selenium.find_element_by_xpath('//*[@id="signup_form"]/button').click()
        body = self.selenium.find_element_by_tag_name('body')
        self.assertIn(u'bugbugbug (0 pts)', body.text)

    @override_settings(DEBUG=True)
    def test_login(self):
        self.selenium.get('%s%s' % (self.live_server_url, '/accounts/login/'))
        self.selenium.find_element_by_name("login").send_keys('bugbug')
        self.selenium.find_element_by_name("password").send_keys('secret')
        self.selenium.find_element_by_xpath('//*[@id="page-wrapper"]/div/form/button').click()
        body = self.selenium.find_element_by_tag_name('body')
        self.assertIn(u'bugbug (0 pts)', body.text)

    @override_settings(DEBUG=True)
    def test_post_bug(self):
        self.selenium.set_page_load_timeout(70)
        self.selenium.get('%s%s' % (self.live_server_url, '/accounts/login/'))
        self.selenium.find_element_by_name("login").send_keys('bugbug')
        self.selenium.find_element_by_name("password").send_keys('secret')
        self.selenium.find_element_by_xpath('//*[@id="page-wrapper"]/div/form/button').click()
        self.selenium.find_element_by_name("url").send_keys('http://www.example.com/')
        self.selenium.find_element_by_id("description").send_keys('Description of bug')
        Imagepath=os.path.abspath('.\\website\\static\img\logo.jpg')
        self.selenium.find_element_by_name("screenshot").send_keys(Imagepath)
        self.selenium.find_element_by_xpath('//*[@id="page-wrapper"]/div/div[2]/div[1]/form/div[4]/button').click()
        body = self.selenium.find_element_by_tag_name('body')
        self.assertIn(u'Description of bug', body.text)<commit_msg>Updated tests.py for the new screenshot button.<commit_after>from selenium import webdriver
from selenium.webdriver.common.keys import Keys
from django.test import LiveServerTestCase
import os
import time
os.environ['DJANGO_LIVE_TEST_SERVER_ADDRESS'] = 'localhost:8082'

from django.test.utils import override_settings
from selenium.webdriver.firefox.webdriver import WebDriver

class MySeleniumTests(LiveServerTestCase):
    fixtures = ['initial_data.json']

    @classmethod
    def setUpClass(cls):
        cls.selenium = WebDriver()
        super(MySeleniumTests, cls).setUpClass()

    @classmethod
    def tearDownClass(cls):
        cls.selenium.quit()
        super(MySeleniumTests, cls).tearDownClass()

    @override_settings(DEBUG=True)
    def test_signup(self):
        self.selenium.get('%s%s' % (self.live_server_url, '/accounts/signup/'))
        self.selenium.find_element_by_name("username").send_keys('bugbugbug')
        self.selenium.find_element_by_name("email").send_keys('bugbugbug@bugbug.com')
        self.selenium.find_element_by_name("password1").send_keys('secret123')
        self.selenium.find_element_by_name("password2").send_keys('secret123')
        self.selenium.find_element_by_xpath('//*[@id="signup_form"]/button').click()
        body = self.selenium.find_element_by_tag_name('body')
        self.assertIn(u'bugbugbug (0 pts)', body.text)

    @override_settings(DEBUG=True)
    def test_login(self):
        self.selenium.get('%s%s' % (self.live_server_url, '/accounts/login/'))
        self.selenium.find_element_by_name("login").send_keys('bugbug')
        self.selenium.find_element_by_name("password").send_keys('secret')
        self.selenium.find_element_by_xpath('//*[@id="page-wrapper"]/div/form/button').click()
        body = self.selenium.find_element_by_tag_name('body')
        self.assertIn(u'bugbug (0 pts)', body.text)

    @override_settings(DEBUG=True)
    def test_post_bug(self):
        self.selenium.set_page_load_timeout(70)
        self.selenium.get('%s%s' % (self.live_server_url, '/accounts/login/'))
        self.selenium.find_element_by_name("login").send_keys('bugbug')
        self.selenium.find_element_by_name("password").send_keys('secret')
        self.selenium.find_element_by_xpath('//*[@id="page-wrapper"]/div/form/button').click()
        self.selenium.find_element_by_name("url").send_keys('http://www.example.com/')
        self.selenium.find_element_by_id("description").send_keys('Description of bug')
        Imagepath=os.path.abspath(os.path.join(os.getcwd(),'website/static/img/logo.jpg'))
        self.selenium.find_element_by_name("screenshot").send_keys(Imagepath)
        self.selenium.find_element_by_xpath('//*[@id="page-wrapper"]/div/div[2]/div[1]/form/div[4]/button').click()
        body = self.selenium.find_element_by_tag_name('body')
        self.assertIn(u'Description of bug', body.text)<|endoftext|>


<commit_before># -*- coding: utf-8 -*-
import logging
import itertools
import httplib as http
from flask import request, redirect
from modularodm import Q

from framework.exceptions import HTTPError
from framework.forms import utils
from framework.routing import proxy_url
from framework.auth import Auth, get_current_user
from framework.auth.decorators import collect_auth, must_be_logged_in
from framework.auth.forms import (RegistrationForm, SignInForm,
                                  ForgotPasswordForm, ResetPasswordForm)

from website.models import Guid
from website.util import web_url_for
from website.project.forms import NewProjectForm
from website.project import model
from website import settings

logger = logging.getLogger('website.views')

def _rescale_ratio(nodes):
    """Get scaling denominator for log lists across a sequence of nodes.

    :param nodes: Nodes
    :return: Max number of logs

    """
    if not nodes:
        return 0
    user = get_current_user()
    counts = [
        len(node.logs)
        for node in nodes
        if node.can_view(Auth(user=user))
    ]
    if counts:
        return float(max(counts))
    return 0.0

def _render_node(node):
    """

    :param node:
    :return:

    """
    return {
        'id': node._primary_key,
        'url': node.url,
        'api_url': node.api_url,
        'primary': node.primary,
    }

def _render_nodes(nodes):
    """

    :param nodes:
    :return:
    """
    ret = {
        'nodes': [
            _render_node(node)
            for node in nodes
        ],
        'rescale_ratio': _rescale_ratio(nodes),
    }
    return ret

def _get_user_activity(node, user, rescale_ratio):

    # Counters
    total_count = len(node.logs)

    # Note: It's typically much faster to find logs of a given node
    # attached to a given user using node.logs.find(...) than by
    # loading the logs into Python and checking each one. However,
    # using deep caching might be even faster down the road.

    ua_count = node.logs.find(Q('user', 'eq', user)).count()
    non_ua_count = total_count - ua_count # base length of blue bar

    # Normalize over all nodes
    ua = ua_count / rescale_ratio * settings.USER_ACTIVITY_MAX_WIDTH
    non_ua = non_ua_count / rescale_ratio * settings.USER_ACTIVITY_MAX_WIDTH

    return ua_count, ua, non_ua

@collect_auth
def index(auth, **kwargs):
    """Redirect to dashboard if user is logged in, else show homepage.

    """
    if auth.user:
        return redirect(web_url_for('dashboard'))
    return {}

@must_be_logged_in
def get_dashboard_nodes(auth, **kwargs):
    user = auth.user

    contributed = user.node__contributed  # nodes user cotributed to

    nodes = contributed.find(
        Q('category', 'eq', 'project') &
        Q('is_deleted', 'eq', False) &
        Q('is_registration', 'eq', False)
    )

    comps = contributed.find(
        # components only
        Q('category', 'ne', 'project') &
        # parent is not in the nodes list
        Q('__backrefs.parent.node.nodes', 'nin', nodes.get_keys()) &
        # exclude deleted nodes
        Q('is_deleted', 'eq', False) &
        # exclude registrations
        Q('is_registration', 'eq', False)
    )

    return _render_nodes(list(nodes) + list(comps))

@must_be_logged_in
def dashboard(**kwargs):
    return {'addons_enabled': kwargs['auth'].user.get_addon_names()}

@must_be_logged_in
def watched_logs_get(**kwargs):
    user = kwargs['auth'].user
    page_num = int(request.args.get('pageNum', '').strip('/') or 0)
    page_size = 10
    offset = page_num * page_size
    recent_log_ids = itertools.islice(user.get_recent_log_ids(), offset, offset + page_size + 1)
    logs = (model.NodeLog.load(id) for id in recent_log_ids)
    watch_logs = []
    has_more_logs = False

    for log in logs:
        if len(watch_logs) < page_size:
            watch_logs.append(log.serialize())
        else:
            has_more_logs =True
            break

    return {"logs": watch_logs, "has_more_logs": has_more_logs}

def reproducibility():
    return redirect('/ezcuj/wiki')

def registration_form():
    return utils.jsonify(RegistrationForm(prefix='register'))

def signin_form():
    return utils.jsonify(SignInForm())

def forgot_password_form():
    return utils.jsonify(ForgotPasswordForm(prefix='forgot_password'))

def reset_password_form():
    return utils.jsonify(ResetPasswordForm())

def new_project_form():
    return utils.jsonify(NewProjectForm())

### GUID ###

def _build_guid_url(url, prefix=None, suffix=None):
    if not url.startswith('/'):
        url = '/' + url
    if not url.endswith('/'):
        url += '/'
    url = (
        (prefix or '') +
        url +
        (suffix or '')
    )
    if not url.endswith('/'):
        url += '/'
    return url

def resolve_guid(guid, suffix=None):
    """Resolve GUID to corresponding URL and return result of appropriate
    view function. This effectively yields a redirect without changing the
    displayed URL of the page.

    :param guid: GUID value (not the object)
    :param suffix: String to append to GUID route
    :return: Werkzeug response

    """
    # Get prefix; handles API routes
    prefix = request.path.split(guid)[0].rstrip('/')

    # Look up GUID
    guid_object = Guid.load(guid)
    if guid_object:
        referent = guid_object.referent
        if referent is None:
            logger.error('Referent of GUID {0} not found'.format(guid))
            raise HTTPError(http.NOT_FOUND)
        mode = referent.redirect_mode
        url = referent.deep_url if mode == 'proxy' else referent.url
        url = _build_guid_url(url, prefix, suffix)
        # Always redirect API URLs; URL should identify endpoint being called
        if prefix or mode == 'redirect':
            if request.query_string:
                url += '?' + request.query_string
            return redirect(url)
        return proxy_url(url)

    # GUID not found; try lower-cased and redirect if exists
    guid_object_lower = Guid.load(guid.lower())
    if guid_object_lower:
        return redirect(
            _build_guid_url(
                guid.lower(), prefix, suffix
            )
        )

    # GUID not found
    raise HTTPError(http.NOT_FOUND)
<commit_msg>Add error handling if a Guid resolves to a non-Guid object<commit_after># -*- coding: utf-8 -*-
import logging
import itertools
import httplib as http
from flask import request, redirect
from modularodm import Q

from framework import sentry
from framework.exceptions import HTTPError
from framework.forms import utils
from framework.routing import proxy_url
from framework.auth import Auth, get_current_user
from framework.auth.decorators import collect_auth, must_be_logged_in
from framework.auth.forms import (RegistrationForm, SignInForm,
                                  ForgotPasswordForm, ResetPasswordForm)
from framework.guid.model import Guid, GuidStoredObject

from website.util import web_url_for
from website.project.forms import NewProjectForm
from website.project import model
from website import settings

logger = logging.getLogger('website.views')

def _rescale_ratio(nodes):
    """Get scaling denominator for log lists across a sequence of nodes.

    :param nodes: Nodes
    :return: Max number of logs

    """
    if not nodes:
        return 0
    user = get_current_user()
    counts = [
        len(node.logs)
        for node in nodes
        if node.can_view(Auth(user=user))
    ]
    if counts:
        return float(max(counts))
    return 0.0

def _render_node(node):
    """

    :param node:
    :return:

    """
    return {
        'id': node._primary_key,
        'url': node.url,
        'api_url': node.api_url,
        'primary': node.primary,
    }

def _render_nodes(nodes):
    """

    :param nodes:
    :return:
    """
    ret = {
        'nodes': [
            _render_node(node)
            for node in nodes
        ],
        'rescale_ratio': _rescale_ratio(nodes),
    }
    return ret

def _get_user_activity(node, user, rescale_ratio):

    # Counters
    total_count = len(node.logs)

    # Note: It's typically much faster to find logs of a given node
    # attached to a given user using node.logs.find(...) than by
    # loading the logs into Python and checking each one. However,
    # using deep caching might be even faster down the road.

    ua_count = node.logs.find(Q('user', 'eq', user)).count()
    non_ua_count = total_count - ua_count # base length of blue bar

    # Normalize over all nodes
    ua = ua_count / rescale_ratio * settings.USER_ACTIVITY_MAX_WIDTH
    non_ua = non_ua_count / rescale_ratio * settings.USER_ACTIVITY_MAX_WIDTH

    return ua_count, ua, non_ua

@collect_auth
def index(auth, **kwargs):
    """Redirect to dashboard if user is logged in, else show homepage.

    """
    if auth.user:
        return redirect(web_url_for('dashboard'))
    return {}

@must_be_logged_in
def get_dashboard_nodes(auth, **kwargs):
    user = auth.user

    contributed = user.node__contributed  # nodes user cotributed to

    nodes = contributed.find(
        Q('category', 'eq', 'project') &
        Q('is_deleted', 'eq', False) &
        Q('is_registration', 'eq', False)
    )

    comps = contributed.find(
        # components only
        Q('category', 'ne', 'project') &
        # parent is not in the nodes list
        Q('__backrefs.parent.node.nodes', 'nin', nodes.get_keys()) &
        # exclude deleted nodes
        Q('is_deleted', 'eq', False) &
        # exclude registrations
        Q('is_registration', 'eq', False)
    )

    return _render_nodes(list(nodes) + list(comps))

@must_be_logged_in
def dashboard(**kwargs):
    return {'addons_enabled': kwargs['auth'].user.get_addon_names()}

@must_be_logged_in
def watched_logs_get(**kwargs):
    user = kwargs['auth'].user
    page_num = int(request.args.get('pageNum', '').strip('/') or 0)
    page_size = 10
    offset = page_num * page_size
    recent_log_ids = itertools.islice(user.get_recent_log_ids(), offset, offset + page_size + 1)
    logs = (model.NodeLog.load(id) for id in recent_log_ids)
    watch_logs = []
    has_more_logs = False

    for log in logs:
        if len(watch_logs) < page_size:
            watch_logs.append(log.serialize())
        else:
            has_more_logs =True
            break

    return {"logs": watch_logs, "has_more_logs": has_more_logs}

def reproducibility():
    return redirect('/ezcuj/wiki')

def registration_form():
    return utils.jsonify(RegistrationForm(prefix='register'))

def signin_form():
    return utils.jsonify(SignInForm())

def forgot_password_form():
    return utils.jsonify(ForgotPasswordForm(prefix='forgot_password'))

def reset_password_form():
    return utils.jsonify(ResetPasswordForm())

def new_project_form():
    return utils.jsonify(NewProjectForm())

### GUID ###

def _build_guid_url(url, prefix=None, suffix=None):
    if not url.startswith('/'):
        url = '/' + url
    if not url.endswith('/'):
        url += '/'
    url = (
        (prefix or '') +
        url +
        (suffix or '')
    )
    if not url.endswith('/'):
        url += '/'
    return url

def resolve_guid(guid, suffix=None):
    """Resolve GUID to corresponding URL and return result of appropriate
    view function. This effectively yields a redirect without changing the
    displayed URL of the page.

    :param guid: GUID value (not the object)
    :param suffix: String to append to GUID route
    :return: Werkzeug response

    """
    # Get prefix; handles API routes
    prefix = request.path.split(guid)[0].rstrip('/')

    # Look up GUID
    guid_object = Guid.load(guid)
    if guid_object:

        # verify that the object is a GuidStoredObject descendant. If a model
        #   was once a descendant but that relationship has changed, it's
        #   possible to have referents that are instances of classes that don't
        #   have a redirect_mode attribute or otherwise don't behave as
        #   expected.
        if not isinstance(guid_object.referent, GuidStoredObject):
            sentry.log_message(
                'Guid `{}` resolved to non-guid object'.format(guid)
            )

            raise HTTPError(http.NOT_FOUND)

        referent = guid_object.referent
        if referent is None:
            logger.error('Referent of GUID {0} not found'.format(guid))
            raise HTTPError(http.NOT_FOUND)
        mode = referent.redirect_mode
        url = referent.deep_url if mode == 'proxy' else referent.url
        url = _build_guid_url(url, prefix, suffix)
        # Always redirect API URLs; URL should identify endpoint being called
        if prefix or mode == 'redirect':
            if request.query_string:
                url += '?' + request.query_string
            return redirect(url)
        return proxy_url(url)

    # GUID not found; try lower-cased and redirect if exists
    guid_object_lower = Guid.load(guid.lower())
    if guid_object_lower:
        return redirect(
            _build_guid_url(
                guid.lower(), prefix, suffix
            )
        )

    # GUID not found
    raise HTTPError(http.NOT_FOUND)
<|endoftext|>


<commit_before>from django.shortcuts import render
from django.http import HttpResponse
from django.core.urlresolvers import reverse
from news.models import Article, Event
from door.models import DoorStatus
from datetime import datetime
from itertools import chain

def index(request):
    number_of_news = 3

    # Sorts the news to show the events nearest in future and then fill in with the newest articles
    event_list = Event.objects.filter(time_end__gte=datetime.now())[0:number_of_news:-1]
    article_list = Article.objects.order_by('-pub_date')[0:number_of_news - len(event_list)]
    news_list = list(chain(event_list, article_list))

    try:
        door_status = DoorStatus.objects.get(name='hackerspace').status
    except DoorStatus.DoesNotExist:
        door_status = True
    context = {
        'news_list': news_list,
        'door_status': door_status,
    }

    return render(request, 'index.html', context)

def opptak(request):
    return HttpResponse(reverse('application_form'))

def test404(request):
    return render(request, '404.html')
<commit_msg>Fix redirect<commit_after>from django.shortcuts import render
from django.http import HttpResponseRedirect
from django.core.urlresolvers import reverse
from news.models import Article, Event
from door.models import DoorStatus
from datetime import datetime
from itertools import chain

def index(request):
    number_of_news = 3

    # Sorts the news to show the events nearest in future and then fill in with the newest articles
    event_list = Event.objects.filter(time_end__gte=datetime.now())[0:number_of_news:-1]
    article_list = Article.objects.order_by('-pub_date')[0:number_of_news - len(event_list)]
    news_list = list(chain(event_list, article_list))

    try:
        door_status = DoorStatus.objects.get(name='hackerspace').status
    except DoorStatus.DoesNotExist:
        door_status = True
    context = {
        'news_list': news_list,
        'door_status': door_status,
    }

    return render(request, 'index.html', context)

def opptak(request):
    return HttpResponseRedirect(reverse('application_form'))

def test404(request):
    return render(request, '404.html')
<|endoftext|>


<commit_before>#!/usr/bin/python3.5

#Author:        Sasan Bahadaran
#Date:          5/1/17
#Organization:  Commerce Data Service
#Description:   This is a bot script for getting Census data from the Census
#Bureau API's and writing it to Wikipedia.

import pywikibot, json, os, requests, argparse, logging, time, json, sys
import mwparserfromhell, datetime, math
from pywikibot.data import api
from util import config_funcs

#get values from CENSUS API.  Return response from first year with valid response starting with
#current year and ending with 2013
def get_census_values(api_url, get_var, for_var, api_key, year=datetime.datetime.today().year):
    try:
        year = datetime.datetime.today().year
        while year >= 2013:
            payload = {'get': get_var, 'for': for_var, 'key': api_key}
            r = requests.get(api_url.replace('XXXX', str(year)), params=payload)
            if r.status_code == 200:
                return r.json(), str(year)
            else:
                logging.info('No API Results for year: {}'.format(year))
                year = year - 1
        else:
            return
    except requests.exceptions.RequestException as e:
        logging.error('General Exception: {}'.format(e))
    except IOError as err:
        logging.error('IOError: {}'.format(err))

def search_for_page_items(template, infobox_keys):
    template_values = {}
    for item, item_keys in infobox_keys.items():
        for key in item_keys:
            if template.has(key):
                # added this condition for cases in which both properties are present
                #if key == 'population_total' and template.has('population_est'):
                #    key = ' population_est'
                template_values[item] = str(template.get(key))
                logging.info('KEY - {} Found'.format(key))
                break
    return template_values

#sort items by population (exluding PR and DC)
def population_rank_sort(pop_list):
    non_states = []
    for i, val in enumerate(pop_list):
        #11 and 72 are the FIPS 5-2 codes for PR and DC
        if val[2] in ['11', '72']:
            non_states.append(pop_list.pop(i))
    pop_list = sorted(pop_list, key=lambda x: int(x[1]), reverse=True)
    ordinal = lambda n: "%d%s" % (n,"tsnrhtdd"[(n//10%10!=1)*(n%10<4)*n%10::4])
    for i,val in enumerate(pop_list):
        val.append(ordinal(i+1))
    #leaving this commented out for now because we decided to exclude these from
    #state processing
    #pop_list.extend(non_states)
    return pop_list

#remove value from dict and return copy rather than mutated dict
def removekey(d, key):
    r = dict(d)
    del r[key]
    return r

# compare value from page to value from API to check for discrepancies
# *note - this function does not compare references, only values
def compare_page_items(api_values, page_values, year):
    for key, val in page_values.items():
        pos = int(key.split(' - ')[1])
        # extact value and normalize whitespace
        current_val = ' '.join(val.split('=', 1)[1].split())
        if pos == 99:
            api_value = year
        else:
            api_value = api_values[pos]
        #if key == 'total_pop - 1':
            # format value correctly for comparison with page content
            #api_value += ' ('+year+' est.)'
            # remove reference, commas, and any
 tags
            #current_val = current_val[:current_val.find('<ref')].replace(',', '').replace('
', ' ')
        if '<ref' in current_val:
            current_val = current_val[:current_val.find('<ref')]
        current_val = current_val.replace(',', '').replace('
', ' ')
        logging.info('KEY: {} | EXISTING VALUE: {} | NEW VALUE: {}'.format(key.split(' - ')[0], current_val, api_value))
        if current_val == api_value:
            logging.info('VALUES MATCH')
            page_values = removekey(page_values, key)
        else:
            logging.info('VALUES DO NOT MATCH')
    return page_values

def create_comment(comment_vals):
    comment = ''
    beg = 'Updating'
    end = 'with latest data from US Census Bureau.'
    comment = beg + 'and'.join(comment_vals) + end
    return comment

def update_page_items(page, text, api_values, page_values, year, reference):
    global num_of_edits
    comment_vals = []
    for key, val in page_values.items():
        pos = int(key.split(' - ')[1])
        if pos == 99:
            new_value = year
        else:
            #new_value = api_values[pos]
            new_value = '{:,}'.format(int(api_values[pos]))
        # add reference and create comment
        if key == 'total_pop - 1':
            new_value = '{:,}'.format(int(new_value))+' ('+year+' est.)'+reference
            comment_vals.append(' total population ')
        else:
            comment_vals.append(' population rank ')
        # add property tag
        new_value = ' '.join(val.split('=', 1)[0].split())+' = '+new_value
        # add new line tag
        new_value += '\n'
        logging.info('FULL EXISTING VALUE: {}\nFULL REPLACEMENT VALUE: {}'.format(val, new_value))
        text = text.replace(val, new_value)
    comment = create_comment(comment_vals)
    if not args.debug:
         page.text = text
         page.save(comment)
         logging.info('Page Successfully Updated')
    else:
         logging.info('DEBUG - Page value will be updated')
    num_of_edits += 1

if __name__ == '__main__':
    scriptpath = os.path.dirname(os.path.abspath(__file__))

    #logging configuration
    logging.basicConfig(
                        filename='logs/wikipedia_bot-log-'+time.strftime('%Y%m%d'),
                        level=logging.INFO,
                        format='%(asctime)s - %(name)s - %(levelname)s -%(message)s',
                        datefmt='%Y%m%d %H:%M:%S'
    )

    parser = argparse.ArgumentParser()
    parser.add_argument(
                        '-m',
                        '--mode',
                        required=True,
                        type=str,
                        choices=['t', 'p'],
                        help='Pass a t flag for test mode or a p flag for production mode'
    )
    parser.add_argument(
                        '-d',
                        '--debug',
                        required=False,
                        action='store_true',
                        default=False,
                        help='Pass this flag to set mode of bot to debug (means there will be no actual changes to the Wiki)'
    )
    parser.add_argument(
                        '-n',
                        '--numedits',
                        required=False,
                        type=int,
                        help='Pass this flag to control the number of edits that the bot makes.'
    )
    args = parser.parse_args()
    logging.info("-------- [SCRIPT ARGUMENTS] --------")
    if args.mode == 't':
        logging.info('      BOT MODE: TEST')
    elif args.mode == 'p':
        logging.info('      BOT MODE: PROD')
    if args.debug:
        logging.info('      !RUNNING IN DEBUG MODE!')
    if args.numedits:
        logging.info('      NUMBER OF EDITS:{}'.format(args.numedits))
    logging.info("----------- [JOB START] -----------")

    # can be used later for config file
    #data_file = os.path.join(scriptpath, "data", "data.json")

    # state set of configs
    #get_var = 'GEONAME,POP'
    #for_var = 'state:*'
    #api_url = 'http://api.census.gov/data/XXXX/pep/population'
    #api_key = config_funcs.getAppConfigParam('API', 'key')
    ## the keys in this dict represent the 'property - position of corresponding value in api response
    ## and the values represent the possible keys in which this property is listed under in the template
    #infobox_keys = {'total_pop - 1': ['population_total', '2010Pop', '2000Pop', 'population_estimate'],
    #        'rank - 3': ['PopRank']
    #        }
    #reference = '<ref name=PopHousingEst>{{{{cite web|url=https://www.census.gov/programs-surveys/popest.html|title=Population'\
    #                    ' and Housing Unit Estimates |date={} |accessdate={}|publisher=[[U.S. Census Bureau]]}}}}</ref>'\
    #                    .format(datetime.datetime.today().strftime('%B %-d, %Y'), datetime.datetime.today().strftime('%B %-d, %Y'))
    ## position of state code in API response
    #code_check_pos = 2
    ##DC and PR
    #exceptions = ['11', '72']
    #key_exceptions = {'Kansas': 'Kansas, United States', 'North Carolina': 'North Carolina, United States',
    #        'Georgia': 'Georgia, United States', 'Washington': 'Washington (state)'}
    #test_data = [['User:Sasan-CDS/sandbox', '555555', '50', '11th']]

    #county parameters
    get_var = 'GEONAME,POP'
    for_var = 'county:*'
    api_url = 'http://api.census.gov/data/XXXX/pep/population'
    api_key = config_funcs.getAppConfigParam('API', 'key')
    # the keys in this dict represent the 'property - position of corresponding value in api response
    # and the values represent the possible keys in which this property is listed under in the template
    infobox_keys = {'population - 1': ['pop', 'population_total', 'population'],
            'population_estimate - 1': ['population_est'],
            'population_as_of - 99': ['popularion_as_of'],
            'population_est_as_of - 99': ['pop_est_as_of', 'census_estimate_yr', 'census estimate yr', 'population_date',
                'census yr']
            }
    reference = '<ref name=PopHousingEst>{{{{cite web|url=https://www.census.gov/programs-surveys/popest.html|title=Population'\
                        ' and Housing Unit Estimates |date={} |accessdate={}|publisher=[[U.S. Census Bureau]]}}}}</ref>'\
                        .format(datetime.datetime.today().strftime('%B %-d, %Y'), datetime.datetime.today().strftime('%B %-d, %Y'))
    # position of value in API response to check for exceptions
    code_check_pos = 2
    #DC and PR
    exceptions = ['11', '72']
    relevant_templates = ['Infobox settlement', 'Infobox U.S. county']
    key_exceptions = {'Winchester city, Virginia': 'Winchester, Virginia', 'Waynesboro city, Virginia': 'Waynesboro, Virginia',
            'Virginia Beach city, Virginia': 'Virginia Beach, Virginia', 'Suffolk city, Virginia': 'Suffolk, Virginia',
            'St. Louis city, Missouri': 'St. Louis, Missouri',
            'Salem city, Virginia': 'Salem, Virginia', 'Roanoke city, Virginia': 'Roanoke, Virginia',
            'Radford city, Virginia': 'Radford, Virginia', 'Portsmouth city, Virginia': 'Portsmouth, Virginia',
            'Poquoson city, Virginia': 'Poquoson, Virginia', 'Newport News city, Virginia': 'Newport News, Virginia',
            'Manassas city, Virginia': 'Manassas, Virginia', 'Lynchburg city, Virginia': 'Lynchburg, Virginia', 'Lexington city, Virginia': 'Lexington, Virginia',
            'Franklin city, Virginia': 'Franklin, Virginia', 'Emporia city, Virginia': 'Emporia, Virginia', 'Covington city, Virginia': 'Covington, Virginia',
            'Buena Vista city, Virginia': 'Buena Vista, Virginia', 'Baltimore city, Maryland': 'Baltimore, Maryland', 'Alexandria city, Virginia': 'Alexandria, Virginia'
            }
    test_data = [['User:Sasan-CDS/sandbox', '555555', '50']]

    num_of_pages_not_found = 0
    num_of_not_founds = 0
    num_of_edits = 0
    site = pywikibot.Site('en', 'wikipedia')
    repo = site.data_repository()

    if args.mode == 'p':
        metric_values, year = get_census_values(api_url, get_var, for_var, api_key)
        #remove header
        metric_values.pop(0)
        #metric_values = population_rank_sort(metric_values)
    else:
        metric_values = test_data
        year = '2016'
    if metric_values:
        logging.info('Number of items in API Response: {}'.format(len(metric_values)))
        for i, api_val in enumerate(metric_values):
            logging.info('API item: {}'.format(api_val))
            #key = api_val[0].split(',')[0]
            key = api_val[0]
            logging.info('[ITEM: {}]'.format(key))
            if key in key_exceptions:
                key = key_exceptions[key]
                logging.info('Key exception found. new key: {}'.format(key))
            if api_val[code_check_pos] in exceptions:
                logging.info('Item being skipped due to configuration settings')
                continue
            page = pywikibot.Page(site, key)
            if page.exists():
                if page.isRedirectPage():
                    page = page.getRedirectTarget()
                text = page.get(get_redirect=True)
                code = mwparserfromhell.parse(text)
                template_values = {}
                template_found = False
                for template in code.filter_templates():
                    if template.name in relevant_templates:
                        template_found = True
                    # if correct template is found, break out of loop
                    if template_values:
                        break
                    else:
                        template_values = search_for_page_items(template, infobox_keys)
                if template_values:
                    if 'population_estimate - 1' in template_values:
                        logging.info('Ignoring population_total and population_as_of properties due to presence of population_estimate property')
                        del template_values['population - 1']
                        if 'population_as_of - 99' in template_values:
                            del template_values['population_as_of - 99']
                        elif 'population_est_as_of - 99' in template_values:
                            del template_values['population_est_as_of - 99']
                    #compare page items
                    template_values = compare_page_items(api_val, template_values, year)
                    if template_values:
                        if args.numedits:
                            if num_of_edits < args.numedits:
                                #update_page_items(page, text, api_val, template_values, year, reference)
                                logging.info('Number of edits: {}'.format(num_of_edits))
                            else:
                                logging.info('Number of maximum edits({}) has been reached and bot will not perform any further updates')
                        else:
                            update_page_items(page, text, api_val, template_values, year, reference)
                            logging.info('Number of edits: {}'.format(num_of_edits))
                    else:
                        logging.info('Nothing to update')
                else:
                    if template_found:
                        logging.warning('No items were found in this page!!!')
                    else:
                        logging.warning('None of the relevant templates were found on this page!!!')
                    num_of_not_founds += 1
            else:
                logging.warning('NO RESULTS FOR: {}'.format(key))
                num_of_pages_not_found += 1
        logging.info('TOTAL NUMBER OF EDITS: {}'.format(num_of_edits))
        logging.info('TOTAL NUMBER OF PAGES WHERE NOTHING FOUND: {}'.format(num_of_not_founds))
        logging.info('TOTAL NUMBER OF PAGES NOT FOUND: {}'.format(num_of_pages_not_found))
    else:
        sys.exit('NO RESULTS FROM THE CENSUS API FOR ANY YEARS.  EXAMINE FOR OTHER ISSUES!')
<commit_msg>cleaned up code and merged some code from state file<commit_after>#!/usr/bin/python3.5

#Author:        Sasan Bahadaran
#Date:          5/1/17
#Organization:  Commerce Data Service
#Description:   This is a bot script for getting Census data from the Census
#Bureau API's and writing it to Wikipedia.

import pywikibot, json, os, requests, argparse, logging, time, json, sys
import mwparserfromhell, datetime, math
from pywikibot.data import api
from util import config_funcs

#get values from CENSUS API.  Return response from first year with valid response starting with
#current year and ending with 2013
def get_census_values(api_url, get_var, for_var, api_key, year=datetime.datetime.today().year):
    try:
        year = datetime.datetime.today().year
        while year >= 2013:
            payload = {'get': get_var, 'for': for_var, 'key': api_key}
            r = requests.get(api_url.replace('XXXX', str(year)), params=payload)
            if r.status_code == 200:
                return r.json(), str(year)
            else:
                logging.info('No API Results for year: {}'.format(year))
                year = year - 1
        else:
            return
    except requests.exceptions.RequestException as e:
        logging.error('General Exception: {}'.format(e))
    except IOError as err:
        logging.error('IOError: {}'.format(err))

def search_for_page_items(template, infobox_keys):
    template_values = {}
    for item, item_keys in infobox_keys.items():
        for key in item_keys:
            if template.has(key):
                template_values[item] = str(template.get(key))
                logging.info('KEY - {} Found'.format(key))
                break
    return template_values

#sort items by population (exluding PR and DC)
def population_rank_sort(pop_list):
    non_states = []
    for i, val in enumerate(pop_list):
        #11 and 72 are the FIPS 5-2 codes for PR and DC
        if val[2] in ['11', '72']:
            non_states.append(pop_list.pop(i))
    pop_list = sorted(pop_list, key=lambda x: int(x[1]), reverse=True)
    ordinal = lambda n: "%d%s" % (n,"tsnrhtdd"[(n//10%10!=1)*(n%10<4)*n%10::4])
    for i,val in enumerate(pop_list):
        val.append(ordinal(i+1))
    #leaving this commented out for now because we decided to exclude these from
    #state processing
    #pop_list.extend(non_states)
    return pop_list

#remove value from dict and return copy rather than mutated dict
def removekey(d, key):
    r = dict(d)
    del r[key]
    return r

# compare value from page to value from API to check for discrepancies
# *note - this function does not compare references, only values
def compare_page_items(api_values, page_values, year):
    for key, val in page_values.items():
        pos = int(key.split(' - ')[1])
        # extact value and normalize whitespace
        current_val = ' '.join(val.split('=', 1)[1].split())
        if pos == 99:
            api_value = year
        else:
            api_value = api_values[pos]
        if 'state' in for_var:
            if key == 'total_pop - 1':
                # format value correctly for comparison with page content
                api_value += ' ('+year+' est.)'
        if '<ref' in current_val:
            current_val = current_val[:current_val.find('<ref')]
        current_val = current_val.replace(',', '').replace('
', ' ')
        logging.info('KEY: {} | EXISTING VALUE: {} | NEW VALUE: {}'.format(key.split(' - ')[0], current_val, api_value))
        if current_val == api_value:
            logging.info('VALUES MATCH')
            page_values = removekey(page_values, key)
        else:
            logging.info('VALUES DO NOT MATCH')
    return page_values

def create_comment(comment_vals):
    comment = ''
    beg = 'Updating'
    end = 'with latest data from US Census Bureau.'
    comment = beg + 'and'.join(comment_vals) + end
    return comment

def update_page_items(page, text, api_values, page_values, year, reference):
    global num_of_edits
    comment_vals = []
    for key, val in page_values.items():
        pos = int(key.split(' - ')[1])
        if pos == 99:
            new_value = year
        else:
            #new_value = api_values[pos]
            new_value = '{:,}'.format(int(api_values[pos]))
        # add reference and create comment
        if key == 'total_pop - 1':
            new_value = '{:,}'.format(int(new_value))+' ('+year+' est.)'+reference
            comment_vals.append(' total population ')
        else:
            comment_vals.append(' population rank ')
        # add property tag
        new_value = ' '.join(val.split('=', 1)[0].split())+' = '+new_value
        # add new line tag
        new_value += '\n'
        logging.info('FULL EXISTING VALUE: {}\nFULL REPLACEMENT VALUE: {}'.format(val, new_value))
        text = text.replace(val, new_value)
    comment = create_comment(comment_vals)
    if not args.debug:
         page.text = text
         page.save(comment)
         logging.info('Page Successfully Updated')
    else:
         logging.info('DEBUG - Page value will be updated')
    num_of_edits += 1

if __name__ == '__main__':
    scriptpath = os.path.dirname(os.path.abspath(__file__))

    #logging configuration
    logging.basicConfig(
                        filename='logs/wikipedia_bot-log-'+time.strftime('%Y%m%d'),
                        level=logging.INFO,
                        format='%(asctime)s - %(name)s - %(levelname)s -%(message)s',
                        datefmt='%Y%m%d %H:%M:%S'
    )

    parser = argparse.ArgumentParser()
    parser.add_argument(
                        '-m',
                        '--mode',
                        required=True,
                        type=str,
                        choices=['t', 'p'],
                        help='Pass a t flag for test mode or a p flag for production mode'
    )
    parser.add_argument(
                        '-d',
                        '--debug',
                        required=False,
                        action='store_true',
                        default=False,
                        help='Pass this flag to set mode of bot to debug (means there will be no actual changes to the Wiki)'
    )
    parser.add_argument(
                        '-n',
                        '--numedits',
                        required=False,
                        type=int,
                        help='Pass this flag to control the number of edits that the bot makes.'
    )
    args = parser.parse_args()
    logging.info("-------- [SCRIPT ARGUMENTS] --------")
    if args.mode == 't':
        logging.info('      BOT MODE: TEST')
    elif args.mode == 'p':
        logging.info('      BOT MODE: PROD')
    if args.debug:
        logging.info('      !RUNNING IN DEBUG MODE!')
    if args.numedits:
        logging.info('      NUMBER OF EDITS:{}'.format(args.numedits))
    logging.info("----------- [JOB START] -----------")

    # state set of configs
    #get_var = 'GEONAME,POP'
    #for_var = 'state:*'
    #api_url = 'http://api.census.gov/data/XXXX/pep/population'
    #api_key = config_funcs.getAppConfigParam('API', 'key')
    ## the keys in this dict represent the 'property - position of corresponding value in api response
    ## and the values represent the possible keys in which this property is listed under in the template
    #infobox_keys = {'total_pop - 1': ['population_total', '2010Pop', '2000Pop', 'population_estimate'],
    #        'rank - 3': ['PopRank']
    #        }
    #reference = '<ref name=PopHousingEst>{{{{cite web|url=https://www.census.gov/programs-surveys/popest.html|title=Population'\
    #                    ' and Housing Unit Estimates |date={} |accessdate={}|publisher=[[U.S. Census Bureau]]}}}}</ref>'\
    #                    .format(datetime.datetime.today().strftime('%B %-d, %Y'), datetime.datetime.today().strftime('%B %-d, %Y'))
    ## position of state code in API response
    #code_check_pos = 2
    ##DC and PR
    #exceptions = ['11', '72']
    #key_exceptions = {'Kansas': 'Kansas, United States', 'North Carolina': 'North Carolina, United States',
    #        'Georgia': 'Georgia, United States', 'Washington': 'Washington (state)'}
    #test_data = [['User:Sasan-CDS/sandbox', '555555', '50', '11th']]

    #county parameters
    get_var = 'GEONAME,POP'
    for_var = 'county:*'
    api_url = 'http://api.census.gov/data/XXXX/pep/population'
    api_key = config_funcs.getAppConfigParam('API', 'key')
    # the keys in this dict represent the 'property - position of corresponding value in api response
    # and the values represent the possible keys in which this property is listed under in the template
    infobox_keys = {'population - 1': ['pop', 'population_total', 'population'],
            'population_estimate - 1': ['population_est'],
            'population_as_of - 99': ['popularion_as_of'],
            'population_est_as_of - 99': ['pop_est_as_of', 'census_estimate_yr', 'census estimate yr', 'population_date',
                'census yr']
            }
    reference = '<ref name=PopHousingEst>{{{{cite web|url=https://www.census.gov/programs-surveys/popest.html|title=Population'\
                        ' and Housing Unit Estimates |date={} |accessdate={}|publisher=[[U.S. Census Bureau]]}}}}</ref>'\
                        .format(datetime.datetime.today().strftime('%B %-d, %Y'), datetime.datetime.today().strftime('%B %-d, %Y'))
    # position of value in API response to check for exceptions
    code_check_pos = 2
    #DC and PR
    exceptions = ['11', '72']
    relevant_templates = ['Infobox settlement', 'Infobox U.S. county']
    key_exceptions = {'Winchester city, Virginia': 'Winchester, Virginia', 'Waynesboro city, Virginia': 'Waynesboro, Virginia',
            'Virginia Beach city, Virginia': 'Virginia Beach, Virginia', 'Suffolk city, Virginia': 'Suffolk, Virginia',
            'St. Louis city, Missouri': 'St. Louis, Missouri','Salem city, Virginia': 'Salem, Virginia',
            'Roanoke city, Virginia': 'Roanoke, Virginia', 'Radford city, Virginia': 'Radford, Virginia',
            'Portsmouth city, Virginia': 'Portsmouth, Virginia', 'Poquoson city, Virginia': 'Poquoson, Virginia',
            'Newport News city, Virginia': 'Newport News, Virginia', 'Manassas city, Virginia': 'Manassas, Virginia',
            'Lynchburg city, Virginia': 'Lynchburg, Virginia', 'Lexington city, Virginia': 'Lexington, Virginia',
            'Franklin city, Virginia': 'Franklin, Virginia', 'Emporia city, Virginia': 'Emporia, Virginia',
            'Covington city, Virginia': 'Covington, Virginia', 'Buena Vista city, Virginia': 'Buena Vista, Virginia',
            'Baltimore city, Maryland': 'Baltimore, Maryland', 'Alexandria city, Virginia': 'Alexandria, Virginia'
            }
    test_data = [['User:Sasan-CDS/sandbox', '555555', '50']]

    num_of_pages_not_found = 0
    num_of_not_founds = 0
    num_of_edits = 0
    site = pywikibot.Site('en', 'wikipedia')
    repo = site.data_repository()

    if args.mode == 'p':
        metric_values, year = get_census_values(api_url, get_var, for_var, api_key)
        #remove header
        metric_values.pop(0)
        #metric_values = population_rank_sort(metric_values)
    else:
        metric_values = test_data
        year = '2016'
    if metric_values:
        logging.info('Number of items in API Response: {}'.format(len(metric_values)))
        for i, api_val in enumerate(metric_values):
            logging.info('API item: {}'.format(api_val))
            #key = api_val[0].split(',')[0]
            key = api_val[0]
            logging.info('[ITEM: {}]'.format(key))
            if key in key_exceptions:
                key = key_exceptions[key]
                logging.info('Key exception found. new key: {}'.format(key))
            if api_val[code_check_pos] in exceptions:
                logging.info('Item being skipped due to configuration settings')
                continue
            page = pywikibot.Page(site, key)
            if page.exists():
                if page.isRedirectPage():
                    page = page.getRedirectTarget()
                text = page.get(get_redirect=True)
                code = mwparserfromhell.parse(text)
                template_values = {}
                template_found = False
                for template in code.filter_templates():
                    if template.name in relevant_templates:
                        template_found = True
                    # if correct template is found, break out of loop
                    if template_values:
                        break
                    else:
                        template_values = search_for_page_items(template, infobox_keys)
                if template_values:
                    if 'population_estimate - 1' in template_values:
                        logging.info('Ignoring population_total and population_as_of properties due to presence of population_estimate property')
                        del template_values['population - 1']
                        if 'population_as_of - 99' in template_values:
                            del template_values['population_as_of - 99']
                        elif 'population_est_as_of - 99' in template_values:
                            del template_values['population_est_as_of - 99']
                    #compare page items
                    template_values = compare_page_items(api_val, template_values, year)
                    if template_values:
                        if args.numedits:
                            if num_of_edits < args.numedits:
                                #update_page_items(page, text, api_val, template_values, year, reference)
                                logging.info('Number of edits: {}'.format(num_of_edits))
                            else:
                                logging.info('Number of maximum edits({}) has been reached and bot will not perform any further updates')
                        else:
                            update_page_items(page, text, api_val, template_values, year, reference)
                            logging.info('Number of edits: {}'.format(num_of_edits))
                    else:
                        logging.info('Nothing to update')
                else:
                    if template_found:
                        logging.warning('No items were found in this page!!!')
                    else:
                        logging.warning('None of the relevant templates were found on this page!!!')
                    num_of_not_founds += 1
            else:
                logging.warning('NO RESULTS FOR: {}'.format(key))
                num_of_pages_not_found += 1
        logging.info('TOTAL NUMBER OF EDITS: {}'.format(num_of_edits))
        logging.info('TOTAL NUMBER OF PAGES WHERE NOTHING FOUND: {}'.format(num_of_not_founds))
        logging.info('TOTAL NUMBER OF PAGES NOT FOUND: {}'.format(num_of_pages_not_found))
    else:
        sys.exit('NO RESULTS FROM THE CENSUS API FOR ANY YEARS.  EXAMINE FOR OTHER ISSUES!')
<|endoftext|>


<commit_before>"""
    wispy.builder
    ~~~~~~~~~~~~~

    Contains utilities for building a CST node
    from the grammars.
"""

# pylint: disable=bad-builtin, missing-docstring
# pylint: disable=no-self-use, no-init, invalid-name

from functools import partial

from . import tree
from . import grammar

__all__ = ['build_tree', 'Builder']

def build_tree(grammar_node):
    """
    Visit the nodes from the given grammar, returning
    a new AST node.
    """
    builder = Builder()
    return builder.visit_scriptblock(grammar_node)

class Builder:
    """ An AST builder, implemented using the visitor pattern. """

    def iter_generic_visit(self, nodes, parent):
        """ Visit each nodes from the given *nodes*, by calling
        :meth:`generic_visit` for each of those nodes.
        """
        meth = partial(self.generic_visit, parent=parent)
        return list(filter(None, map(meth, nodes)))

    def generic_visit(self, node, parent):
        """ A generic visit method, which visits the given node,
        if the builder implements it. If so, it should return
        a new AST node representing the type of the node.
        """

        cls = node.__class__.__name__.replace("<", "").replace(">", "")
        visit_name = 'visit_' + cls.lower()
        visit_method = getattr(self, visit_name, None)
        if visit_method:
            return visit_method(node, parent)

    def visit_children(self, node):
        """ Visit the children of the given node, by calling
        :meth:`generic_visit` for each of them.
        Returns a list of AST nodes, representing the children.
        """
        return self.iter_generic_visit(node, node)

    def visit_scriptblock(self, node):
        """
        Visit a ScriptBlock grammar node and return the AST
        representing it.
        """
        stmts = node.find_all(grammar.Statement)
        blocks = node.find_all(grammar.NamedBlock)

        newnode = tree.ScriptBlock()
        newnode.body = self.visit_children(node)
        newnode.statements = self.iter_generic_visit(stmts, newnode)
        newnode.named_blocks = self.iter_generic_visit(blocks, newnode)
        newnode.grammar = node
        return newnode

    def visit_namedblock(self, node, parent):
        """ Visit a NamedBlock. """
        stmts = node.find_all(grammar.Statement)
        newnode = tree.NamedBlock()
        newnode.parent = parent
        newnode.block_name = tree.Name(value=node[0].string.lower())
        newnode.statements = self.iter_generic_visit(stmts, newnode)
        return newnode

    def visit_statement(self, node, parent):
        return self.generic_visit(node[0], parent)

    def visit_functionstatement(self, node, parent):
        newnode = tree.FunctionStatement()
        newnode.parent = parent
        newnode.type = tree.Name(value=node[0].string.lower())
        newnode.name = tree.Name(value=node[2].string)
        newnode.body = self.visit_scriptblock(node[8])
        newnode.body.parent = newnode
        newnode.params = self.generic_visit(node[4], newnode)
        return newnode

    def visit_ifstatement(self, node, parent):
        stmts = node.find_all(grammar.Statement)
        elifs = node.find_all(grammar.ElseIfClause)
        orelse = node.find_all(grammar.ElseClause)

        newnode = tree.IfStatement()
        newnode.parent = parent
        newnode.test = self.generic_visit(node[4], newnode)
        newnode.body = self.iter_generic_visit(stmts, newnode)
        newnode.elifs = self.iter_generic_visit(elifs, newnode)
        newnode.orelse = self.iter_generic_visit(orelse, newnode)
        return newnode

    def visit_elseclause(self, node, parent):
        # TODO
        pass

    def visit_elseifclause(self, node, parent):
        # TODO
        pass

    def visit_functionparameterdeclaration(self, node, parent):
        script_params = node[2].find_all(grammar.ScriptParameter)
        return self.iter_generic_visit(script_params, parent)

    def visit_scriptparameter(self, node, parent):
        newnode = tree.Parameter()
        attributes = node.find_all(grammar.Attribute)
        newnode.parent = parent
        newnode.variable = self.generic_visit(node[3], newnode)
        newnode.default = self.generic_visit(node[4], newnode)
        newnode.attributes = self.iter_generic_visit(
            attributes, newnode.variable)
        return newnode

    def visit_typespec(self, node, parent):
        newnode = tree.TypeSpec()
        newnode.parent = parent
        newnode.name = tree.Name(value=node[0].string)
        specs = node[1].find_all(grammar.TypeSpec) or []
        newnode.types = self.iter_generic_visit(specs, newnode)
        return newnode
<commit_msg>The visited methods must have underscores between terms.<commit_after>"""
    wispy.builder
    ~~~~~~~~~~~~~

    Contains utilities for building a CST node
    from the grammars.
"""

# pylint: disable=bad-builtin, missing-docstring
# pylint: disable=no-self-use, no-init, invalid-name

from functools import partial

from . import tree
from . import grammar

__all__ = ['build_tree', 'Builder']

def to_underscore(string):
    new_string = []
    for char in string:
        if char.isupper():
            char = char.lower()
            new_string.append("_")
        new_string.append(char)
    return "".join(new_string).strip("_")

def build_tree(grammar_node):
    """
    Visit the nodes from the given grammar, returning
    a new AST node.
    """
    builder = Builder()
    return builder.visit_script_block(grammar_node)

class Builder:
    """ An AST builder, implemented using the visitor pattern. """

    def iter_generic_visit(self, nodes, parent):
        """ Visit each nodes from the given *nodes*, by calling
        :meth:`generic_visit` for each of those nodes.
        """
        meth = partial(self.generic_visit, parent=parent)
        return list(filter(None, map(meth, nodes)))

    def generic_visit(self, node, parent):
        """ A generic visit method, which visits the given node,
        if the builder implements it. If so, it should return
        a new AST node representing the type of the node.
        """

        cls = node.__class__.__name__.replace("<", "").replace(">", "")
        visit_name = 'visit_' + to_underscore(cls)
        visit_method = getattr(self, visit_name, None)
        if visit_method:
            return visit_method(node, parent)

    def visit_children(self, node):
        """ Visit the children of the given node, by calling
        :meth:`generic_visit` for each of them.
        Returns a list of AST nodes, representing the children.
        """
        return self.iter_generic_visit(node, node)

    def visit_script_block(self, node):
        """
        Visit a ScriptBlock grammar node and return the AST
        representing it.
        """
        stmts = node.find_all(grammar.Statement)
        blocks = node.find_all(grammar.NamedBlock)

        newnode = tree.ScriptBlock()
        newnode.body = self.visit_children(node)
        newnode.statements = self.iter_generic_visit(stmts, newnode)
        newnode.named_blocks = self.iter_generic_visit(blocks, newnode)
        newnode.grammar = node
        return newnode

    def visit_named_block(self, node, parent):
        """ Visit a NamedBlock. """
        stmts = node.find_all(grammar.Statement)
        newnode = tree.NamedBlock()
        newnode.parent = parent
        newnode.block_name = tree.Name(value=node[0].string.lower())
        newnode.statements = self.iter_generic_visit(stmts, newnode)
        return newnode

    def visit_statement(self, node, parent):
        return self.generic_visit(node[0], parent)

    def visit_function_statement(self, node, parent):
        newnode = tree.FunctionStatement()
        newnode.parent = parent
        newnode.type = tree.Name(value=node[0].string.lower())
        newnode.name = tree.Name(value=node[2].string)
        newnode.body = self.visit_script_block(node[8])
        newnode.body.parent = newnode
        newnode.params = self.generic_visit(node[4], newnode)
        return newnode

    def visit_if_statement(self, node, parent):
        stmts = node.find_all(grammar.Statement)
        elifs = node.find_all(grammar.ElseIfClause)
        orelse = node.find_all(grammar.ElseClause)

        newnode = tree.IfStatement()
        newnode.parent = parent
        newnode.test = self.generic_visit(node[4], newnode)
        newnode.body = self.iter_generic_visit(stmts, newnode)
        newnode.elifs = self.iter_generic_visit(elifs, newnode)
        newnode.orelse = self.iter_generic_visit(orelse, newnode)
        return newnode

    def visit_else_clause(self, node, parent):
        # TODO
        pass

    def visit_elseif_clause(self, node, parent):
        # TODO
        pass

    def visit_function_parameter_declaration(self, node, parent):
        script_params = node[2].find_all(grammar.ScriptParameter)
        return self.iter_generic_visit(script_params, parent)

    def visit_script_parameter(self, node, parent):
        newnode = tree.Parameter()
        attributes = node.find_all(grammar.Attribute)
        newnode.parent = parent
        newnode.variable = self.generic_visit(node[3], newnode)
        newnode.default = self.generic_visit(node[4], newnode)
        newnode.attributes = self.iter_generic_visit(
            attributes, newnode.variable)
        return newnode

    def visit_type_spec(self, node, parent):
        newnode = tree.TypeSpec()
        newnode.parent = parent
        newnode.name = tree.Name(value=node[0].string)
        specs = node[1].find_all(grammar.TypeSpec) or []
        newnode.types = self.iter_generic_visit(specs, newnode)
        return newnode
<|endoftext|>


<commit_before>"""pylint plugin to do static analysis on stbt scripts

Used by "stbt lint".

Documentation on Abstract Syntax Tree traversal with python/pylint:

* http://docs.pylint.org/extend.html#writing-your-own-checker
* http://hg.logilab.org/review/pylint/file/default/examples/custom.py
* http://docs.python.org/2/library/ast.html

"""
from __future__ import unicode_literals
from __future__ import print_function
from __future__ import division
from __future__ import absolute_import
from builtins import *  # pylint:disable=redefined-builtin,unused-wildcard-import,wildcard-import,wrong-import-order
from future.utils import string_types

import os
import re
import subprocess

from astroid import (
    Assert, BinOp, Call, ClassDef, Const, Expr, FunctionDef, Keyword, MANAGER,
    Name, Raise, YES)
from pylint.checkers import BaseChecker
from pylint.interfaces import IAstroidChecker

class StbtChecker(BaseChecker):
    __implements__ = IAstroidChecker
    name = 'stb-tester'
    msgs = {
        # Range 70xx reserved for custom checkers: www.logilab.org/ticket/68057
        # When you add a new checker update the docstring in ../stbt_lint.py
        'E7001': ('Image "%s" not found on disk',
                  'stbt-missing-image',
                  'The image path given to "stbt.match" '
                  '(and similar functions) does not exist on disk.'),
        'E7002': ('"%s" return value not used (missing "assert"?)',
                  'stbt-unused-return-value',
                  "This function does not raise an exception on failure but "
                  "you aren't using its return value. Perhaps you've forgotten "
                  'to use "assert".'),
        'E7003': ('"wait_until" argument "%s" isn\'t callable',
                  'stbt-wait-until-callable',
                  'The argument given to "wait_until" must be a callable '
                  '(such as a function or a lambda expression).'),
        'E7004': ('"%s" missing "frame" argument',
                  'stbt-frame-object-missing-frame',
                  'FrameObject properties must always provide "self._frame" as '
                  'the "frame" parameter to functions such as "stbt.match".'),
        'E7005': ('Image "%s" not committed to git',
                  'stbt-uncommitted-image',
                  'The image path given to "stbt.match" '
                  '(and similar functions) exists on disk, '
                  "but isn't committed to git."),
        'E7006': ('FrameObject properties must use "self._frame", not '
                  '"get_frame()"',
                  'stbt-frame-object-get-frame',
                  'FrameObject properties must use "self._frame", not '
                  '"stbt.get_frame()".'),
        'E7007': ('FrameObject properties must not use "%s"',
                  'stbt-frame-object-property-press',
                  'FrameObject properties must not have side-effects that '
                  'change the state of the device-under-test by calling '
                  '"stbt.press()" or "stbt.press_and_wait()".'),
        'E7008': ('"assert True" has no effect',
                  'stbt-assert-true',
                  '"assert True" has no effect; consider replacing it with a '
                  'comment or a call to "logging.info()".'),
    }

    def visit_const(self, node):
        if (isinstance(node.value, string_types) and
                re.search(r'.+\.png$', node.value) and
                "\n" not in node.value and
                not _is_uri(node.value) and
                not _is_calculated_value(node) and
                not _is_pattern_value(node) and
                not _is_whitelisted_name(node.value) and
                not _in_whitelisted_functions(node)):
            path = _find_file(node.value, node)
            if os.path.isfile(path):
                if _is_file_uncommitted(path):
                    self.add_message('E7005', node=node,
                                     args=os.path.relpath(path))
            else:
                self.add_message('E7001', node=node, args=os.path.relpath(path))

    def visit_call(self, node):
        if re.search(r"\b(is_screen_black|match|match_text|ocr|press_and_wait|"
                     r"wait_until)$",
                     node.func.as_string()):
            if isinstance(node.parent, Expr):
                for inferred in _infer(node.func):
                    if inferred.root().name in (
                            '_stbt.black', '_stbt.match', '_stbt.ocr',
                            '_stbt.transition', '_stbt.wait'):
                        self.add_message(
                            'E7002', node=node, args=node.func.as_string())

        if re.search(r"\bwait_until$", node.func.as_string()):
            if node.args:
                arg = node.args[0]
                if not _is_callable(arg):
                    self.add_message('E7003', node=node, args=arg.as_string())

        if _in_frameobject(node) and _in_property(node):
            if re.search(r"\bget_frame$", node.func.as_string()):
                self.add_message('E7006', node=node)

            if re.search(
                    r"\b(press|press_and_wait|pressing|press_until_match)$",
                    node.func.as_string()):
                self.add_message('E7007', node=node, args=node.func.as_string())

            for funcdef in _infer(node.func):
                argnames = _get_argnames(funcdef)
                if "frame" in argnames:
                    index = argnames.index("frame")
                    args = [a for a in node.args if not isinstance(a, Keyword)]
                    if hasattr(node, "keywords"):  # astroid >= 1.4
                        keywords = node.keywords or []
                        kwargs = [k.arg for k in keywords]
                    else:  # astroid < 1.4
                        kwargs = [a.arg for a in node.args
                                  if isinstance(a, Keyword)]
                    if len(args) <= index and "frame" not in kwargs:
                        self.add_message('E7004', node=node,
                                         args=node.as_string())

    def visit_assert(self, assertion):
        if isinstance(assertion.test, Const) and assertion.test.value is True:
            if assertion.fail:
                self.add_message("E7008", node=assertion)
            else:
                self.add_message("E7008", node=assertion)

def _transform_assert_false_into_raise(assertion):
    if isinstance(assertion.test, Const) and assertion.test.value is False:
        out = Raise(lineno=assertion.lineno,
                    col_offset=assertion.col_offset,
                    parent=assertion.parent)
        exc = Call(parent=out)
        if assertion.fail:
            args = [assertion.fail]
            args[0].parent = exc
        else:
            args = []
        exc.postinit(Name("AssertionError", parent=exc), args)
        out.postinit(exc, None)
        return out

MANAGER.register_transform(Assert, _transform_assert_false_into_raise)

def _is_callable(node):
    failed_to_infer = True
    for inferred in _infer(node):
        failed_to_infer = False
        if inferred.callable():
            return True
    if failed_to_infer:
        if (isinstance(node, Call) and
                _is_function_named(node.func, "functools.partial")):
            return True
    return False

def _in_frameobject(node):
    while node is not None:
        if isinstance(node, ClassDef):
            if "_stbt.frameobject.FrameObject" in [
                    base.qname() for base in node.ancestors()]:
                return True
        node = node.parent
    return False

def _in_property(node):
    while node is not None:
        if isinstance(node, FunctionDef):
            if ("__builtin__.property" in node.decoratornames() or
                    "builtins.property" in node.decoratornames()):
                return True
        node = node.parent
    return False

def _get_argnames(node):
    if isinstance(node, FunctionDef):
        return node.argnames()
    if isinstance(node, ClassDef) and node.newstyle:
        for method in node.methods():  # pylint:disable=redefined-outer-name
            if method.name == "__init__":
                return method.argnames()[1:]
    return []

def _is_uri(filename):
    return re.search(r"^[a-zA-Z][a-zA-Z0-9+.-]*:", filename)

def _is_calculated_value(node):
    return (
        isinstance(node.parent, BinOp) or
        (isinstance(node.parent, Call) and
         node.parent.func.as_string().split(".")[-1] in ("join", "replace")))

def _is_pattern_value(node):
    return re.search(r'\*', node.value)

def _is_whitelisted_name(filename):
    return filename == 'screenshot.png'

def _in_whitelisted_functions(node):
    return (
        isinstance(node.parent, Call) and
        any(_is_function_named(node.parent.func, x) for x in (
            "cv2.imwrite",
            "imwrite",  # handles "from cv2 import imwrite" with OpenCV 3.x
            "re.match",
            "re.search",
            "re.sub",
            "stbt.save_frame",
            "_stbt.imgutils.save_frame",  # "from stbt import save_frame"
        )))

def _is_function_named(func, name):
    if func.as_string() == name:
        return True
    for funcdef in _infer(func):
        if (isinstance(funcdef, FunctionDef) and
                ".".join((funcdef.parent.name, funcdef.name)) == name):
            return True
    return False

def _find_file(filename, node):
    """Resolves `filename` on stbt's image search path

    See `stbt.load_image` for stbt's image lookup algorithm.
    """
    return os.path.join(
        os.path.dirname(node.root().file),
        filename)

def _is_file_uncommitted(filename):
    if not _in_git_repo():
        return False
    try:
        subprocess.check_output(
            ["git", "ls-files", "--error-unmatch", filename],
            stderr=subprocess.STDOUT)
        return False
    except subprocess.CalledProcessError:
        return True

__in_git_repo = None

def _in_git_repo():
    global __in_git_repo
    if __in_git_repo is None:
        try:
            subprocess.check_output(["git", "rev-parse", "--show-toplevel"],
                                    stderr=subprocess.STDOUT)
            __in_git_repo = True
        except (OSError, subprocess.CalledProcessError):
            __in_git_repo = False
    return __in_git_repo

def _infer(node):
    try:
        for inferred in node.infer():
            # Sometimes when `infer()` fails it returns `YES` which returns
            # True to everything, including `callable()` and
            # `isinstance(YES, <anything else>)`, so it isn't useful.
            if inferred == YES:
                continue
            yield inferred
    except Exception:  # pylint:disable=broad-except
        pass

def register(linter):
    linter.register_checker(StbtChecker(linter))
<commit_msg>pylint_plugin: Support pylint 2 / astroid 2<commit_after>"""pylint plugin to do static analysis on stbt scripts

Used by "stbt lint".

Documentation on Abstract Syntax Tree traversal with python/pylint:

* http://docs.pylint.org/extend.html#writing-your-own-checker
* http://hg.logilab.org/review/pylint/file/default/examples/custom.py
* http://docs.python.org/2/library/ast.html

"""
from __future__ import unicode_literals
from __future__ import print_function
from __future__ import division
from __future__ import absolute_import
from builtins import *  # pylint:disable=redefined-builtin,unused-wildcard-import,wildcard-import,wrong-import-order
from future.utils import string_types

import os
import re
import subprocess

from astroid import (
    Assert, BinOp, Call, ClassDef, Const, Expr, FunctionDef, Keyword, MANAGER,
    Name, Raise, Uninferable)
from pylint.checkers import BaseChecker
from pylint.interfaces import IAstroidChecker

class StbtChecker(BaseChecker):
    __implements__ = IAstroidChecker
    name = 'stb-tester'
    msgs = {
        # Range 70xx reserved for custom checkers: www.logilab.org/ticket/68057
        # When you add a new checker update the docstring in ../stbt_lint.py
        'E7001': ('Image "%s" not found on disk',
                  'stbt-missing-image',
                  'The image path given to "stbt.match" '
                  '(and similar functions) does not exist on disk.'),
        'E7002': ('"%s" return value not used (missing "assert"?)',
                  'stbt-unused-return-value',
                  "This function does not raise an exception on failure but "
                  "you aren't using its return value. Perhaps you've forgotten "
                  'to use "assert".'),
        'E7003': ('"wait_until" argument "%s" isn\'t callable',
                  'stbt-wait-until-callable',
                  'The argument given to "wait_until" must be a callable '
                  '(such as a function or a lambda expression).'),
        'E7004': ('"%s" missing "frame" argument',
                  'stbt-frame-object-missing-frame',
                  'FrameObject properties must always provide "self._frame" as '
                  'the "frame" parameter to functions such as "stbt.match".'),
        'E7005': ('Image "%s" not committed to git',
                  'stbt-uncommitted-image',
                  'The image path given to "stbt.match" '
                  '(and similar functions) exists on disk, '
                  "but isn't committed to git."),
        'E7006': ('FrameObject properties must use "self._frame", not '
                  '"get_frame()"',
                  'stbt-frame-object-get-frame',
                  'FrameObject properties must use "self._frame", not '
                  '"stbt.get_frame()".'),
        'E7007': ('FrameObject properties must not use "%s"',
                  'stbt-frame-object-property-press',
                  'FrameObject properties must not have side-effects that '
                  'change the state of the device-under-test by calling '
                  '"stbt.press()" or "stbt.press_and_wait()".'),
        'E7008': ('"assert True" has no effect',
                  'stbt-assert-true',
                  '"assert True" has no effect; consider replacing it with a '
                  'comment or a call to "logging.info()".'),
    }

    def visit_const(self, node):
        if (isinstance(node.value, string_types) and
                re.search(r'.+\.png$', node.value) and
                "\n" not in node.value and
                not _is_uri(node.value) and
                not _is_calculated_value(node) and
                not _is_pattern_value(node) and
                not _is_whitelisted_name(node.value) and
                not _in_whitelisted_functions(node)):
            path = _find_file(node.value, node)
            if os.path.isfile(path):
                if _is_file_uncommitted(path):
                    self.add_message('E7005', node=node,
                                     args=os.path.relpath(path))
            else:
                self.add_message('E7001', node=node, args=os.path.relpath(path))

    def visit_call(self, node):
        if re.search(r"\b(is_screen_black|match|match_text|ocr|press_and_wait|"
                     r"wait_until)$",
                     node.func.as_string()):
            if isinstance(node.parent, Expr):
                for inferred in _infer(node.func):
                    if inferred.root().name in (
                            '_stbt.black', '_stbt.match', '_stbt.ocr',
                            '_stbt.transition', '_stbt.wait'):
                        self.add_message(
                            'E7002', node=node, args=node.func.as_string())

        if re.search(r"\bwait_until$", node.func.as_string()):
            if node.args:
                arg = node.args[0]
                if not _is_callable(arg):
                    self.add_message('E7003', node=node, args=arg.as_string())

        if _in_frameobject(node) and _in_property(node):
            if re.search(r"\bget_frame$", node.func.as_string()):
                self.add_message('E7006', node=node)

            if re.search(
                    r"\b(press|press_and_wait|pressing|press_until_match)$",
                    node.func.as_string()):
                self.add_message('E7007', node=node, args=node.func.as_string())

            for funcdef in _infer(node.func):
                argnames = _get_argnames(funcdef)
                if "frame" in argnames:
                    index = argnames.index("frame")
                    args = [a for a in node.args if not isinstance(a, Keyword)]
                    if hasattr(node, "keywords"):  # astroid >= 1.4
                        keywords = node.keywords or []
                        kwargs = [k.arg for k in keywords]
                    else:  # astroid < 1.4
                        kwargs = [a.arg for a in node.args
                                  if isinstance(a, Keyword)]
                    if len(args) <= index and "frame" not in kwargs:
                        self.add_message('E7004', node=node,
                                         args=node.as_string())

    def visit_assert(self, assertion):
        if isinstance(assertion.test, Const) and assertion.test.value is True:
            if assertion.fail:
                self.add_message("E7008", node=assertion)
            else:
                self.add_message("E7008", node=assertion)

def _transform_assert_false_into_raise(assertion):
    if isinstance(assertion.test, Const) and assertion.test.value is False:
        out = Raise(lineno=assertion.lineno,
                    col_offset=assertion.col_offset,
                    parent=assertion.parent)
        exc = Call(parent=out)
        if assertion.fail:
            args = [assertion.fail]
            args[0].parent = exc
        else:
            args = []
        exc.postinit(Name("AssertionError", parent=exc), args)
        out.postinit(exc, None)
        return out

MANAGER.register_transform(Assert, _transform_assert_false_into_raise)

def _is_callable(node):
    failed_to_infer = True
    for inferred in _infer(node):
        failed_to_infer = False
        if inferred.callable():
            return True
    if failed_to_infer:
        if (isinstance(node, Call) and
                _is_function_named(node.func, "functools.partial")):
            return True
    return False

def _in_frameobject(node):
    while node is not None:
        if isinstance(node, ClassDef):
            if "_stbt.frameobject.FrameObject" in [
                    base.qname() for base in node.ancestors()]:
                return True
        node = node.parent
    return False

def _in_property(node):
    while node is not None:
        if isinstance(node, FunctionDef):
            if ("__builtin__.property" in node.decoratornames() or
                    "builtins.property" in node.decoratornames()):
                return True
        node = node.parent
    return False

def _get_argnames(node):
    if isinstance(node, FunctionDef):
        return node.argnames()
    if isinstance(node, ClassDef) and node.newstyle:
        for method in node.methods():  # pylint:disable=redefined-outer-name
            if method.name == "__init__":
                return method.argnames()[1:]
    return []

def _is_uri(filename):
    return re.search(r"^[a-zA-Z][a-zA-Z0-9+.-]*:", filename)

def _is_calculated_value(node):
    return (
        isinstance(node.parent, BinOp) or
        (isinstance(node.parent, Call) and
         node.parent.func.as_string().split(".")[-1] in ("join", "replace")))

def _is_pattern_value(node):
    return re.search(r'\*', node.value)

def _is_whitelisted_name(filename):
    return filename == 'screenshot.png'

def _in_whitelisted_functions(node):
    return (
        isinstance(node.parent, Call) and
        any(_is_function_named(node.parent.func, x) for x in (
            "cv2.imwrite",
            "imwrite",  # handles "from cv2 import imwrite" with OpenCV 3.x
            "re.match",
            "re.search",
            "re.sub",
            "stbt.save_frame",
            "_stbt.imgutils.save_frame",  # "from stbt import save_frame"
        )))

def _is_function_named(func, name):
    if func.as_string() == name:
        return True
    for funcdef in _infer(func):
        if (isinstance(funcdef, FunctionDef) and
                ".".join((funcdef.parent.name, funcdef.name)) == name):
            return True
    return False

def _find_file(filename, node):
    """Resolves `filename` on stbt's image search path

    See `stbt.load_image` for stbt's image lookup algorithm.
    """
    return os.path.join(
        os.path.dirname(node.root().file),
        filename)

def _is_file_uncommitted(filename):
    if not _in_git_repo():
        return False
    try:
        subprocess.check_output(
            ["git", "ls-files", "--error-unmatch", filename],
            stderr=subprocess.STDOUT)
        return False
    except subprocess.CalledProcessError:
        return True

__in_git_repo = None

def _in_git_repo():
    global __in_git_repo
    if __in_git_repo is None:
        try:
            subprocess.check_output(["git", "rev-parse", "--show-toplevel"],
                                    stderr=subprocess.STDOUT)
            __in_git_repo = True
        except (OSError, subprocess.CalledProcessError):
            __in_git_repo = False
    return __in_git_repo

def _infer(node):
    try:
        for inferred in node.infer():
            # When `infer()` fails it returns this singleton:
            if inferred is Uninferable:
                continue
            yield inferred
    except Exception:  # pylint:disable=broad-except
        pass

def register(linter):
    linter.register_checker(StbtChecker(linter))
<|endoftext|>


<commit_before>#!/usr/bin/python
"""
Allows using any of the stbt remote control backends remotely using the lirc
protocol.

Presents the same socket protocol as lircd but sending keypresses using any of
stbt's controls.  This allows for example controlling a roku over its HTTP
interface from some software that only speaks lirc.

Example usage:

    $ stbt control-relay file:example

Listens on `/var/run/lirc/lircd` for lirc clients.  Keypress sent will be
written to the file example.  So

    $ irsend SEND_ONCE stbt KEY_UP

Will write the text "KEY_UP" to the file `example`.

    $ stbt control-relay --socket=lircd.sock roku:192.168.1.13

Listens on lircd.sock and will forward keypresses to the roku at 192.168.1.13
using its HTTP protocol.  So

    $ irsend -d lircd.sock SEND_ONCE stbt KEY_OK

Will press KEY_OK on the roku device.
"""
from __future__ import unicode_literals
from __future__ import print_function
from __future__ import division
from __future__ import absolute_import
from builtins import *  # pylint:disable=redefined-builtin,unused-wildcard-import,wildcard-import,wrong-import-order
import argparse
import logging
import os
import re
import signal
import socket
import sys

from _stbt.control import uri_to_control
from _stbt.utils import to_bytes

def main(argv):
    parser = argparse.ArgumentParser(
        epilog=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
    parser.add_argument(
        "--socket", default="/var/run/lirc/lircd", help="""LIRC socket to read
        remote control presses from (defaults to %(default)s).""")
    parser.add_argument("-v", "--verbose", action="store_true")
    parser.add_argument("output", help="""Remote control configuration to
        transmit on. Values are the same as stbt run's --control.""")
    args = parser.parse_args(argv[1:])

    logging.basicConfig(
        format="%(levelname)s: %(message)s",
        level=logging.DEBUG if args.verbose else logging.INFO)

    signal.signal(signal.SIGTERM, lambda _signo, _stack_frame: sys.exit(0))

    if os.environ.get('LISTEN_FDS') == '1' and \
            os.environ.get('LISTEN_PID') == str(os.getpid()):
        s = socket.fromfd(3, socket.AF_UNIX, socket.SOCK_STREAM)
    else:
        s = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM)
        s.bind(args.socket)
        s.listen(5)

    control = uri_to_control(args.output)

    logging.info("stbt-control-relay started up with output '%s'", args.output)

    while True:
        conn, _ = s.accept()
        f = conn.makefile('rb', 0)
        while True:
            cmd = f.readline()
            if not cmd:
                break
            cmd = cmd.rstrip(b"\n")
            m = re.match(br"(?P<action>SEND_ONCE|SEND_START|SEND_STOP) "
                         br"(?P<ctrl>\S+) (?P<key>\S+)", cmd)
            if not m:
                logging.error("Invalid command: %s", cmd)
                send_response(conn, cmd, success=False,
                              data=b"Invalid command: %s" % cmd)
                continue
            action = m.group("action")
            key = m.group("key")
            logging.debug("Received %s %s", action, key)
            try:
                if action == b"SEND_ONCE":
                    control.press(key)
                elif action == b"SEND_START":
                    control.keydown(key)
                elif action == b"SEND_STOP":
                    control.keyup(key)
            except Exception as e:  # pylint: disable=broad-except
                logging.error("Error pressing key %r: %s", key, e,
                              exc_info=True)
                send_response(conn, cmd, success=False, data=to_bytes(str(e)))
                continue
            send_response(conn, cmd, success=True)

def send_response(sock, request, success, data=b""):
    # See http://www.lirc.org/html/lircd.html
    message = b"BEGIN\n%s\n%s\n" % (
        request,
        b"SUCCESS" if success else b"ERROR")
    if data:
        data = data.split(b"\n")
        message += b"DATA\n%d\n%s\n" % (len(data), b"\n".join(data))
    message += b"END\n"

    try:
        sock.sendall(message)
    except Exception:  # pylint: disable=broad-except
        pass

if __name__ == "__main__":
    sys.exit(main(sys.argv))
<commit_msg>Python3: stbt-control-relay: Pass a unicode keyname to `press()`<commit_after>#!/usr/bin/python
"""
Allows using any of the stbt remote control backends remotely using the lirc
protocol.

Presents the same socket protocol as lircd but sending keypresses using any of
stbt's controls.  This allows for example controlling a roku over its HTTP
interface from some software that only speaks lirc.

Example usage:

    $ stbt control-relay file:example

Listens on `/var/run/lirc/lircd` for lirc clients.  Keypress sent will be
written to the file example.  So

    $ irsend SEND_ONCE stbt KEY_UP

Will write the text "KEY_UP" to the file `example`.

    $ stbt control-relay --socket=lircd.sock roku:192.168.1.13

Listens on lircd.sock and will forward keypresses to the roku at 192.168.1.13
using its HTTP protocol.  So

    $ irsend -d lircd.sock SEND_ONCE stbt KEY_OK

Will press KEY_OK on the roku device.
"""
from __future__ import unicode_literals
from __future__ import print_function
from __future__ import division
from __future__ import absolute_import
from builtins import *  # pylint:disable=redefined-builtin,unused-wildcard-import,wildcard-import,wrong-import-order
import argparse
import logging
import os
import re
import signal
import socket
import sys

from _stbt.control import uri_to_control
from _stbt.utils import to_bytes

def main(argv):
    parser = argparse.ArgumentParser(
        epilog=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
    parser.add_argument(
        "--socket", default="/var/run/lirc/lircd", help="""LIRC socket to read
        remote control presses from (defaults to %(default)s).""")
    parser.add_argument("-v", "--verbose", action="store_true")
    parser.add_argument("output", help="""Remote control configuration to
        transmit on. Values are the same as stbt run's --control.""")
    args = parser.parse_args(argv[1:])

    logging.basicConfig(
        format="%(levelname)s: %(message)s",
        level=logging.DEBUG if args.verbose else logging.INFO)

    signal.signal(signal.SIGTERM, lambda _signo, _stack_frame: sys.exit(0))

    if os.environ.get('LISTEN_FDS') == '1' and \
            os.environ.get('LISTEN_PID') == str(os.getpid()):
        s = socket.fromfd(3, socket.AF_UNIX, socket.SOCK_STREAM)
    else:
        s = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM)
        s.bind(args.socket)
        s.listen(5)

    control = uri_to_control(args.output)

    logging.info("stbt-control-relay started up with output '%s'", args.output)

    while True:
        conn, _ = s.accept()
        f = conn.makefile('rb', 0)
        while True:
            cmd = f.readline()
            if not cmd:
                break
            cmd = cmd.rstrip(b"\n")
            m = re.match(br"(?P<action>SEND_ONCE|SEND_START|SEND_STOP) "
                         br"(?P<ctrl>\S+) (?P<key>\S+)", cmd)
            if not m:
                logging.error("Invalid command: %s", cmd)
                send_response(conn, cmd, success=False,
                              data=b"Invalid command: %s" % cmd)
                continue
            action = m.group("action")
            key = m.group("key")
            logging.debug("Received %s %s", action, key)
            try:
                key = key.decode("utf-8")
                if action == b"SEND_ONCE":
                    control.press(key)
                elif action == b"SEND_START":
                    control.keydown(key)
                elif action == b"SEND_STOP":
                    control.keyup(key)
            except Exception as e:  # pylint: disable=broad-except
                logging.error("Error pressing key %r: %s", key, e,
                              exc_info=True)
                send_response(conn, cmd, success=False, data=to_bytes(str(e)))
                continue
            send_response(conn, cmd, success=True)

def send_response(sock, request, success, data=b""):
    # See http://www.lirc.org/html/lircd.html
    message = b"BEGIN\n%s\n%s\n" % (
        request,
        b"SUCCESS" if success else b"ERROR")
    if data:
        data = data.split(b"\n")
        message += b"DATA\n%d\n%s\n" % (len(data), b"\n".join(data))
    message += b"END\n"

    try:
        sock.sendall(message)
    except Exception:  # pylint: disable=broad-except
        pass

if __name__ == "__main__":
    sys.exit(main(sys.argv))
<|endoftext|>