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

import os
from ConfigManager import *

c = config = ConfigManager()

#
# General
#
general = c.add_section('general')

general['renderer'] = StringOption(
    """ Renderer to use for conversion """,
    options = '--renderer',
    default = 'XHTML',
)

general['theme'] = StringOption(
    """ Theme for the renderer to use """,
    options = '--theme',
    default = 'default',
)

general['copy-theme-extras'] = BooleanOption(
    """  Copy files associated with the theme to the output directory """,
    options = '--copy-theme-extras !--no-theme-extras',
    default = True,
)

general['kpsewhich'] = StringOption(
    """ Program which locates LaTeX files and packages """,
    options = '--kpsewhich',
    default = 'kpsewhich',
)

def readconfig(file):
    """ Read a configuration file """
    if not os.path.isfile(file):
        print >>sys.stderr, "WARNING: Could not load config file '%s'" % file
        return
    config.read(file)

general['config'] = StringOption(
    """
    Load additional configuration file

    This configuration file will be loaded during the processing of
    command-line options.  However, configuration files cannot override
    command-line options.  So no matter what is in the configuration file,
    if an option is specified on the command-line, it will always win.
    To eliminate any confusion, this should generally be the first
    argument.

    """,
    options = '--config -c',
    callback = readconfig,
)

#
# Links
#

links = c.add_section('links')

def getstring(data):
    """ Return a string which may or may not be quoted """
    data = list(data.lstrip())

    if not data:
        return ''

    s = []
    for quote in '\'"':
        if data[0] == quote:
            data.pop(0)
            while data:
                if data[0] == quote:
                    data.pop(0)
                    break
                s.append(data.pop(0))
            return ''.join(s), ''.join(data)

    while data:
        if data[0] != ' ':
            s.append(data.pop(0))
            continue
        break
    return ''.join(s), ''.join(data)

def setlinks(data):
    """ Set links in the configuration """
    data = data[1:-1].strip()
    key, data = getstring(data)
    url, data = getstring(data)
    title, data = getstring(data)
    if not title:
        title = url
        url = ''
    if title:
        links['%s-title' % key] = StringOption(default=title)
    if url:
        links['%s-url' % key] = StringOption(default=url)

links[';links'] = CompoundOption(
    """ Set links for use in navigation """,
    options = '--link',
    category = 'document',
    callback = setlinks,
)

#
# Counters
#

counters = c.add_section('counters')

def setcounter(data):
    """ Set counters in the configuration """
    value = [x for x in data[1:-1].split() if x]
    if len(value) == 2:
        counters[value[0]] = IntegerOption(default=int(value[1]))

counters[';counters'] = CompoundOption(
    """ Set initial counter values """,
    options = '--counter',
    category = 'document',
    callback = setcounter,
)

#
# Files
#

files = c.add_section('files')
c.add_category('files', 'File Handling Options')

files['input-encoding'] = StringOption(
    """ Input file encoding """,
    options = '--input-encoding',
    default = 'utf-8',
    category = 'files',
)

files['output-encoding'] = StringOption(
    """ Output file encoding """,
    options = '--output-encoding',
    default = 'utf-8',
    category = 'files',
)

files['escape-high-chars'] = BooleanOption(
    """ Escape characters that are higher than 7-bit """,
    options = '--escape-high-chars',
    default = False,
    category = 'files',
)

files['split-level'] = IntegerOption(
    """ Highest section level that generates a new file """,
    options = '--split-level',
    default = 2,
    category = 'files',
)

files['filename'] = StringOption(
    """ Template for output filenames """,
    options = '--filename',
    default = 'index [$id, sect$num(4)]',
    category = 'files',
)

files['bad-chars'] = StringOption(
    """ Characters that should not be allowed in a filename """,
    options = '--bad-filename-chars',
    default = ': #$%^&*!~`"\'=?/{}[]()|<>;\\,.',
    category = 'files',
)

files['bad-chars-sub'] = StringOption(
    """ Character that should be used instead of an illegal character """,
    options = '--bad-filename-chars-sub',
    default = '-',
    category = 'files',
)

files['directory'] = StringOption(
    """ Directory to put output files into """,
    options = '--dir -d',
    category = 'files',
    default = '$jobname',
)

#
# Images
#

images = c.add_section('images')
c.add_category('images', 'Image Generation Options')

images['base-url'] = StringOption(
    """ Base URL for all images """,
    options = '--image-base-url',
    category = 'images',
)

images['enabled'] = BooleanOption(
    """ Enable image generation """,
    options = '--enable-images !--disable-images',
    default = True,
    category = 'images',
)

images['imager'] = StringOption(
    """ LaTeX to image program """,
    options = '--imager',
    default = 'dvipng dvi2bitmap gsdvipng gspdfpng OSXCoreGraphics',
    category = 'images',
)

images['vector-imager'] = StringOption(
    """ LaTeX to vector image program """,
    options = '--vector-imager',
    default = 'none dvisvgm',
    category = 'images',
)

images['filenames'] = StringOption(
    """ Template for image filenames """,
    options = '--image-filenames',
    default = 'images/img-$num(4)',
    category = 'images',
)

images['baseline-padding'] = IntegerOption(
    """ Amount to pad the image below the baseline """,
    options = '--image-baseline-padding',
    default = 0,
    category = 'images',
)

images['compiler'] = StringOption(
    """ LaTeX command to use when compiling image document """,
    options = '--image-compiler',
    category = 'images',
)

images['cache'] = BooleanOption(
    """  Enable image caching between runs """,
    options = '--enable-image-cache !--disable-image-cache',
    default = True,
    category = 'images',
)

#
# Document
#

doc = c.add_section('document')
c.add_category('document', 'Document Options')

doc['base-url'] = StringOption(
    """ Base URL for inter-node links """,
    options = '--base-url',
    category = 'document',
)

doc['title'] = StringOption(
    """
    Title for the document

    This option specifies a title to use instead of the title
    specified in the LaTeX document.

    """,
    options = '--title',
    category = 'document',
)

doc['toc-depth'] = IntegerOption(
    """ Number of levels to display in the table of contents """,
    options = '--toc-depth',
    category = 'document',
    default = 3,
)

doc['toc-non-files'] = BooleanOption(
    """ Display sections that do not create files in the table of contents """,
    options = '--toc-non-files',
    category = 'document',
    default = False,
)

doc['sec-num-depth'] = IntegerOption(
    """ Maximum section depth to display section numbers """,
    options = '--sec-num-depth',
    category = 'document',
    default = 2,
)

doc['index-columns'] = IntegerOption(
    """ Number of columns to split the index entries into """,
    options = '--index-columns',
    category = 'document',
    default = 2,
)

config.read('~/.plasTeXrc')
config.read('/usr/local/etc/plasTeXrc')
config.read(os.path.join(os.path.dirname(__file__), 'plasTeXrc'))

del c
<commit_msg>Added support for Windows in the imager<commit_after>#!/usr/bin/env python

import os
from ConfigManager import *

c = config = ConfigManager()

#
# General
#
general = c.add_section('general')

general['renderer'] = StringOption(
    """ Renderer to use for conversion """,
    options = '--renderer',
    default = 'XHTML',
)

general['theme'] = StringOption(
    """ Theme for the renderer to use """,
    options = '--theme',
    default = 'default',
)

general['copy-theme-extras'] = BooleanOption(
    """  Copy files associated with the theme to the output directory """,
    options = '--copy-theme-extras !--no-theme-extras',
    default = True,
)

general['kpsewhich'] = StringOption(
    """ Program which locates LaTeX files and packages """,
    options = '--kpsewhich',
    default = 'kpsewhich',
)

def readconfig(file):
    """ Read a configuration file """
    if not os.path.isfile(file):
        print >>sys.stderr, "WARNING: Could not load config file '%s'" % file
        return
    config.read(file)

general['config'] = StringOption(
    """
    Load additional configuration file

    This configuration file will be loaded during the processing of
    command-line options.  However, configuration files cannot override
    command-line options.  So no matter what is in the configuration file,
    if an option is specified on the command-line, it will always win.
    To eliminate any confusion, this should generally be the first
    argument.

    """,
    options = '--config -c',
    callback = readconfig,
)

#
# Links
#

links = c.add_section('links')

def getstring(data):
    """ Return a string which may or may not be quoted """
    data = list(data.lstrip())

    if not data:
        return ''

    s = []
    for quote in '\'"':
        if data[0] == quote:
            data.pop(0)
            while data:
                if data[0] == quote:
                    data.pop(0)
                    break
                s.append(data.pop(0))
            return ''.join(s), ''.join(data)

    while data:
        if data[0] != ' ':
            s.append(data.pop(0))
            continue
        break
    return ''.join(s), ''.join(data)

def setlinks(data):
    """ Set links in the configuration """
    data = data[1:-1].strip()
    key, data = getstring(data)
    url, data = getstring(data)
    title, data = getstring(data)
    if not title:
        title = url
        url = ''
    if title:
        links['%s-title' % key] = StringOption(default=title)
    if url:
        links['%s-url' % key] = StringOption(default=url)

links[';links'] = CompoundOption(
    """ Set links for use in navigation """,
    options = '--link',
    category = 'document',
    callback = setlinks,
)

#
# Counters
#

counters = c.add_section('counters')

def setcounter(data):
    """ Set counters in the configuration """
    value = [x for x in data[1:-1].split() if x]
    if len(value) == 2:
        counters[value[0]] = IntegerOption(default=int(value[1]))

counters[';counters'] = CompoundOption(
    """ Set initial counter values """,
    options = '--counter',
    category = 'document',
    callback = setcounter,
)

#
# Files
#

files = c.add_section('files')
c.add_category('files', 'File Handling Options')

files['input-encoding'] = StringOption(
    """ Input file encoding """,
    options = '--input-encoding',
    default = 'utf-8',
    category = 'files',
)

files['output-encoding'] = StringOption(
    """ Output file encoding """,
    options = '--output-encoding',
    default = 'utf-8',
    category = 'files',
)

files['escape-high-chars'] = BooleanOption(
    """ Escape characters that are higher than 7-bit """,
    options = '--escape-high-chars',
    default = False,
    category = 'files',
)

files['split-level'] = IntegerOption(
    """ Highest section level that generates a new file """,
    options = '--split-level',
    default = 2,
    category = 'files',
)

files['filename'] = StringOption(
    """ Template for output filenames """,
    options = '--filename',
    default = 'index [$id, sect$num(4)]',
    category = 'files',
)

files['bad-chars'] = StringOption(
    """ Characters that should not be allowed in a filename """,
    options = '--bad-filename-chars',
    default = ': #$%^&*!~`"\'=?/{}[]()|<>;\\,.',
    category = 'files',
)

files['bad-chars-sub'] = StringOption(
    """ Character that should be used instead of an illegal character """,
    options = '--bad-filename-chars-sub',
    default = '-',
    category = 'files',
)

files['directory'] = StringOption(
    """ Directory to put output files into """,
    options = '--dir -d',
    category = 'files',
    default = '$jobname',
)

#
# Images
#

images = c.add_section('images')
c.add_category('images', 'Image Generation Options')

images['base-url'] = StringOption(
    """ Base URL for all images """,
    options = '--image-base-url',
    category = 'images',
)

images['enabled'] = BooleanOption(
    """ Enable image generation """,
    options = '--enable-images !--disable-images',
    default = True,
    category = 'images',
)

images['imager'] = StringOption(
    """ LaTeX to image program """,
    options = '--imager',
    default = 'dvipng dvi2bitmap gspdfpng gsdvipng OSXCoreGraphics',
    category = 'images',
)

images['vector-imager'] = StringOption(
    """ LaTeX to vector image program """,
    options = '--vector-imager',
    default = 'none dvisvgm',
    category = 'images',
)

images['filenames'] = StringOption(
    """ Template for image filenames """,
    options = '--image-filenames',
    default = 'images/img-$num(4)',
    category = 'images',
)

images['baseline-padding'] = IntegerOption(
    """ Amount to pad the image below the baseline """,
    options = '--image-baseline-padding',
    default = 0,
    category = 'images',
)

images['compiler'] = StringOption(
    """ LaTeX command to use when compiling image document """,
    options = '--image-compiler',
    category = 'images',
)

images['cache'] = BooleanOption(
    """  Enable image caching between runs """,
    options = '--enable-image-cache !--disable-image-cache',
    default = True,
    category = 'images',
)

#
# Document
#

doc = c.add_section('document')
c.add_category('document', 'Document Options')

doc['base-url'] = StringOption(
    """ Base URL for inter-node links """,
    options = '--base-url',
    category = 'document',
)

doc['title'] = StringOption(
    """
    Title for the document

    This option specifies a title to use instead of the title
    specified in the LaTeX document.

    """,
    options = '--title',
    category = 'document',
)

doc['toc-depth'] = IntegerOption(
    """ Number of levels to display in the table of contents """,
    options = '--toc-depth',
    category = 'document',
    default = 3,
)

doc['toc-non-files'] = BooleanOption(
    """ Display sections that do not create files in the table of contents """,
    options = '--toc-non-files',
    category = 'document',
    default = False,
)

doc['sec-num-depth'] = IntegerOption(
    """ Maximum section depth to display section numbers """,
    options = '--sec-num-depth',
    category = 'document',
    default = 2,
)

doc['index-columns'] = IntegerOption(
    """ Number of columns to split the index entries into """,
    options = '--index-columns',
    category = 'document',
    default = 2,
)

config.read('~/.plasTeXrc')
config.read('/usr/local/etc/plasTeXrc')
config.read(os.path.join(os.path.dirname(__file__), 'plasTeXrc'))

del c
<|endoftext|>


<commit_before>"""
    This module implements the request manager functionality.

    Job Queue and Processing
    ^^^^^^^^^^^^^^^^^^^^^^^^

    As requests are issued via http to the API a process queue will store all
    active jobs. Processes will be created and assume one of the following
    states throughout their existence: ::

        * 'pending' - The request has yet to be begin being processed
        * 'running' - The request is being processed
        * 'success' - The request has finished processing and is exposed at
            the url
        * 'failure' - The result has finished processing but dailed to expose
            results

    When a process a request is received and a job is created to service that
    request it enters the 'pending' state. If the job returns without
    exception it enters the 'success' state, otherwise it enters the 'failure'
    state.  The job remains in either of these states until it is cleared
    from the process queue.

    Response Data
    ^^^^^^^^^^^^^

    As requests are made to the API the data generated and formatted as JSON.
    The definition of is as follows: ::

        {   header : header_list,
            cohort_expr : cohort_gen_timestamp : metric : timeseries :
            aggregator : date_start : date_end : [ metric_param : ]* : data
        }

    Where each component is defined: ::

        header_str := list(str), list of header values
        cohort_expr := str, cohort ID expression
        cohort_gen_timestamp := str, cohort generation timestamp (earliest of
            all cohorts in expression)
        metric := str, user metric handle
        timeseries := boolean, indicates if this is a timeseries
        aggregator := str, aggregator used
        date_start := str, start datetime of request
        date_end := str, end datetime of request
        metric_param := -, optional metric parameters
        data := list(tuple), set of data points

    Request data is mapped to a query via metric objects and hashed in the
    dictionary `pkl_data`.

"""

__author__ = {
    "ryan faulkner": "rfaulkner@wikimedia.org"
}
__date__ = "2013-03-05"
__license__ = "GPL (version 2 or later)"

from flask import jsonify, make_response

from user_metrics.config import logging
from user_metrics.api import pkl_data
from user_metrics.api.engine import MAX_CONCURRENT_JOBS, \
    QUEUE_WAIT, MW_UID_REGEX
from user_metrics.api.engine.data import get_users, set_data
from user_metrics.api.engine.request_meta import QUERY_PARAMS_BY_METRIC, \
    RequestMetaFactory
from user_metrics.metrics.users import MediaWikiUser
from user_metrics.metrics.metrics_manager import process_data_request

from multiprocessing import Process, Queue
from json import loads
from collections import namedtuple
from re import search
from os import getpid

job_item_type = namedtuple('JobItem', 'id process request queue')

def job_control(request_queue):
    """
        Controls the execution of user metrics requests

        Parameters
        ~~~~~~~~~~

        request_queue : multiprocessing.Queue
           Queues incoming API requests.

    """

    # Store executed and pending jobs respectively
    job_queue = list()
    wait_queue = list()

    # Global job ID number
    job_id = 0

    # Tallies the number of concurrently running jobs
    concurrent_jobs = 0

    logging.debug('{0} :: {1}  - STARTING...'
    .format(__name__, job_control.__name__))

    while 1:

        # Request Queue Processing
        # ------------------------

        try:
            # Pull an item off of the queue
            req_item = request_queue.get(timeout=QUEUE_WAIT)

            logging.debug('{0} :: {1}  Pull item from request queue -> {2}'
            .format(__name__, job_control.__name__, str(req_item)))

        except Exception:
            req_item = None
            logging.debug('{0} :: {1}  - Listening ...'
            .format(__name__, job_control.__name__))

        # Process complete jobs
        # ---------------------

        for job_item in job_queue:
            # Look for completed jobs
            if not job_item.process.is_alive():

                # Pull data off of the queue and add it to the queue data
                queue_data = loads(job_item.queue.get().data)
                response = make_response(jsonify(queue_data))

                set_data(job_item.request, response, pkl_data)
                # response_queue.put(response)
                del job_queue[job_queue.index(job_item)]

                concurrent_jobs -= 1

                logging.debug('{0} :: {1}  - RUN -> RESPONSE {2}'
                .format(__name__, job_control.__name__, str(job_item)))

        # Process pending jobs
        # --------------------

        for wait_req in wait_queue:
            if concurrent_jobs <= MAX_CONCURRENT_JOBS:
                # prepare job from item

                req_q = Queue()
                proc = Process(target=process_metrics, args=(req_q, wait_req))
                proc.start()

                job_item = job_item_type(job_id, proc, wait_req, req_q)
                job_queue.append(job_item)
                # job_queue.append(wait_req)

                concurrent_jobs += 1
                job_id += 1

                logging.debug('{0} :: {1}  - WAIT -> RUN {2}'
                .format(__name__, job_control.__name__, str(wait_req)))

        # Add newest job to the queue
        # ---------------------------

        if req_item and concurrent_jobs <= MAX_CONCURRENT_JOBS:
            rm = RequestMetaFactory(req_item['cohort'],
                                    req_item['cohort_refresh_ts'],
                                    req_item['metric'])
            logging.debug('{0} :: {1}  - REQUEST -> WAIT {2}'
            .format(__name__, job_control.__name__, str(rm)))
            wait_queue.append(rm)

    logging.debug('{0} :: {1}  - FINISHING.'
    .format(__name__, job_control.__name__))

def process_metrics(p, rm):
    """ Worker process for requests -
        this will typically operate in a forked process """

    logging.info(__name__ + ' :: START JOB %s (PID = %s)' % (str(rm),
                                                             getpid()))

    # obtain user list - handle the case where a lone user ID is passed
    if search(MW_UID_REGEX, str(rm.cohort_expr)):
        users = [rm.cohort_expr]
    # Special case where user lists are to be generated based on registered
    # user reg dates from the logging table -- see src/metrics/users.py
    elif rm.cohort_expr == 'all':
        users = MediaWikiUser(query_type=1)
    else:
        users = get_users(rm.cohort_expr)

    # unpack RequestMeta into dict using MEDIATOR
    args = {attr.metric_var: getattr(rm, attr.query_var)
            for attr in QUERY_PARAMS_BY_METRIC[rm.metric]}
    logging.info(__name__ + ' :: Calling %s with args = %s.' % (rm.metric,
                                                                str(args)))

    # process request
    results = process_data_request(rm.metric, users, **args)

    p.put(jsonify(results))
    logging.info(__name__ + ' :: END JOB %s (PID = %s)' % (str(rm), getpid()))

<commit_msg>rm. flask dependencies for data formatting in request_manager<commit_after>"""
    This module implements the request manager functionality.

    Job Queue and Processing
    ^^^^^^^^^^^^^^^^^^^^^^^^

    As requests are issued via http to the API a process queue will store all
    active jobs. Processes will be created and assume one of the following
    states throughout their existence: ::

        * 'pending' - The request has yet to be begin being processed
        * 'running' - The request is being processed
        * 'success' - The request has finished processing and is exposed at
            the url
        * 'failure' - The result has finished processing but dailed to expose
            results

    When a process a request is received and a job is created to service that
    request it enters the 'pending' state. If the job returns without
    exception it enters the 'success' state, otherwise it enters the 'failure'
    state.  The job remains in either of these states until it is cleared
    from the process queue.

    Response Data
    ^^^^^^^^^^^^^

    As requests are made to the API the data generated and formatted as JSON.
    The definition of is as follows: ::

        {   header : header_list,
            cohort_expr : cohort_gen_timestamp : metric : timeseries :
            aggregator : date_start : date_end : [ metric_param : ]* : data
        }

    Where each component is defined: ::

        header_str := list(str), list of header values
        cohort_expr := str, cohort ID expression
        cohort_gen_timestamp := str, cohort generation timestamp (earliest of
            all cohorts in expression)
        metric := str, user metric handle
        timeseries := boolean, indicates if this is a timeseries
        aggregator := str, aggregator used
        date_start := str, start datetime of request
        date_end := str, end datetime of request
        metric_param := -, optional metric parameters
        data := list(tuple), set of data points

    Request data is mapped to a query via metric objects and hashed in the
    dictionary `pkl_data`.

"""

__author__ = {
    "ryan faulkner": "rfaulkner@wikimedia.org"
}
__date__ = "2013-03-05"
__license__ = "GPL (version 2 or later)"

from user_metrics.config import logging
from user_metrics.api import pkl_data
from user_metrics.api.engine import MAX_CONCURRENT_JOBS, \
    QUEUE_WAIT, MW_UID_REGEX
from user_metrics.api.engine.data import get_users, set_data
from user_metrics.api.engine.request_meta import QUERY_PARAMS_BY_METRIC, \
    RequestMetaFactory
from user_metrics.metrics.users import MediaWikiUser
from user_metrics.metrics.metrics_manager import process_data_request

from multiprocessing import Process, Queue
from collections import namedtuple
from re import search
from os import getpid

job_item_type = namedtuple('JobItem', 'id process request queue')

def job_control(request_queue):
    """
        Controls the execution of user metrics requests

        Parameters
        ~~~~~~~~~~

        request_queue : multiprocessing.Queue
           Queues incoming API requests.

    """

    # Store executed and pending jobs respectively
    job_queue = list()
    wait_queue = list()

    # Global job ID number
    job_id = 0

    # Tallies the number of concurrently running jobs
    concurrent_jobs = 0

    logging.debug('{0} :: {1}  - STARTING...'
    .format(__name__, job_control.__name__))

    while 1:

        # Request Queue Processing
        # ------------------------

        try:
            # Pull an item off of the queue
            req_item = request_queue.get(timeout=QUEUE_WAIT)

            logging.debug('{0} :: {1}  Pull item from request queue -> {2}'
            .format(__name__, job_control.__name__, str(req_item)))

        except Exception:
            req_item = None
            logging.debug('{0} :: {1}  - Listening ...'
            .format(__name__, job_control.__name__))

        # Process complete jobs
        # ---------------------

        for job_item in job_queue:
            # Look for completed jobs
            if not job_item.process.is_alive():

                # Pull data off of the queue and add it to the queue data
                data = job_item.queue.get()
                set_data(job_item.request, data, pkl_data)

                del job_queue[job_queue.index(job_item)]

                concurrent_jobs -= 1

                logging.debug('{0} :: {1}  - RUN -> RESPONSE {2}, ' \
                              'Concurrent jobs = {3}'\
                .format(__name__, job_control.__name__, str(job_item),
                        concurrent_jobs))

        # Process pending jobs
        # --------------------

        for wait_req in wait_queue:
            if concurrent_jobs <= MAX_CONCURRENT_JOBS:
                # prepare job from item

                req_q = Queue()
                proc = Process(target=process_metrics, args=(req_q, wait_req))
                proc.start()

                job_item = job_item_type(job_id, proc, wait_req, req_q)
                job_queue.append(job_item)

                concurrent_jobs += 1
                job_id += 1

                logging.debug('{0} :: {1}  - WAIT -> RUN {2}, ' \
                              'Concurrent jobs = {3}'\
                .format(__name__, job_control.__name__, str(wait_req),
                        concurrent_jobs))

        # Add newest job to the queue
        # ---------------------------

        if req_item and concurrent_jobs <= MAX_CONCURRENT_JOBS:
            rm = RequestMetaFactory(req_item['cohort'],
                                    req_item['cohort_refresh_ts'],
                                    req_item['metric'])
            logging.debug('{0} :: {1}  - REQUEST -> WAIT {2}'
            .format(__name__, job_control.__name__, str(rm)))
            wait_queue.append(rm)

    logging.debug('{0} :: {1}  - FINISHING.'
    .format(__name__, job_control.__name__))

def process_metrics(p, rm):
    """ Worker process for requests -
        this will typically operate in a forked process """

    logging.info(__name__ + ' :: START JOB %s (PID = %s)' % (str(rm),
                                                             getpid()))

    # obtain user list - handle the case where a lone user ID is passed
    if search(MW_UID_REGEX, str(rm.cohort_expr)):
        users = [rm.cohort_expr]
    # Special case where user lists are to be generated based on registered
    # user reg dates from the logging table -- see src/metrics/users.py
    elif rm.cohort_expr == 'all':
        users = MediaWikiUser(query_type=1)
    else:
        users = get_users(rm.cohort_expr)

    # unpack RequestMeta into dict using MEDIATOR
    args = {attr.metric_var: getattr(rm, attr.query_var)
            for attr in QUERY_PARAMS_BY_METRIC[rm.metric]}
    logging.info(__name__ + ' :: Calling %s with args = %s.' % (rm.metric,
                                                                str(args)))

    # process request
    results = process_data_request(rm.metric, users, **args)
    p.put(results)

    logging.info(__name__ + ' :: END JOB %s (PID = %s)' % (str(rm), getpid()))

<|endoftext|>


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

"""
Helper functions for working with network manager.
"""

import collections
from gi.repository import GLib as glib
from gi.repository import NM as nm
import logging
import socket
import struct
import uuid

logger = logging.getLogger(__name__)

CONNECTION_TYPE_NAMES = collections.OrderedDict([
    ('802-3-ethernet', 'Ethernet'),
    ('802-11-wireless', 'Wi-Fi'),
    ('pppoe', 'PPPoE')
])

class ConnectionNotFound(Exception):
    """Network connection with a given name could not be found."""
    pass

class DeviceNotFound(Exception):
    """Network device for specified operation could not be found."""
    pass

def ipv4_string_to_int(address):
    """Return an integer equivalent of a string contain IPv4 address."""
    return struct.unpack("=I", socket.inet_aton(address))[0]

def _callback(source_object, result, user_data):
    """Called when an operation is completed."""
    del source_object  # Unused
    del result  # Unused
    del user_data  # Unused

def _commit_callback(connection, error, data=None):
    """Called when the connection changes are committed."""
    del connection
    del error
    del data

def get_interface_list(device_type):
    """Get a list of network interface available on the system."""
    interfaces = {}
    for device in nm.Client.new(None).get_devices():
        if device.get_device_type() == device_type:
            interfaces[device.get_iface()] = device.get_hw_address()

    return interfaces

def get_connection_list():
    """Get a list of active and available connections."""
    active_uuids = []
    client = nm.Client.new(None)
    for connection in client.get_active_connections():
        active_uuids.append(connection.get_uuid())

    connections = []
    for connection in client.get_connections():
        # Display a friendly type name if known.
        connection_type = connection.get_connection_type()
        connection_type = CONNECTION_TYPE_NAMES.get(connection_type,
                                                    connection_type)
        connections.append({
            'name': connection.get_id(),
            'uuid': connection.get_uuid(),
            'type': connection_type,
            'is_active': connection.get_uuid() in active_uuids,
        })
    connections.sort(key=lambda connection: connection['is_active'],
                     reverse=True)
    return connections

def get_connection(connection_uuid):
    """Return connection with matching uuid.

    Raise ConnectionNotFound if a connection with that uuid is not
    found.
    """
    client = nm.Client.new(None)
    connection = client.get_connection_by_uuid(connection_uuid)
    if not connection:
        raise ConnectionNotFound(connection_uuid)

    return connection

def get_active_connection(connection_uuid):
    """Return active connection with matching UUID.

    Raise ConnectionNotFound if a connection with that uuid is not
    found.
    """
    connections = nm.Client.new(None).get_active_connections()
    connections = {connection.get_uuid(): connection
                   for connection in connections}
    try:
        return connections[connection_uuid]
    except KeyError:
        raise ConnectionNotFound(connection_uuid)

def _update_common_settings(connection, connection_uuid, name, type_,
                            interface, zone, ipv4_method, ipv4_address):
    """Create/edit basic settings for network manager connections."""
    if not connection:
        connection = nm.SimpleConnection.new()

    # Connection
    settings = connection.get_setting_connection()
    if not settings:
        settings = nm.SettingConnection.new()
        connection.add_setting(settings)

    settings.set_property(nm.SETTING_CONNECTION_UUID, connection_uuid)
    settings.set_property(nm.SETTING_CONNECTION_ID, name)
    settings.set_property(nm.SETTING_CONNECTION_TYPE, type_)
    settings.set_property(nm.SETTING_CONNECTION_INTERFACE_NAME, interface)
    settings.set_property(nm.SETTING_CONNECTION_ZONE, zone)

    # IPv4
    settings = connection.get_setting_ip4_config()
    if not settings:
        settings = nm.SettingIP4Config.new()
        connection.add_setting(settings)

    settings.set_property(nm.SETTING_IP_CONFIG_METHOD, ipv4_method)
    if ipv4_method == nm.SETTING_IP4_CONFIG_METHOD_MANUAL and ipv4_address:
        ipv4_address_int = ipv4_string_to_int(ipv4_address)
        ipv4_prefix = nm.utils_ip4_get_default_prefix(ipv4_address_int)

        address = nm.IPAddress.new(socket.AF_INET, ipv4_address, ipv4_prefix)
        settings.add_address(address)

        settings.set_property(nm.SETTING_IP_CONFIG_GATEWAY, '0.0.0.0')
    else:
        settings.clear_addresses()

    return connection

def _update_ethernet_settings(connection, connection_uuid, name, interface,
                              zone, ipv4_method, ipv4_address):
    """Create/edit ethernet settings for network manager connections."""
    type_ = '802-3-ethernet'

    connection = _update_common_settings(connection, connection_uuid, name,
                                         type_, interface, zone, ipv4_method,
                                         ipv4_address)

    # Ethernet
    settings = connection.get_setting_wired()
    if not settings:
        settings = nm.SettingWired.new()
        connection.add_setting(settings)

    return connection

def _update_pppoe_settings(connection, connection_uuid, name, interface, zone,
                           username, password):
    """Create/edit ethernet settings for network manager connections."""
    type_ = 'pppoe'

    connection = _update_common_settings(connection, connection_uuid, name,
                                         type_, interface, zone, 'auto',
                                         '0.0.0.0')

    #pppoe
    settings = connection.get_setting_pppoe()
    if not settings:
        settings = nm.SettingPppoe.new()
        connection.add_setting(settings)

    settings.set_property(nm.SETTING_PPPOE_USERNAME, username)
    settings.set_property(nm.SETTING_PPPOE_PASSWORD, password)

    settings = connection.get_setting_ppp()
    if not settings:
        settings = nm.SettingPpp.new()
        connection.add_setting(settings)

    #ToDo: make this configurable?
    #ToDo: apt-get install ppp pppoe
    settings.set_property(nm.SETTING_PPP_LCP_ECHO_FAILURE, 5)
    settings.set_property(nm.SETTING_PPP_LCP_ECHO_INTERVAL, 30)

    return connection

def add_pppoe_connection(name, interface, zone, username, password):
    """Add an automatic pppoe connection in network manager.

    Return the UUID for the connection.
    """
    connection_uuid = str(uuid.uuid4())
    connection = _update_pppoe_settings(
        None, connection_uuid, name, interface, zone, username, password)
    client = nm.Client.new(None)
    client.add_connection_async(connection, True, None, _callback, None)
    return connection_uuid

def edit_pppoe_connection(connection, name, interface, zone, username,
                          password):
    """Edit an existing pppoe connection in network manager."""
    _update_pppoe_settings(
        connection, connection.get_uuid(), name, interface, zone, username,
        password)
    connection.commit_changes(True)

def add_ethernet_connection(name, interface, zone, ipv4_method, ipv4_address):
    """Add an automatic ethernet connection in network manager.

    Return the UUID for the connection.
    """
    connection_uuid = str(uuid.uuid4())
    connection = _update_ethernet_settings(
        None, connection_uuid, name, interface, zone, ipv4_method,
        ipv4_address)
    client = nm.Client.new(None)
    client.add_connection_async(connection, True, None, _callback, None)
    return connection_uuid

def edit_ethernet_connection(connection, name, interface, zone, ipv4_method,
                             ipv4_address):
    """Edit an existing ethernet connection in network manager."""
    _update_ethernet_settings(
        connection, connection.get_uuid(), name, interface, zone, ipv4_method,
        ipv4_address)
    connection.commit_changes(True)

def _update_wifi_settings(connection, connection_uuid, name, interface, zone,
                          ssid, mode, auth_mode, passphrase, ipv4_method,
                          ipv4_address):
    """Create/edit wifi settings for network manager connections."""
    type_ = '802-11-wireless'
    key_mgmt = 'wpa-psk'

    connection = _update_common_settings(connection, connection_uuid, name,
                                         type_, interface, zone, ipv4_method,
                                         ipv4_address)

    # Wireless
    settings = connection.get_setting_wireless()
    if not settings:
        settings = nm.SettingWireless.new()
        connection.add_setting(settings)

    ssid_gbytes = glib.Bytes.new(ssid.encode())
    settings.set_property(nm.SETTING_WIRELESS_SSID, ssid_gbytes)
    settings.set_property(nm.SETTING_WIRELESS_MODE, mode)

    # Wireless Security
    if auth_mode == 'wpa' and passphrase:
        settings = connection.get_setting_wireless_security()
        if not settings:
            settings = nm.SettingWirelessSecurity.new()
            connection.add_setting(settings)

        settings.set_property(nm.SETTING_WIRELESS_SECURITY_KEY_MGMT, key_mgmt)
        settings.set_property(nm.SETTING_WIRELESS_SECURITY_PSK, passphrase)
    else:
        connection.remove_setting(nm.SettingWirelessSecurity)

    return connection

def add_wifi_connection(name, interface, zone, ssid, mode, auth_mode,
                        passphrase, ipv4_method, ipv4_address):
    """Add an automatic Wi-Fi connection in network manager.

    Return the UUID for the connection.
    """
    connection_uuid = str(uuid.uuid4())
    connection = _update_wifi_settings(
        None, connection_uuid, name, interface, zone, ssid, mode, auth_mode,
        passphrase, ipv4_method, ipv4_address)
    client = nm.Client.new(None)
    client.add_connection_async(connection, True, None, _callback, None)
    return connection_uuid

def edit_wifi_connection(connection, name, interface, zone, ssid, mode,
                         auth_mode, passphrase, ipv4_method, ipv4_address):
    """Edit an existing wifi connection in network manager."""
    _update_wifi_settings(
        connection, connection.get_uuid(), name, interface, zone, ssid, mode,
        auth_mode, passphrase, ipv4_method, ipv4_address)
    connection.commit_changes(True)

def activate_connection(connection_uuid):
    """Find and activate a network connection."""
    # Find the connection
    connection = get_connection(connection_uuid)
    interface = connection.get_interface_name()
    client = nm.Client.new(None)
    for device in client.get_devices():
        if device.get_iface() == interface:
            client.activate_connection_async(connection,
                                             device,
                                             '/', None, _callback,
                                             None)
    return connection

def deactivate_connection(connection_uuid):
    """Find and de-activate a network connection."""
    active_connection = get_active_connection(connection_uuid)
    nm.Client.new(None).deactivate_connection(active_connection)
    return active_connection

def delete_connection(connection_uuid):
    """Delete an exiting connection from network manager.

    Raise ConnectionNotFound if connection does not exist.
    """
    connection = get_connection(connection_uuid)
    name = connection.get_id()
    connection.delete()
    return name

def wifi_scan():
    """Scan for available access points across all Wi-Fi devices."""
    access_points = []
    for device in nm.Client.new(None).get_devices():
        if device.get_device_type() != nm.DeviceType.WIFI:
            continue

        for access_point in device.get_access_points():
            # Retrieve the bytes in SSID.  Don't convert to utf-8 or
            # escape it in any way as it may contain null bytes.  When
            # this is used in the URL it will be escaped properly and
            # unescaped when taken as view function's argument.
            ssid = access_point.get_ssid()
            ssid_string = ssid.get_data() if ssid else ''
            access_points.append({
                'ssid': ssid_string,
                'strength': access_point.get_strength()})

    return access_points
<commit_msg>networks: Refactor common settings updation<commit_after>#
# This file is part of Plinth.
#
# This program is free software: you can redistribute it and/or modify
# it under the terms of the GNU Affero General Public License as
# published by the Free Software Foundation, either version 3 of the
# License, or (at your option) any later version.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU Affero General Public License for more details.
#
# You should have received a copy of the GNU Affero General Public License
# along with this program.  If not, see <http://www.gnu.org/licenses/>.
#

"""
Helper functions for working with network manager.
"""

import collections
from gi.repository import GLib as glib
from gi.repository import NM as nm
import logging
import socket
import struct
import uuid

logger = logging.getLogger(__name__)

CONNECTION_TYPE_NAMES = collections.OrderedDict([
    ('802-3-ethernet', 'Ethernet'),
    ('802-11-wireless', 'Wi-Fi'),
    ('pppoe', 'PPPoE')
])

class ConnectionNotFound(Exception):
    """Network connection with a given name could not be found."""
    pass

class DeviceNotFound(Exception):
    """Network device for specified operation could not be found."""
    pass

def ipv4_string_to_int(address):
    """Return an integer equivalent of a string contain IPv4 address."""
    return struct.unpack("=I", socket.inet_aton(address))[0]

def _callback(source_object, result, user_data):
    """Called when an operation is completed."""
    del source_object  # Unused
    del result  # Unused
    del user_data  # Unused

def _commit_callback(connection, error, data=None):
    """Called when the connection changes are committed."""
    del connection
    del error
    del data

def get_interface_list(device_type):
    """Get a list of network interface available on the system."""
    interfaces = {}
    for device in nm.Client.new(None).get_devices():
        if device.get_device_type() == device_type:
            interfaces[device.get_iface()] = device.get_hw_address()

    return interfaces

def get_connection_list():
    """Get a list of active and available connections."""
    active_uuids = []
    client = nm.Client.new(None)
    for connection in client.get_active_connections():
        active_uuids.append(connection.get_uuid())

    connections = []
    for connection in client.get_connections():
        # Display a friendly type name if known.
        connection_type = connection.get_connection_type()
        connection_type = CONNECTION_TYPE_NAMES.get(connection_type,
                                                    connection_type)
        connections.append({
            'name': connection.get_id(),
            'uuid': connection.get_uuid(),
            'type': connection_type,
            'is_active': connection.get_uuid() in active_uuids,
        })
    connections.sort(key=lambda connection: connection['is_active'],
                     reverse=True)
    return connections

def get_connection(connection_uuid):
    """Return connection with matching uuid.

    Raise ConnectionNotFound if a connection with that uuid is not
    found.
    """
    client = nm.Client.new(None)
    connection = client.get_connection_by_uuid(connection_uuid)
    if not connection:
        raise ConnectionNotFound(connection_uuid)

    return connection

def get_active_connection(connection_uuid):
    """Return active connection with matching UUID.

    Raise ConnectionNotFound if a connection with that uuid is not
    found.
    """
    connections = nm.Client.new(None).get_active_connections()
    connections = {connection.get_uuid(): connection
                   for connection in connections}
    try:
        return connections[connection_uuid]
    except KeyError:
        raise ConnectionNotFound(connection_uuid)

def _update_common_settings(connection, connection_uuid, name, type_, interface,
                            zone):
    """Create/edit basic settings for network manager connections.

    Return newly created connection object if connection is None.
    """
    if not connection:
        connection = nm.SimpleConnection.new()

    settings = connection.get_setting_connection()
    if not settings:
        settings = nm.SettingConnection.new()
        connection.add_setting(settings)

    settings.set_property(nm.SETTING_CONNECTION_UUID, connection_uuid)
    settings.set_property(nm.SETTING_CONNECTION_ID, name)
    settings.set_property(nm.SETTING_CONNECTION_TYPE, type_)
    settings.set_property(nm.SETTING_CONNECTION_INTERFACE_NAME, interface)
    settings.set_property(nm.SETTING_CONNECTION_ZONE, zone)

    return connection

def _update_ipv4_settings(connection, ipv4_method, ipv4_address):
    """Edit IPv4 settings for network manager connections."""
    settings = connection.get_setting_ip4_config()
    if not settings:
        settings = nm.SettingIP4Config.new()
        connection.add_setting(settings)

    settings.set_property(nm.SETTING_IP_CONFIG_METHOD, ipv4_method)
    if ipv4_method == nm.SETTING_IP4_CONFIG_METHOD_MANUAL and ipv4_address:
        ipv4_address_int = ipv4_string_to_int(ipv4_address)
        ipv4_prefix = nm.utils_ip4_get_default_prefix(ipv4_address_int)

        address = nm.IPAddress.new(socket.AF_INET, ipv4_address, ipv4_prefix)
        settings.add_address(address)

        settings.set_property(nm.SETTING_IP_CONFIG_GATEWAY, '0.0.0.0')
    else:
        settings.clear_addresses()

def _update_ethernet_settings(connection, connection_uuid, name, interface,
                              zone, ipv4_method, ipv4_address):
    """Create/edit ethernet settings for network manager connections."""
    type_ = '802-3-ethernet'

    connection = _update_common_settings(connection, connection_uuid, name,
                                         type_, interface, zone)
    _update_ipv4_settings(connection, ipv4_method, ipv4_address)

    # Ethernet
    settings = connection.get_setting_wired()
    if not settings:
        settings = nm.SettingWired.new()
        connection.add_setting(settings)

    return connection

def _update_pppoe_settings(connection, connection_uuid, name, interface, zone,
                           username, password):
    """Create/edit ethernet settings for network manager connections."""
    type_ = 'pppoe'

    connection = _update_common_settings(connection, connection_uuid, name,
                                         type_, interface, zone)

    #pppoe
    settings = connection.get_setting_pppoe()
    if not settings:
        settings = nm.SettingPppoe.new()
        connection.add_setting(settings)

    settings.set_property(nm.SETTING_PPPOE_USERNAME, username)
    settings.set_property(nm.SETTING_PPPOE_PASSWORD, password)

    settings = connection.get_setting_ppp()
    if not settings:
        settings = nm.SettingPpp.new()
        connection.add_setting(settings)

    #ToDo: make this configurable?
    #ToDo: apt-get install ppp pppoe
    settings.set_property(nm.SETTING_PPP_LCP_ECHO_FAILURE, 5)
    settings.set_property(nm.SETTING_PPP_LCP_ECHO_INTERVAL, 30)

    return connection

def add_pppoe_connection(name, interface, zone, username, password):
    """Add an automatic pppoe connection in network manager.

    Return the UUID for the connection.
    """
    connection_uuid = str(uuid.uuid4())
    connection = _update_pppoe_settings(
        None, connection_uuid, name, interface, zone, username, password)
    client = nm.Client.new(None)
    client.add_connection_async(connection, True, None, _callback, None)
    return connection_uuid

def edit_pppoe_connection(connection, name, interface, zone, username,
                          password):
    """Edit an existing pppoe connection in network manager."""
    _update_pppoe_settings(
        connection, connection.get_uuid(), name, interface, zone, username,
        password)
    connection.commit_changes(True)

def add_ethernet_connection(name, interface, zone, ipv4_method, ipv4_address):
    """Add an automatic ethernet connection in network manager.

    Return the UUID for the connection.
    """
    connection_uuid = str(uuid.uuid4())
    connection = _update_ethernet_settings(
        None, connection_uuid, name, interface, zone, ipv4_method,
        ipv4_address)
    client = nm.Client.new(None)
    client.add_connection_async(connection, True, None, _callback, None)
    return connection_uuid

def edit_ethernet_connection(connection, name, interface, zone, ipv4_method,
                             ipv4_address):
    """Edit an existing ethernet connection in network manager."""
    _update_ethernet_settings(
        connection, connection.get_uuid(), name, interface, zone, ipv4_method,
        ipv4_address)
    connection.commit_changes(True)

def _update_wifi_settings(connection, connection_uuid, name, interface, zone,
                          ssid, mode, auth_mode, passphrase, ipv4_method,
                          ipv4_address):
    """Create/edit wifi settings for network manager connections."""
    type_ = '802-11-wireless'
    key_mgmt = 'wpa-psk'

    connection = _update_common_settings(connection, connection_uuid, name,
                                         type_, interface, zone)
    _update_ipv4_settings(connection, ipv4_method, ipv4_address)

    # Wireless
    settings = connection.get_setting_wireless()
    if not settings:
        settings = nm.SettingWireless.new()
        connection.add_setting(settings)

    ssid_gbytes = glib.Bytes.new(ssid.encode())
    settings.set_property(nm.SETTING_WIRELESS_SSID, ssid_gbytes)
    settings.set_property(nm.SETTING_WIRELESS_MODE, mode)

    # Wireless Security
    if auth_mode == 'wpa' and passphrase:
        settings = connection.get_setting_wireless_security()
        if not settings:
            settings = nm.SettingWirelessSecurity.new()
            connection.add_setting(settings)

        settings.set_property(nm.SETTING_WIRELESS_SECURITY_KEY_MGMT, key_mgmt)
        settings.set_property(nm.SETTING_WIRELESS_SECURITY_PSK, passphrase)
    else:
        connection.remove_setting(nm.SettingWirelessSecurity)

    return connection

def add_wifi_connection(name, interface, zone, ssid, mode, auth_mode,
                        passphrase, ipv4_method, ipv4_address):
    """Add an automatic Wi-Fi connection in network manager.

    Return the UUID for the connection.
    """
    connection_uuid = str(uuid.uuid4())
    connection = _update_wifi_settings(
        None, connection_uuid, name, interface, zone, ssid, mode, auth_mode,
        passphrase, ipv4_method, ipv4_address)
    client = nm.Client.new(None)
    client.add_connection_async(connection, True, None, _callback, None)
    return connection_uuid

def edit_wifi_connection(connection, name, interface, zone, ssid, mode,
                         auth_mode, passphrase, ipv4_method, ipv4_address):
    """Edit an existing wifi connection in network manager."""
    _update_wifi_settings(
        connection, connection.get_uuid(), name, interface, zone, ssid, mode,
        auth_mode, passphrase, ipv4_method, ipv4_address)
    connection.commit_changes(True)

def activate_connection(connection_uuid):
    """Find and activate a network connection."""
    # Find the connection
    connection = get_connection(connection_uuid)
    interface = connection.get_interface_name()
    client = nm.Client.new(None)
    for device in client.get_devices():
        if device.get_iface() == interface:
            client.activate_connection_async(connection,
                                             device,
                                             '/', None, _callback,
                                             None)
    return connection

def deactivate_connection(connection_uuid):
    """Find and de-activate a network connection."""
    active_connection = get_active_connection(connection_uuid)
    nm.Client.new(None).deactivate_connection(active_connection)
    return active_connection

def delete_connection(connection_uuid):
    """Delete an exiting connection from network manager.

    Raise ConnectionNotFound if connection does not exist.
    """
    connection = get_connection(connection_uuid)
    name = connection.get_id()
    connection.delete()
    return name

def wifi_scan():
    """Scan for available access points across all Wi-Fi devices."""
    access_points = []
    for device in nm.Client.new(None).get_devices():
        if device.get_device_type() != nm.DeviceType.WIFI:
            continue

        for access_point in device.get_access_points():
            # Retrieve the bytes in SSID.  Don't convert to utf-8 or
            # escape it in any way as it may contain null bytes.  When
            # this is used in the URL it will be escaped properly and
            # unescaped when taken as view function's argument.
            ssid = access_point.get_ssid()
            ssid_string = ssid.get_data() if ssid else ''
            access_points.append({
                'ssid': ssid_string,
                'strength': access_point.get_strength()})

    return access_points
<|endoftext|>


<commit_before>from __future__ import with_statement
import collections
import os
import signal
import sys
import thread
import time
import argparse
from plop import platform

class Collector(object):
    MODES = {
        'prof': (platform.ITIMER_PROF, signal.SIGPROF),
        'virtual': (platform.ITIMER_VIRTUAL, signal.SIGVTALRM),
        'real': (platform.ITIMER_REAL, signal.SIGALRM),
    }

    def __init__(self, interval=0.01, mode='virtual'):
        self.interval = interval
        self.mode = mode
        assert mode in Collector.MODES
        timer, sig = Collector.MODES[self.mode]
        signal.signal(sig, self.handler)
        signal.siginterrupt(sig, False)
        self.reset()

    def reset(self):
        self.stacks = list()
        self.samples_remaining = 0
        self.stopping = False
        self.stopped = False

        self.samples_taken = 0
        self.sample_time = 0

    def start(self, duration=30.0):
        self.stopping = False
        self.stopped = False
        self.samples_remaining = int(duration / self.interval)
        timer, sig = Collector.MODES[self.mode]
        platform.setitimer(timer, self.interval, self.interval)

    def stop(self):
        self.stopping = True
        self.wait()

    def wait(self):
        while not self.stopped:
            pass  # need busy wait; ITIMER_PROF doesn't proceed while sleeping

    def handler(self, sig, current_frame):
        start = time.time()
        self.samples_remaining -= 1
        if self.samples_remaining <= 0 or self.stopping:
            platform.setitimer(Collector.MODES[self.mode][0], 0, 0)
            self.stopped = True
            return
        current_tid = thread.get_ident()
        for tid, frame in sys._current_frames().items():
            if tid == current_tid:
                frame = current_frame
            frames = []
            while frame is not None:
                code = frame.f_code
                frames.append((code.co_filename, code.co_firstlineno, code.co_name))
                frame = frame.f_back
            self.stacks.append(frames)
        end = time.time()
        self.samples_taken += 1
        self.sample_time += (end - start)

class CollectorFormatter(object):
    """
    Abstract class for output formats
    """
    def format(self, collector):
        raise Exception("not implemented")

    def store(self, collector, filename):
        with open(filename, "wb") as f:
            f.write(self.format(collector))

class PlopFormatter(CollectorFormatter):
    """
    Formats stack frames for plop.viewer
    """
    def __init__(self, max_stacks=50):
        self.max_stacks = 50

    def format(self, collector):
        # defaultdict instead of counter for pre-2.7 compatibility
        stack_counts = collections.defaultdict(int)
        for frames in collector.stacks:
            stack_counts[tuple(frames)] += 1
        stack_counts = dict(sorted(stack_counts.iteritems(),
                                key=lambda kv: -kv[1])[:self.max_stacks])
        return repr(stack_counts)

class FlamegraphFormatter(CollectorFormatter):
    """
    Creates Flamegraph files
    """
    def format(self, collector):
        output = ""
        previous = None
        previous_count = 1
        for stack in collector.stacks:
            current = self.format_flame(stack)
            if current == previous:
                previous_count += 1
            else:
                output += "%s %d\n" % (previous, previous_count)
                previous_count = 1
                previous = current
        output += "%s %d\n" % (previous, previous_count)
        return output

    def format_flame(self, stack):
        funcs = map("{0[2]} ({0[0]}:{0[1]})".format, reversed(stack))
        return ";".join(funcs)

def main():
    # TODO: more options, refactor this into somewhere shared
    # between tornado.autoreload and auto2to3
    parser = argparse.ArgumentParser(description="Plop: Python Low-Overhead Profiler",
                                     prog="python -m plop.collector",
                                     formatter_class=argparse.ArgumentDefaultsHelpFormatter)
    parser.add_argument("--format", "-f", help="Output format",
                        choices=["plop", "flamegraph"], default="plop")
    parser.add_argument("--module", "-m", help="Execute target as a module",
                        action="store_const", const=True, default=False)
    parser.add_argument("--mode", help="Interval timer mode to use, see `man 2 setitimer`",
                        choices=["prof", "real", "virtual"], default="prof")
    parser.add_argument("--interval", help="Timer interval in seconds", default=0.01, type=float)
    parser.add_argument("--duration", help="Profiling duration in seconds", default=3600,
                        type=int)
    parser.add_argument("--max-stacks", help=("Number of most frequent stacks to store."
                        " Ignored for Flamegraph output."), type=int, default=50)
    parser.add_argument("target", help="Module or script to run")
    parser.add_argument("arguments", nargs=argparse.REMAINDER,
                        help="Pass-through arguments for the profiled application")
    args = parser.parse_args()
    sys.argv = [args.target] + args.arguments

    if args.format == "flamegraph":
        extension = "flame"
        formatter = FlamegraphFormatter()
    elif args.format == "plop":
        extension = "plop"
        formatter = PlopFormatter(max_stacks=args.max_stacks)
    else:
        sys.stderr.write("Unhandled output format: %s" % args.format)
        sys.stderr.flush()
        sys.exit(1)

    if not os.path.exists('profiles'):
        os.mkdir('profiles')
    filename = 'profiles/%s-%s.%s' % (args.target, time.strftime('%Y%m%d-%H%M-%S'),
                                      extension)

    collector = Collector(mode=args.mode, interval=args.interval)
    collector.start(duration=args.duration)
    exit_code = 0
    try:
        if args.module:
            import runpy
            runpy.run_module(args.target, run_name="__main__", alter_sys=True)
        else:
            with open(args.target) as f:
                global __file__
                __file__ = args.target
                # Use globals as our "locals" dictionary so that
                # something that tries to import __main__ (e.g. the unittest
                # module) will see the right things.
                exec f.read() in globals(), globals()
    except SystemExit, e:
        exit_code = e.code
    collector.stop()
    if collector.samples_taken:
        formatter.store(collector, filename)
        print "profile output saved to %s" % filename
        overhead = float(collector.sample_time) / collector.samples_taken
        print "overhead was %s per sample (%s%%)" % (
            overhead, overhead / collector.interval)
    else:
        print "no samples collected; program was too fast"
    sys.exit(exit_code)

if __name__ == '__main__':
    main()
<commit_msg>Remove __package__ from the script being run.<commit_after>from __future__ import with_statement
import collections
import os
import signal
import sys
import thread
import time
import argparse
import plop.platform

class Collector(object):
    MODES = {
        'prof': (plop.platform.ITIMER_PROF, signal.SIGPROF),
        'virtual': (plop.platform.ITIMER_VIRTUAL, signal.SIGVTALRM),
        'real': (plop.platform.ITIMER_REAL, signal.SIGALRM),
    }

    def __init__(self, interval=0.01, mode='virtual'):
        self.interval = interval
        self.mode = mode
        assert mode in Collector.MODES
        timer, sig = Collector.MODES[self.mode]
        signal.signal(sig, self.handler)
        signal.siginterrupt(sig, False)
        self.reset()

    def reset(self):
        self.stacks = list()
        self.samples_remaining = 0
        self.stopping = False
        self.stopped = False

        self.samples_taken = 0
        self.sample_time = 0

    def start(self, duration=30.0):
        self.stopping = False
        self.stopped = False
        self.samples_remaining = int(duration / self.interval)
        timer, sig = Collector.MODES[self.mode]
        plop.platform.setitimer(timer, self.interval, self.interval)

    def stop(self):
        self.stopping = True
        self.wait()

    def wait(self):
        while not self.stopped:
            pass  # need busy wait; ITIMER_PROF doesn't proceed while sleeping

    def handler(self, sig, current_frame):
        start = time.time()
        self.samples_remaining -= 1
        if self.samples_remaining <= 0 or self.stopping:
            plop.platform.setitimer(Collector.MODES[self.mode][0], 0, 0)
            self.stopped = True
            return
        current_tid = thread.get_ident()
        for tid, frame in sys._current_frames().items():
            if tid == current_tid:
                frame = current_frame
            frames = []
            while frame is not None:
                code = frame.f_code
                frames.append((code.co_filename, code.co_firstlineno, code.co_name))
                frame = frame.f_back
            self.stacks.append(frames)
        end = time.time()
        self.samples_taken += 1
        self.sample_time += (end - start)

class CollectorFormatter(object):
    """
    Abstract class for output formats
    """
    def format(self, collector):
        raise Exception("not implemented")

    def store(self, collector, filename):
        with open(filename, "wb") as f:
            f.write(self.format(collector))

class PlopFormatter(CollectorFormatter):
    """
    Formats stack frames for plop.viewer
    """
    def __init__(self, max_stacks=50):
        self.max_stacks = 50

    def format(self, collector):
        # defaultdict instead of counter for pre-2.7 compatibility
        stack_counts = collections.defaultdict(int)
        for frames in collector.stacks:
            stack_counts[tuple(frames)] += 1
        stack_counts = dict(sorted(stack_counts.iteritems(),
                                   key=lambda kv: -kv[1])[:self.max_stacks])
        return repr(stack_counts)

class FlamegraphFormatter(CollectorFormatter):
    """
    Creates Flamegraph files
    """
    def format(self, collector):
        output = ""
        previous = None
        previous_count = 1
        for stack in collector.stacks:
            current = self.format_flame(stack)
            if current == previous:
                previous_count += 1
            else:
                output += "%s %d\n" % (previous, previous_count)
                previous_count = 1
                previous = current
        output += "%s %d\n" % (previous, previous_count)
        return output

    def format_flame(self, stack):
        funcs = map("{0[2]} ({0[0]}:{0[1]})".format, reversed(stack))
        return ";".join(funcs)

def main():
    # TODO: more options, refactor this into somewhere shared
    # between tornado.autoreload and auto2to3
    parser = argparse.ArgumentParser(description="Plop: Python Low-Overhead Profiler",
                                     prog="python -m plop.collector",
                                     formatter_class=argparse.ArgumentDefaultsHelpFormatter)
    parser.add_argument("--format", "-f", help="Output format",
                        choices=["plop", "flamegraph"], default="plop")
    parser.add_argument("--module", "-m", help="Execute target as a module",
                        action="store_const", const=True, default=False)
    parser.add_argument("--mode", help="Interval timer mode to use, see `man 2 setitimer`",
                        choices=["prof", "real", "virtual"], default="prof")
    parser.add_argument("--interval", help="Timer interval in seconds", default=0.01, type=float)
    parser.add_argument("--duration", help="Profiling duration in seconds", default=3600,
                        type=int)
    parser.add_argument("--max-stacks", help=("Number of most frequent stacks to store."
                        " Ignored for Flamegraph output."), type=int, default=50)
    parser.add_argument("target", help="Module or script to run")
    parser.add_argument("arguments", nargs=argparse.REMAINDER,
                        help="Pass-through arguments for the profiled application")
    args = parser.parse_args()
    sys.argv = [args.target] + args.arguments

    if args.format == "flamegraph":
        extension = "flame"
        formatter = FlamegraphFormatter()
    elif args.format == "plop":
        extension = "plop"
        formatter = PlopFormatter(max_stacks=args.max_stacks)
    else:
        sys.stderr.write("Unhandled output format: %s" % args.format)
        sys.stderr.flush()
        sys.exit(1)

    if not os.path.exists('profiles'):
        os.mkdir('profiles')
    filename = 'profiles/%s-%s.%s' % (args.target, time.strftime('%Y%m%d-%H%M-%S'),
                                      extension)

    collector = Collector(mode=args.mode, interval=args.interval)
    collector.start(duration=args.duration)
    exit_code = 0
    try:
        if args.module:
            import runpy
            runpy.run_module(args.target, run_name="__main__", alter_sys=True)
        else:
            with open(args.target) as f:
                # Execute the script in our namespace instead of creating
                # a new one so that something that tries to import __main__
                # (e.g. the unittest module) will see names defined in the
                # script instead of just those defined in this module.
                global __file__
                __file__ = args.target
                # If __package__ is defined, imports may be incorrectly
                # interpreted as relative to this module.
                global __package__
                del __package__
                exec f.read() in globals(), globals()
    except SystemExit, e:
        exit_code = e.code
    collector.stop()
    if collector.samples_taken:
        formatter.store(collector, filename)
        print "profile output saved to %s" % filename
        overhead = float(collector.sample_time) / collector.samples_taken
        print "overhead was %s per sample (%s%%)" % (
            overhead, overhead / collector.interval)
    else:
        print "no samples collected; program was too fast"
    sys.exit(exit_code)

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


<commit_before>###
# Copyright (c) 2002-2004, Jeremiah Fincher
# All rights reserved.
#
# Redistribution and use in source and binary forms, with or without
# modification, are permitted provided that the following conditions are met:
#
#   * Redistributions of source code must retain the above copyright notice,
#     this list of conditions, and the following disclaimer.
#   * Redistributions in binary form must reproduce the above copyright notice,
#     this list of conditions, and the following disclaimer in the
#     documentation and/or other materials provided with the distribution.
#   * Neither the name of the author of this software nor the name of
#     contributors to this software may be used to endorse or promote products
#     derived from this software without specific prior written consent.
#
# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
# ARE DISCLAIMED.  IN NO EVENT SHALL THE COPYRIGHT OWNER OR CONTRIBUTORS BE
# LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR
# CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF
# SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS
# INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN
# CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE)
# ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE
# POSSIBILITY OF SUCH DAMAGE.
###

"""
Silently listens to a channel, building an SQL database of Markov Chains for
later hijinks.  To read more about Markov Chains, check out
<http://www.cs.bell-labs.com/cm/cs/pearls/sec153.html>.  When the database is
large enough, you can have it make fun little random messages from it.
"""

__revision__ = "$Id$"

import supybot.plugins as plugins

import sets
import Queue
import anydbm
import random
import os.path
import threading

import supybot.world as world
import supybot.ircmsgs as ircmsgs
import supybot.ircutils as ircutils
import supybot.privmsgs as privmsgs
import supybot.callbacks as callbacks

class MarkovDBInterface(object):
    def close(self):
        pass

    def addPair(self, channel, first, second, follower,
                isFirst=False, isLast=False):
        pass

    def getFirstPair(self, channel):
        pass

    def getPair(self, channel, first, second):
        # Returns (follower, last) tuple.
        pass

    def firsts(self, channel):
        pass

    def lasts(self, channel):
        pass

    def pairs(self, channel):
        pass

    def follows(self, channel):
        pass

class SqliteMarkovDB(object):
    def __init__(self, filename):
        self.dbs = ircutils.IrcDict()
        self.filename = filename

    def close(self):
        for db in self.dbs.values():
            db.close()

    def _getDb(self, channel):
        try:
            import sqlite
        except ImportError:
            raise callbacks.Error, 'You need to have PySQLite installed to '\
                                   'use this plugin.  Download it at '\
                                   '<http://pysqlite.sf.net/>'
        if channel not in self.dbs:
            filename = plugins.makeChannelFilename(self.filename, channel)
            if os.path.exists(filename):
                self.dbs[channel] = sqlite.connect(filename)
                return self.dbs[channel]
            #else:
            self.dbs[channel] = sqlite.connect(filename)
            cursor = self.dbs[channel].cursor()
            # TODO Finish the rest of the implementation
        return self.dbs[channel]

    def addPair(self, channel, first, second, follower,
                isFirst=False, isLast=False):
        pass

    def getFirstPair(self, channel):
        pass

    def getFollower(self, channel, first, second):
        # Returns (follower, last) tuple.
        pass

    def firsts(self, channel):
        pass

    def lasts(self, channel):
        pass

    def pairs(self, channel):
        pass

    def follows(self, channel):
        pass

class DbmMarkovDB(object):
    def __init__(self, filename):
        self.dbs = ircutils.IrcDict()
        ## Stupid anydbm seems to append .db to the end of this.
        #self.filename = filename.replace('.db', '')
        self.filename = filename

    def close(self):
        for db in self.dbs.values():
            db.close()

    def _getDb(self, channel):
        if channel not in self.dbs:
            filename = plugins.makeChannelFilename(self.filename, channel)
            # To keep the code simpler for addPair, I decided not to make
            # self.dbs[channel]['firsts'] and ['lasts'].  Instead, we'll pad
            # the words list being sent to addPair such that ['\n \n'] will be
            # ['firsts'] and ['\n'] will be ['lasts'].  This also means isFirst
            # and isLast aren't necessary, but they'll be left alone in case
            # one of the other Db formats uses them or someone decides that I
            # was wrong and changes my code.
            self.dbs[channel] = anydbm.open(filename, 'c')
        return self.dbs[channel]

    def _flush(self, db):
        if hasattr(db, 'sync'):
            db.sync()
        if hasattr(db, 'flush'):
            db.flush()

    def addPair(self, channel, first, second, follower,
                isFirst=False, isLast=False):
        db = self._getDb(channel)
        combined = self._combine(first, second)
        if db.has_key(combined): # EW!
            db[combined] = ' '.join([db[combined], follower])
        else:
            db[combined] = follower
        if follower == '\n':
            if db.has_key('\n'):
                db['\n'] = ' '.join([db['\n'], second])
            else:
                db['\n'] = second
        self._flush(db)

    def getFirstPair(self, channel):
        db = self._getDb(channel)
        firsts = db['\n \n'].split()
        if firsts:
            if firsts:
                return ('\n', random.choice(firsts))
            else:
                raise KeyError, 'No firsts for %s.' % channel
        else:
            raise KeyError, 'No firsts for %s.' % channel

    def _combine(self, first, second):
        return '%s %s' % (first, second)

    def getFollower(self, channel, first, second):
        db = self._getDb(channel)
        followers = db[self._combine(first, second)]
        follower = random.choice(followers.split(' '))
        return (follower, follower == '\n')

    def firsts(self, channel):
        db = self._getDb(channel)
        if db.has_key('\n \n'):
            return len(sets.Set(db['\n \n'].split()))
        else:
            return 0

    def lasts(self, channel):
        db = self._getDb(channel)
        if db.has_key('\n'):
            return len(sets.Set(db['\n'].split()))
        else:
            return 0

    def pairs(self, channel):
        db = self._getDb(channel)
        pairs = [k for k in db.keys() if '\n' not in k]
        return len(pairs)

    def follows(self, channel):
        db = self._getDb(channel)
        follows = [len(v.split()) for (k,v) in db.iteritems() if '\n' not in k]
        return sum(follows)

MarkovDB = plugins.DB('Markov',
                      {'anydbm': DbmMarkovDB})

class MarkovWorkQueue(threading.Thread):
    def __init__(self, *args, **kwargs):
        name = 'Thread #%s (MarkovWorkQueue)' % world.threadsSpawned
        world.threadsSpawned += 1
        threading.Thread.__init__(self, name=name)
        self.db = MarkovDB(*args, **kwargs)
        self.q = Queue.Queue()
        self.killed = False
        self.setDaemon(True)
        self.start()

    def die(self):
        self.killed = True
        self.q.put(None)

    def enqueue(self, f):
        self.q.put(f)

    def run(self):
        while not self.killed:
            f = self.q.get()
            if f is not None:
                f(self.db)
        self.db.close()

class Markov(callbacks.Privmsg):
    def __init__(self):
        self.q = MarkovWorkQueue()
        callbacks.Privmsg.__init__(self)

    def die(self):
        self.q.die()

    def tokenize(self, m):
        if ircmsgs.isAction(m):
            return ircmsgs.unAction(m).split()
        elif ircmsgs.isCtcp(m):
            return []
        else:
            return m.args[1].split()

    def doPrivmsg(self, irc, msg):
        channel = msg.args[0]
        if ircutils.isChannel(channel):
            words = self.tokenize(msg)
            words.insert(0, '\n')
            words.insert(0, '\n')
            words.append('\n')
            # This shouldn't happen often (CTCP messages being the possible exception)
            if not words or len(words) == 3:
                return
            def doPrivmsg(db):
                for (first, second, follower) in window(words, 3):
                    db.addPair(channel, first, second, follower)
            self.q.enqueue(doPrivmsg)

    def markov(self, irc, msg, args):
        """[<channel>] [word1 word2]

        Returns a randomly-generated Markov Chain generated sentence from the
        data kept on <channel> (which is only necessary if not sent in the
        channel itself).  If word1 and word2 are specified, they will be used
        to start the Markov chain.
        """
        channel = privmsgs.getChannel(msg, args)
        (word1, word2) = privmsgs.getArgs(args, required=0, optional=2)
        def markov(db):
            if word1 and word2:
                givenArgs = True
                words = [word1, word2]
            else:
                givenArgs = False
                try:
                    # words is of the form ['\r', word]
                    words = list(db.getFirstPair(channel))
                except KeyError:
                    irc.error('I don\'t have any first pairs for %s.' %channel)
                    return
            follower = words[-1]
            last = False
            resp = []
            while not last:
                resp.append(follower)
                try:
                    (follower,last) = db.getFollower(channel, words[-2],
                                                     words[-1])
                except KeyError:
                    irc.error('I found a broken link in the Markov chain.  '
                              'Maybe I received two bad links to start the '
                              'chain.')
                    return
                words.append(follower)
            if givenArgs:
                irc.reply(' '.join(words[:-1]))
            else:
                irc.reply(' '.join(resp))
        self.q.enqueue(markov)

    def firsts(self, irc, msg, args):
        """[<channel>]

        Returns the number of Markov's first links in the database for
        <channel>.
        """
        channel = privmsgs.getChannel(msg, args)
        def firsts(db):
            s = 'There are %s firsts in my Markov database for %s.'
            irc.reply(s % (db.firsts(channel), channel))
        self.q.enqueue(firsts)

    def lasts(self, irc, msg, args):
        """[<channel>]

        Returns the number of Markov's last links in the database for
        <channel>.
        """
        channel = privmsgs.getChannel(msg, args)
        def lasts(db):
            s = 'There are %s lasts in my Markov database for %s.'
            irc.reply(s % (db.lasts(channel), channel))
        self.q.enqueue(lasts)

    def pairs(self, irc, msg, args):
        """[<channel>]

        Returns the number of Markov's chain links in the database for
        <channel>.
        """
        channel = privmsgs.getChannel(msg, args)
        def pairs(db):
            s = 'There are %s pairs in my Markov database for %s.'
            irc.reply(s % (db.pairs(channel), channel))
        self.q.enqueue(pairs)

    def follows(self, irc, msg, args):
        """[<channel>]

        Returns the number of Markov's third links in the database for
        <channel>.
        """
        channel = privmsgs.getChannel(msg, args)
        def follows(db):
            s = 'There are %s follows in my Markov database for %s.'
            irc.reply(s % (db.follows(channel), channel))
        self.q.enqueue(follows)

Class = Markov

# vim:set shiftwidth=4 tabstop=8 expandtab textwidth=78:
<commit_msg>Commit my long sought after *cough*Skorobeus*cough* randomSpeaking updates<commit_after>###
# Copyright (c) 2002-2004, Jeremiah Fincher
# All rights reserved.
#
# Redistribution and use in source and binary forms, with or without
# modification, are permitted provided that the following conditions are met:
#
#   * Redistributions of source code must retain the above copyright notice,
#     this list of conditions, and the following disclaimer.
#   * Redistributions in binary form must reproduce the above copyright notice,
#     this list of conditions, and the following disclaimer in the
#     documentation and/or other materials provided with the distribution.
#   * Neither the name of the author of this software nor the name of
#     contributors to this software may be used to endorse or promote products
#     derived from this software without specific prior written consent.
#
# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
# ARE DISCLAIMED.  IN NO EVENT SHALL THE COPYRIGHT OWNER OR CONTRIBUTORS BE
# LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR
# CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF
# SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS
# INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN
# CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE)
# ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE
# POSSIBILITY OF SUCH DAMAGE.
###

"""
Silently listens to a channel, building a database of Markov Chains for
later hijinks.  To read more about Markov Chains, check out
<http://www.cs.bell-labs.com/cm/cs/pearls/sec153.html>.  When the database is
large enough, you can have it make fun little random messages from it.
"""

__revision__ = "$Id$"

import supybot.plugins as plugins

import sets
import time
import Queue
import anydbm
import random
import os.path
import threading

import supybot.conf as conf
import supybot.world as world
from supybot.commands import *
import supybot.ircmsgs as ircmsgs
import supybot.ircutils as ircutils
import supybot.registry as registry
import supybot.schedule as schedule
import supybot.callbacks as callbacks

class Probability(registry.Float):
    """Value must be a floating-point number between 0 and 1."""
    def setValue(self, v):
        if v < 0 or v > 1:
            self.error()
        else:
            registry.Float.setValue(self, float(v))

conf.registerPlugin('Markov')
conf.registerGroup(conf.supybot.plugins.Markov, 'randomSpeaking')
conf.registerChannelValue(conf.supybot.plugins.Markov.randomSpeaking,
    'probability', Probability(0, """Determines the probability that will be
    checked against to determine whether the bot should randomly say
    something.  If 0, the bot will never say anything on it's own.  If 1, the
    bot will speak every time we make a check."""))
conf.registerChannelValue(conf.supybot.plugins.Markov.randomSpeaking,
    'maxDelay', registry.PositiveInteger(10, """Determines the upper bound for
    how long the bot will wait before randomly speaking.  The delay is a
    randomly generated number of seconds below the value of this config
    variable."""))
conf.registerChannelValue(conf.supybot.plugins.Markov.randomSpeaking,
    'throttleTime', registry.PositiveInteger(300, """Determines the minimum
    number of seconds between the bot randomly speaking."""))

class MarkovDBInterface(object):
    def close(self):
        pass

    def addPair(self, channel, first, second, follower,
                isFirst=False, isLast=False):
        pass

    def getFirstPair(self, channel):
        pass

    def getPair(self, channel, first, second):
        # Returns (follower, last) tuple.
        pass

    def firsts(self, channel):
        pass

    def lasts(self, channel):
        pass

    def pairs(self, channel):
        pass

    def follows(self, channel):
        pass

class SqliteMarkovDB(object):
    def __init__(self, filename):
        self.dbs = ircutils.IrcDict()
        self.filename = filename

    def close(self):
        for db in self.dbs.values():
            db.close()

    def _getDb(self, channel):
        try:
            import sqlite
        except ImportError:
            raise callbacks.Error, 'You need to have PySQLite installed to '\
                                   'use this plugin.  Download it at '\
                                   '<http://pysqlite.sf.net/>'
        if channel not in self.dbs:
            filename = plugins.makeChannelFilename(self.filename, channel)
            if os.path.exists(filename):
                self.dbs[channel] = sqlite.connect(filename)
                return self.dbs[channel]
            #else:
            self.dbs[channel] = sqlite.connect(filename)
            cursor = self.dbs[channel].cursor()
            # TODO Finish the rest of the implementation
        return self.dbs[channel]

    def addPair(self, channel, first, second, follower,
                isFirst=False, isLast=False):
        pass

    def getFirstPair(self, channel):
        pass

    def getFollower(self, channel, first, second):
        # Returns (follower, last) tuple.
        pass

    def firsts(self, channel):
        pass

    def lasts(self, channel):
        pass

    def pairs(self, channel):
        pass

    def follows(self, channel):
        pass

class DbmMarkovDB(object):
    def __init__(self, filename):
        self.dbs = ircutils.IrcDict()
        ## Stupid anydbm seems to append .db to the end of this.
        #self.filename = filename.replace('.db', '')
        self.filename = filename

    def close(self):
        for db in self.dbs.values():
            db.close()

    def _getDb(self, channel):
        if channel not in self.dbs:
            filename = plugins.makeChannelFilename(self.filename, channel)
            # To keep the code simpler for addPair, I decided not to make
            # self.dbs[channel]['firsts'] and ['lasts'].  Instead, we'll pad
            # the words list being sent to addPair such that ['\n \n'] will be
            # ['firsts'] and ['\n'] will be ['lasts'].  This also means isFirst
            # and isLast aren't necessary, but they'll be left alone in case
            # one of the other Db formats uses them or someone decides that I
            # was wrong and changes my code.
            self.dbs[channel] = anydbm.open(filename, 'c')
        return self.dbs[channel]

    def _flush(self, db):
        if hasattr(db, 'sync'):
            db.sync()
        if hasattr(db, 'flush'):
            db.flush()

    def addPair(self, channel, first, second, follower,
                isFirst=False, isLast=False):
        db = self._getDb(channel)
        combined = self._combine(first, second)
        if db.has_key(combined): # EW!
            db[combined] = ' '.join([db[combined], follower])
        else:
            db[combined] = follower
        if follower == '\n':
            if db.has_key('\n'):
                db['\n'] = ' '.join([db['\n'], second])
            else:
                db['\n'] = second
        self._flush(db)

    def getFirstPair(self, channel):
        db = self._getDb(channel)
        firsts = db['\n \n'].split()
        if firsts:
            if firsts:
                return ('\n', random.choice(firsts))
            else:
                raise KeyError, 'No firsts for %s.' % channel
        else:
            raise KeyError, 'No firsts for %s.' % channel

    def _combine(self, first, second):
        return '%s %s' % (first, second)

    def getFollower(self, channel, first, second):
        db = self._getDb(channel)
        followers = db[self._combine(first, second)]
        follower = random.choice(followers.split(' '))
        return (follower, follower == '\n')

    def firsts(self, channel):
        db = self._getDb(channel)
        if db.has_key('\n \n'):
            return len(sets.Set(db['\n \n'].split()))
        else:
            return 0

    def lasts(self, channel):
        db = self._getDb(channel)
        if db.has_key('\n'):
            return len(sets.Set(db['\n'].split()))
        else:
            return 0

    def pairs(self, channel):
        db = self._getDb(channel)
        pairs = [k for k in db.keys() if '\n' not in k]
        return len(pairs)

    def follows(self, channel):
        db = self._getDb(channel)
        follows = [len(v.split()) for (k,v) in db.iteritems() if '\n' not in k]
        return sum(follows)

MarkovDB = plugins.DB('Markov',
                      {'anydbm': DbmMarkovDB})

class MarkovWorkQueue(threading.Thread):
    def __init__(self, *args, **kwargs):
        name = 'Thread #%s (MarkovWorkQueue)' % world.threadsSpawned
        world.threadsSpawned += 1
        threading.Thread.__init__(self, name=name)
        self.db = MarkovDB(*args, **kwargs)
        self.q = Queue.Queue()
        self.killed = False
        self.setDaemon(True)
        self.start()

    def die(self):
        self.killed = True
        self.q.put(None)

    def enqueue(self, f):
        self.q.put(f)

    def run(self):
        while not self.killed:
            f = self.q.get()
            if f is not None:
                f(self.db)
        self.db.close()

class Markov(callbacks.Privmsg):
    def __init__(self):
        self.q = MarkovWorkQueue()
        self.__parent = super(Markov, self)
        self.__parent.__init__()
        self.lastSpoke = time.time()

    def die(self):
        self.q.die()
        self.__parent.die()

    def tokenize(self, m):
        if ircmsgs.isAction(m):
            return ircmsgs.unAction(m).split()
        elif ircmsgs.isCtcp(m):
            return []
        else:
            return m.args[1].split()

    def doPrivmsg(self, irc, msg):
        channel = msg.args[0]
        if irc.isChannel(channel):
            canSpeak = False
            now = time.time()
            throttle = self.registryValue('randomSpeaking.throttleTime',
                                          channel)
            prob = self.registryValue('randomSpeaking.probability', channel)
            delay = self.registryValue('randomSpeaking.maxDelay', channel)
            irc = callbacks.SimpleProxy(irc, msg)
            if now > self.lastSpoke + throttle:
                canSpeak = True
            if canSpeak and random.random() < prob:
                f = self._markov(channel, irc, private=True, to=channel)
                schedule.addEvent(lambda: self.q.enqueue(f), now + delay)
                self.lastSpoke = now + delay
            words = self.tokenize(msg)
            words.insert(0, '\n')
            words.insert(0, '\n')
            words.append('\n')
            # This shouldn't happen often (CTCP messages being the possible exception)
            if not words or len(words) == 3:
                return
            def doPrivmsg(db):
                for (first, second, follower) in window(words, 3):
                    db.addPair(channel, first, second, follower)
            self.q.enqueue(doPrivmsg)

    def _markov(self, channel, irc, word1=None, word2=None, **kwargs):
        def f(db):
            if word1 and word2:
                givenArgs = True
                words = [word1, word2]
            else:
                givenArgs = False
                try:
                    # words is of the form ['\r', word]
                    words = list(db.getFirstPair(channel))
                except KeyError:
                    irc.error('I don\'t have any first pairs for %s.' %channel)
                    return
            follower = words[-1]
            last = False
            resp = []
            while not last:
                resp.append(follower)
                try:
                    (follower,last) = db.getFollower(channel, words[-2],
                                                     words[-1])
                except KeyError:
                    irc.error('I found a broken link in the Markov chain.  '
                              'Maybe I received two bad links to start the '
                              'chain.')
                    return
                words.append(follower)
            if givenArgs:
                irc.reply(' '.join(words[:-1]), **kwargs)
            else:
                irc.reply(' '.join(resp), **kwargs)
        return f

    def markov(self, irc, msg, args, channel, word1, word2):
        """[<channel>] [word1 word2]

        Returns a randomly-generated Markov Chain generated sentence from the
        data kept on <channel> (which is only necessary if not sent in the
        channel itself).  If word1 and word2 are specified, they will be used
        to start the Markov chain.
        """
        f = self._markov(channel, irc, word1, word2)
        self.q.enqueue(f)
    markov = wrap(markov, ['channel', optional('something'),
                           additional('something')])

    def firsts(self, irc, msg, args, channel):
        """[<channel>]

        Returns the number of Markov's first links in the database for
        <channel>.
        """
        def firsts(db):
            s = 'There are %s firsts in my Markov database for %s.'
            irc.reply(s % (db.firsts(channel), channel))
        self.q.enqueue(firsts)
    firsts = wrap(firsts, ['channel'])

    def lasts(self, irc, msg, args, channel):
        """[<channel>]

        Returns the number of Markov's last links in the database for
        <channel>.
        """
        def lasts(db):
            s = 'There are %s lasts in my Markov database for %s.'
            irc.reply(s % (db.lasts(channel), channel))
        self.q.enqueue(lasts)
    lasts = wrap(lasts, ['channel'])

    def pairs(self, irc, msg, args, channel):
        """[<channel>]

        Returns the number of Markov's chain links in the database for
        <channel>.
        """
        def pairs(db):
            s = 'There are %s pairs in my Markov database for %s.'
            irc.reply(s % (db.pairs(channel), channel))
        self.q.enqueue(pairs)
    pairs = wrap(pairs, ['channel'])

    def follows(self, irc, msg, args, channel):
        """[<channel>]

        Returns the number of Markov's third links in the database for
        <channel>.
        """
        def follows(db):
            s = 'There are %s follows in my Markov database for %s.'
            irc.reply(s % (db.follows(channel), channel))
        self.q.enqueue(follows)
    follows = wrap(follows, ['channel'])

Class = Markov

# vim:set shiftwidth=4 tabstop=8 expandtab textwidth=78:
<|endoftext|>