<commit_before>import os
import tarfile
import zipfile
import json
import redis
from datetime import datetime
# flask/external lib
from flask import Blueprint, render_template, request, jsonify, current_app, g, url_for, redirect
from rq import Queue
from werkzeug.utils import secure_filename
from flask_recaptcha import ReCaptcha
# spfy code
from modules.spfy import spfy
from routes.utility_functions import handle_tar, handle_zip, fix_uri
from modules.groupComparisons.sparql_queries import get_all_attribute_types, get_attribute_values
bp = Blueprint('main', __name__)
from modules.groupComparisons.fishers import fishers

@bp.route('/api/v0/newgroupcomparison', methods=['POST'])
def handle_group_comparison_submission():
    query = json.load(request.json)
    print query
    queryAttributeUriA = query[0]['attribute']
    queryAttributeUriB = query[1]['attribute']
    queryAttributeTypeUriA = query[0]['relation']
    queryAttributeTypeUriB = query[1]['relation']
    targetUri = 'https://www.github.com/superphy#Marker'
    f = fishers(queryAttibuteUriA, queryAttibuteUriB, targetUri, queryAttributeTypeUriA, queryAttributeTypeUriB)
    return jsonify(f.to_json())

@bp.route('/api/v0/get_attribute_values/type/<path:attributetype>')
def call_get_attribute_values(attributetype):
    '''
    Front-End API:
    Get all attribute values for a given attribute type.
    '''
    # workaround: Flask's path converter allows slashes, but only a SINGLE slash
    # this adds the second slash
    # also convert to a rdflib.URIRef object
    uri = fix_uri(attributetype)
    return jsonify(get_attribute_values(attributeTypeUri=uri))

@bp.route('/api/v0/get_all_attribute_types')
def call_get_all_atribute_types():
    '''
    Front-End API:
    Get all possible attribute types.
    '''
    return jsonify(get_all_attribute_types())

def fetch_job(job_id):
    '''
    Iterates through all queues looking for the job.
    '''
    redis_url = current_app.config['REDIS_URL']
    redis_connection = redis.from_url(redis_url)
    queues = current_app.config['QUEUES_SPFY']
    for queue in queues:
        q = Queue(queue, connection=redis_connection)
        job = q.fetch_job(job_id)
        if job is not None:
            return job
    return "fudge muffins!", 500

@bp.route('/results/<job_id>')
def job_status(job_id):
    job = fetch_job(job_id)
    if job.is_finished:
        return jsonify(job.result)
    elif job.is_failed:
        return job.exc_info, 415
    else:
        return "Still pending", 202

@bp.route('/upload', methods=['POST'])
def upload():
    recaptcha = ReCaptcha(app=current_app)
    if recaptcha.verify():
        form = request.form
        options = {}
        #defaults
        options['amr']=True
        options['vf']=True
        options['serotype']=True
        options['pi']=90

        # processing form data
        for key, value in form.items():
            #we need to convert lower-case true/false in js to upper case in python
                #remember, we also have numbers
            if not value.isdigit():
                if value.lower() == 'false':
                    value = False
                else:
                    value = True
                if key == 'options.amr':
                    options['amr']=value
                if key == 'options.vf':
                    options['vf']=value
                if key == 'options.serotype':
                    options['serotype']=value
            else:
                if key =='options.pi':
                    options['pi']=int(value)

        # get a list of files submitted
        uploaded_files = request.files.getlist("file")
        print uploaded_files

        #set up constants for identifying this sessions
        now = datetime.now()
        now = now.strftime("%Y-%m-%d-%H-%M-%S-%f")
        jobs_dict = {}

        for file in uploaded_files:
            if file:
                # for saving file
                filename = os.path.join(current_app.config[
                                        'DATASTORE'], now + '-' + secure_filename(file.filename))
                file.save(filename)
                print 'IVE SAVED YO FILE AT', str(filename)

                if tarfile.is_tarfile(filename):
                    # set filename to dir for spfy call
                    filename = handle_tar(filename, now)
                elif zipfile.is_zipfile(filename):
                    filename = handle_zip(filename, now)

                # for enqueing task
                jobs_enqueued = spfy(
                    {'i': filename, 'disable_serotype': False, 'disable_amr': False, 'disable_vf': False, 'pi':options['pi'], 'options':options})
                jobs_dict.update(jobs_enqueued)

                d = dict(jobs_dict)
                #strip jobs that the user doesn't want to see
                # we run them anyways cause we want the data analyzed on our end
                for job_id, descrip_dict in jobs_dict.items():
                    if (not options['serotype']) and (not options['vf']):
                        if descrip_dict['analysis'] == 'Virulence Factors and Serotype':
                            del d[job_id]
                    if (not options['amr']):
                        if descrip_dict['analysis'] == 'Antimicrobial Resistance':
                            del d[job_id]
                jobs_dict = d
        return jsonify(jobs_dict)
    else:
        return "Captcha Failed Verification", 500

@bp.route('/', methods=['GET', 'POST'])
def index():
    return render_template("index.html")
<commit_msg>FIX: flask already converts json to python<commit_after>import os
import tarfile
import zipfile
import json
import redis
from datetime import datetime
# flask/external lib
from flask import Blueprint, render_template, request, jsonify, current_app, g, url_for, redirect
from rq import Queue
from werkzeug.utils import secure_filename
from flask_recaptcha import ReCaptcha
# spfy code
from modules.spfy import spfy
from routes.utility_functions import handle_tar, handle_zip, fix_uri
from modules.groupComparisons.sparql_queries import get_all_attribute_types, get_attribute_values
bp = Blueprint('main', __name__)
from modules.groupComparisons.fishers import fishers

@bp.route('/api/v0/newgroupcomparison', methods=['POST'])
def handle_group_comparison_submission():
    query = json.load
    print query
    queryAttributeUriA = query[0]['attribute']
    queryAttributeUriB = query[1]['attribute']
    queryAttributeTypeUriA = query[0]['relation']
    queryAttributeTypeUriB = query[1]['relation']
    targetUri = 'https://www.github.com/superphy#Marker'
    f = fishers(queryAttibuteUriA, queryAttibuteUriB, targetUri, queryAttributeTypeUriA, queryAttributeTypeUriB)
    return jsonify(f.to_json())

@bp.route('/api/v0/get_attribute_values/type/<path:attributetype>')
def call_get_attribute_values(attributetype):
    '''
    Front-End API:
    Get all attribute values for a given attribute type.
    '''
    # workaround: Flask's path converter allows slashes, but only a SINGLE slash
    # this adds the second slash
    # also convert to a rdflib.URIRef object
    uri = fix_uri(attributetype)
    return jsonify(get_attribute_values(attributeTypeUri=uri))

@bp.route('/api/v0/get_all_attribute_types')
def call_get_all_atribute_types():
    '''
    Front-End API:
    Get all possible attribute types.
    '''
    return jsonify(get_all_attribute_types())

def fetch_job(job_id):
    '''
    Iterates through all queues looking for the job.
    '''
    redis_url = current_app.config['REDIS_URL']
    redis_connection = redis.from_url(redis_url)
    queues = current_app.config['QUEUES_SPFY']
    for queue in queues:
        q = Queue(queue, connection=redis_connection)
        job = q.fetch_job(job_id)
        if job is not None:
            return job
    return "fudge muffins!", 500

@bp.route('/results/<job_id>')
def job_status(job_id):
    job = fetch_job(job_id)
    if job.is_finished:
        return jsonify(job.result)
    elif job.is_failed:
        return job.exc_info, 415
    else:
        return "Still pending", 202

@bp.route('/upload', methods=['POST'])
def upload():
    recaptcha = ReCaptcha(app=current_app)
    if recaptcha.verify():
        form = request.form
        options = {}
        #defaults
        options['amr']=True
        options['vf']=True
        options['serotype']=True
        options['pi']=90

        # processing form data
        for key, value in form.items():
            #we need to convert lower-case true/false in js to upper case in python
                #remember, we also have numbers
            if not value.isdigit():
                if value.lower() == 'false':
                    value = False
                else:
                    value = True
                if key == 'options.amr':
                    options['amr']=value
                if key == 'options.vf':
                    options['vf']=value
                if key == 'options.serotype':
                    options['serotype']=value
            else:
                if key =='options.pi':
                    options['pi']=int(value)

        # get a list of files submitted
        uploaded_files = request.files.getlist("file")
        print uploaded_files

        #set up constants for identifying this sessions
        now = datetime.now()
        now = now.strftime("%Y-%m-%d-%H-%M-%S-%f")
        jobs_dict = {}

        for file in uploaded_files:
            if file:
                # for saving file
                filename = os.path.join(current_app.config[
                                        'DATASTORE'], now + '-' + secure_filename(file.filename))
                file.save(filename)
                print 'IVE SAVED YO FILE AT', str(filename)

                if tarfile.is_tarfile(filename):
                    # set filename to dir for spfy call
                    filename = handle_tar(filename, now)
                elif zipfile.is_zipfile(filename):
                    filename = handle_zip(filename, now)

                # for enqueing task
                jobs_enqueued = spfy(
                    {'i': filename, 'disable_serotype': False, 'disable_amr': False, 'disable_vf': False, 'pi':options['pi'], 'options':options})
                jobs_dict.update(jobs_enqueued)

                d = dict(jobs_dict)
                #strip jobs that the user doesn't want to see
                # we run them anyways cause we want the data analyzed on our end
                for job_id, descrip_dict in jobs_dict.items():
                    if (not options['serotype']) and (not options['vf']):
                        if descrip_dict['analysis'] == 'Virulence Factors and Serotype':
                            del d[job_id]
                    if (not options['amr']):
                        if descrip_dict['analysis'] == 'Antimicrobial Resistance':
                            del d[job_id]
                jobs_dict = d
        return jsonify(jobs_dict)
    else:
        return "Captcha Failed Verification", 500

@bp.route('/', methods=['GET', 'POST'])
def index():
    return render_template("index.html")
<|endoftext|>


<commit_before>import logging
import typing

import enum

import functools

import itertools

from katagawa import kg as md_kg
from katagawa.orm import query as md_query
from katagawa.backends.base import BaseTransaction
from katagawa.orm.schema import TableRow
from katagawa.orm.schema.row import NO_VALUE

logger = logging.getLogger(__name__)

class SessionState(enum.Enum):
    NOT_READY = 0
    READY = 1
    CLOSED = 2

# decorators
def enforce_open(func):
    @functools.wraps(func)
    def wrapper(self, *args, **kwargs):
        if self._state is not SessionState.READY:
            raise RuntimeError("Session is not ready or closed")
        else:
            return func(self, *args, **kwargs)

    return wrapper

class Session(object):
    """
    Sessions act as a temporary window into the database. They are responsible for creating queries,
    inserting and updating rows, etc.

    Sessions are bound to a :class:`.Katagawa` instance which they use to get a transaction and
    execute queries in.

    .. code-block:: python
        # get a session from our db interface
        sess = db.get_session()
    """

    def __init__(self, bind: 'md_kg.Katagawa'):
        """
        :param bind: The :class:`.Katagawa` instance we are bound to.
        """
        self.bind = bind

        #: The current state for the session.
        self._state = SessionState.NOT_READY

        #: The current :class:`.BaseTransaction` this Session is associated with.
        #: The transaction is used for making queries and inserts, etc.
        self.transaction = None  # type: BaseTransaction

        #: The current list of "new" objects.
        #: These are table rows that are ready to be inserted.
        self.new = []

        #: The current list of "dirty" objects.
        #: These are table rows that are ready to be updated.
        self.dirty = []

        #: The current list of "deleted" objects.
        #: These are table rows that are ready to be deleted.
        self.deleted = []

    async def __aenter__(self) -> 'Session':
        await self.start()
        return self

    async def __aexit__(self, exc_type, exc_val, exc_tb) -> 'Session':
        try:
            if exc_type is None:
                await self.commit()
            else:
                await self.rollback()
        finally:
            await self.close()

        return False

    # Query builders
    def select(self, table) -> 'md_query.SelectQuery':
        """
        Creates a new SELECT query that can be built upon.

        :param table: The :class:`.Table` to select.
        :return: A new :class:`.SelectQuery`.
        """
        q = md_query.SelectQuery(self)
        q.set_table(table)
        return q

    def _generate_inserts(self) -> typing.List[typing.Tuple[str, typing.Dict[str, typing.Any]]]:
        """
        Generates INSERT INTO queries for the current session.
        """
        queries = []
        counter = itertools.count()

        # group the rows
        for tbl, rows in itertools.groupby(self.new, lambda r: r._table):
            rows = list(rows)
            # insert into the quoted table
            base_query = "INSERT INTO {} ".format(tbl.__quoted_name__)
            # add the columns (quoted names)
            base_query += "({})".format(", ".join(column.quoted_name
                                                  for column in tbl.iter_columns()))
            # get the values
            base_query += " VALUES "

            # build the params dict and VALUES section
            # this is done by looping over the rows
            # looping over the columns of the row's table
            # checking for the new value, then emitting it
            # if no value is available, it will emit DEFAULT
            params = {}
            value_sets = []
            for row in rows:
                prms_so_far = []
                assert isinstance(row, TableRow)
                #row._validate()
                for column in tbl.iter_columns():
                    value = row.get_column_value(column, return_default=False)
                    if value is NO_VALUE:
                        prms_so_far.append("DEFAULT")
                    else:
                        # emit a new param
                        number = next(counter)
                        name = "param_{}".format(number)
                        param_name = self.bind.emit_param(name)
                        # set the params to value
                        # then add the {param_name} to the VALUES
                        params[name] = value
                        prms_so_far.append(param_name)

                # join the params together
                value_sets.append("({})".format(", ".join(prms_so_far)))

            # join all of the value sets together
            base_query += ", ".join(value_sets)
            base_query += ";"
            queries.append((base_query, params, rows))

        return queries

    async def start(self) -> 'Session':
        """
        Starts the session, acquiring a transaction connection which will be used to modify the DB.
        This **must** be called before using the session.

        .. code-block:: python
            sess = db.get_session()
            await sess.start()

        .. note::
            When using ``async with``, this is automatically called.
        """
        if self._state is not SessionState.NOT_READY:
            raise RuntimeError("Session must not be ready or closed")

        self.transaction = self.bind.get_transaction()
        await self.transaction.begin()

        self._state = SessionState.READY
        return self

    @enforce_open
    async def commit(self):
        """
        Commits the current session, running inserts/updates/deletes.

        This will **not** close the session; it can be re-used after a commit.
        """
        # TODO: Update generation
        # TODO: Checkpoints
        inserts = self._generate_inserts()
        for query, params, rows in inserts:
            await self.execute(query, params=params)
            for row in rows:
                self.new.remove(row)

        await self.transaction.commit()
        return self

    @enforce_open
    async def rollback(self, checkpoint: str = None):
        """
        Rolls the current session back.
        This is useful if an error occurs inside your code.

        :param checkpoint: The checkpoint to roll back to, if applicable.
        """
        await self.transaction.rollback(checkpoint=checkpoint)
        return self

    @enforce_open
    async def close(self):
        """
        Closes the current session.

        .. warning::
            This will **NOT COMMIT ANY DATA**. Old data will die.
        """
        await self.transaction.close()
        self._state = SessionState.CLOSED
        del self.transaction

    @enforce_open
    async def execute(self, sql: str, params: typing.Union[typing.Mapping[str, typing.Any],
                                                           typing.Iterable[typing.Any]] = None):
        """
        Executes SQL inside the current session.

        This is part of the **low-level API.**

        :param sql: The SQL to execute.
        :param params: The parameters to use inside the query.
        """
        return await self.transaction.execute(sql, params)

    @enforce_open
    async def cursor(self, sql: str, params: typing.Union[typing.Mapping[str, typing.Any],
                                                          typing.Iterable[typing.Any]] = None):
        """
        Executes SQL inside the current session, and returns a new :class:`.BaseResultSet.`

        :param sql: The SQL to execute.
        :param params: The parameters to use inside the query.
        """
        return await self.transaction.cursor(sql, params)

    def add(self, row: TableRow) -> 'Session':
        """
        Adds a row to this session, storing it for a later commit.

        :param row: The :class:`.TableRow` to add to this session.
        :return: This session.
        """
        if row._TableRow__existed is True:
            self.dirty.append(row)
        else:
            self.new.append(row)

        return self

    insert = add
<commit_msg>Rollback and re-raise in the case that commit() fails.<commit_after>import logging
import typing

import enum

import functools

import itertools

from katagawa import kg as md_kg
from katagawa.orm import query as md_query
from katagawa.backends.base import BaseTransaction
from katagawa.orm.schema import TableRow
from katagawa.orm.schema.row import NO_VALUE

logger = logging.getLogger(__name__)

class SessionState(enum.Enum):
    NOT_READY = 0
    READY = 1
    CLOSED = 2

# decorators
def enforce_open(func):
    @functools.wraps(func)
    def wrapper(self, *args, **kwargs):
        if self._state is not SessionState.READY:
            raise RuntimeError("Session is not ready or closed")
        else:
            return func(self, *args, **kwargs)

    return wrapper

class Session(object):
    """
    Sessions act as a temporary window into the database. They are responsible for creating queries,
    inserting and updating rows, etc.

    Sessions are bound to a :class:`.Katagawa` instance which they use to get a transaction and
    execute queries in.

    .. code-block:: python
        # get a session from our db interface
        sess = db.get_session()
    """

    def __init__(self, bind: 'md_kg.Katagawa'):
        """
        :param bind: The :class:`.Katagawa` instance we are bound to.
        """
        self.bind = bind

        #: The current state for the session.
        self._state = SessionState.NOT_READY

        #: The current :class:`.BaseTransaction` this Session is associated with.
        #: The transaction is used for making queries and inserts, etc.
        self.transaction = None  # type: BaseTransaction

        #: The current list of "new" objects.
        #: These are table rows that are ready to be inserted.
        self.new = []

        #: The current list of "dirty" objects.
        #: These are table rows that are ready to be updated.
        self.dirty = []

        #: The current list of "deleted" objects.
        #: These are table rows that are ready to be deleted.
        self.deleted = []

    async def __aenter__(self) -> 'Session':
        await self.start()
        return self

    async def __aexit__(self, exc_type, exc_val, exc_tb) -> 'Session':
        try:
            if exc_type is None:
                await self.commit()
            else:
                await self.rollback()
        except Exception:
            # error in committing, etc
            await self.rollback()
            raise
        finally:
            await self.close()

        return False

    # Query builders
    def select(self, table) -> 'md_query.SelectQuery':
        """
        Creates a new SELECT query that can be built upon.

        :param table: The :class:`.Table` to select.
        :return: A new :class:`.SelectQuery`.
        """
        q = md_query.SelectQuery(self)
        q.set_table(table)
        return q

    def _generate_inserts(self) -> typing.List[typing.Tuple[str, typing.Dict[str, typing.Any]]]:
        """
        Generates INSERT INTO queries for the current session.
        """
        queries = []
        counter = itertools.count()

        # group the rows
        for tbl, rows in itertools.groupby(self.new, lambda r: r._table):
            rows = list(rows)
            # insert into the quoted table
            base_query = "INSERT INTO {} ".format(tbl.__quoted_name__)
            # add the columns (quoted names)
            base_query += "({})".format(", ".join(column.quoted_name
                                                  for column in tbl.iter_columns()))
            # get the values
            base_query += " VALUES "

            # build the params dict and VALUES section
            # this is done by looping over the rows
            # looping over the columns of the row's table
            # checking for the new value, then emitting it
            # if no value is available, it will emit DEFAULT
            params = {}
            value_sets = []
            for row in rows:
                prms_so_far = []
                assert isinstance(row, TableRow)
                #row._validate()
                for column in tbl.iter_columns():
                    value = row.get_column_value(column, return_default=False)
                    if value is NO_VALUE:
                        prms_so_far.append("DEFAULT")
                    else:
                        # emit a new param
                        number = next(counter)
                        name = "param_{}".format(number)
                        param_name = self.bind.emit_param(name)
                        # set the params to value
                        # then add the {param_name} to the VALUES
                        params[name] = value
                        prms_so_far.append(param_name)

                # join the params together
                value_sets.append("({})".format(", ".join(prms_so_far)))

            # join all of the value sets together
            base_query += ", ".join(value_sets)
            base_query += ";"
            queries.append((base_query, params, rows))

        return queries

    async def start(self) -> 'Session':
        """
        Starts the session, acquiring a transaction connection which will be used to modify the DB.
        This **must** be called before using the session.

        .. code-block:: python
            sess = db.get_session()
            await sess.start()

        .. note::
            When using ``async with``, this is automatically called.
        """
        if self._state is not SessionState.NOT_READY:
            raise RuntimeError("Session must not be ready or closed")

        self.transaction = self.bind.get_transaction()
        await self.transaction.begin()

        self._state = SessionState.READY
        return self

    @enforce_open
    async def commit(self):
        """
        Commits the current session, running inserts/updates/deletes.

        This will **not** close the session; it can be re-used after a commit.
        """
        # TODO: Update generation
        # TODO: Checkpoints
        inserts = self._generate_inserts()
        for query, params, rows in inserts:
            await self.execute(query, params=params)
            for row in rows:
                self.new.remove(row)

        await self.transaction.commit()
        return self

    @enforce_open
    async def rollback(self, checkpoint: str = None):
        """
        Rolls the current session back.
        This is useful if an error occurs inside your code.

        :param checkpoint: The checkpoint to roll back to, if applicable.
        """
        await self.transaction.rollback(checkpoint=checkpoint)
        return self

    @enforce_open
    async def close(self):
        """
        Closes the current session.

        .. warning::
            This will **NOT COMMIT ANY DATA**. Old data will die.
        """
        await self.transaction.close()
        self._state = SessionState.CLOSED
        del self.transaction

    @enforce_open
    async def execute(self, sql: str, params: typing.Union[typing.Mapping[str, typing.Any],
                                                           typing.Iterable[typing.Any]] = None):
        """
        Executes SQL inside the current session.

        This is part of the **low-level API.**

        :param sql: The SQL to execute.
        :param params: The parameters to use inside the query.
        """
        return await self.transaction.execute(sql, params)

    @enforce_open
    async def cursor(self, sql: str, params: typing.Union[typing.Mapping[str, typing.Any],
                                                          typing.Iterable[typing.Any]] = None):
        """
        Executes SQL inside the current session, and returns a new :class:`.BaseResultSet.`

        :param sql: The SQL to execute.
        :param params: The parameters to use inside the query.
        """
        return await self.transaction.cursor(sql, params)

    def add(self, row: TableRow) -> 'Session':
        """
        Adds a row to this session, storing it for a later commit.

        :param row: The :class:`.TableRow` to add to this session.
        :return: This session.
        """
        if row._TableRow__existed is True:
            self.dirty.append(row)
        else:
            self.new.append(row)

        return self

    insert = add
<|endoftext|>


<commit_before>#
#  Copyright (C) 2010-2011, 2011 Canonical Ltd. All Rights Reserved
#
#  This file is part of txzookeeper.
#
#  Authors:
#   Kapil Thangavelu
#
#  txzookeeper is free software: you can redistribute it and/or modify
#  it under the terms of the GNU Lesser General Public License as published by
#  the Free Software Foundation, either version 3 of the License, or
#  (at your option) any later version.
#
#  txzookeeper 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 Lesser General Public License for more details.
#
#  You should have received a copy of the GNU Lesser General Public License
#  along with txzookeeper.  If not, see <http://www.gnu.org/licenses/>.
#

import os
import os.path
import shutil
import subprocess
import tempfile

from itertools import chain
from collections import namedtuple
from glob import glob

ServerInfo = namedtuple(
    "ServerInfo", "server_id client_port election_port leader_port")

class ManagedZooKeeper(object):
    """Class to manage the running of a ZooKeeper instance for testing.

    Note: no attempt is made to probe the ZooKeeper instance is
    actually available, or that the selected port is free. In the
    future, we may want to do that, especially when run in a
    Hudson/Buildbot context, to ensure more test robustness."""

    def __init__(self, software_path, server_info, peers=()):
        """Define the ZooKeeper test instance.

        @param install_path: The path to the install for ZK
        @param port: The port to run the managed ZK instance
        """
        self.install_path = software_path
        self.server_info = server_info
        self.host = "127.0.0.1"
        self.peers = peers
        self.working_path = tempfile.mkdtemp()
        self._running = False

    def run(self):
        """Run the ZooKeeper instance under a temporary directory.

        Writes ZK log messages to zookeeper.log in the current directory.
        """
        if self.running:
            return
        config_path = os.path.join(self.working_path, "zoo.cfg")
        log_path = os.path.join(self.working_path, "log")
        log4j_path = os.path.join(self.working_path, "log4j.properties")
        data_path = os.path.join(self.working_path, "data")

        # various setup steps
        if not os.path.exists(self.working_path):
            os.mkdir(self.working_path)
        if not os.path.exists(log_path):
            os.mkdir(log_path)
        if not os.path.exists(data_path):
            os.mkdir(data_path)

        with open(config_path, "w") as config:
            config.write("""
tickTime=2000
dataDir=%s
clientPort=%s
maxClientCnxns=0
""" % (data_path, self.server_info.client_port))

        # setup a replicated setup if peers are specified
        if self.peers:
            servers_cfg = []
            for p in chain((self.server_info,), self.peers):
                servers_cfg.append("server.%s=localhost:%s:%s" % (
                    p.server_id, p.leader_port, p.election_port))

            with open(config_path, "a") as config:
                config.write("""
initLimit=4
syncLimit=2
%s
""" % ("\n".join(servers_cfg)))

        # Write server ids into datadir
        with open(os.path.join(data_path, "myid"), "w") as myid_file:
            myid_file.write(str(self.server_info.server_id))

        with open(log4j_path, "w") as log4j:
            log4j.write("""
# DEFAULT: console appender only
log4j.rootLogger=INFO, ROLLINGFILE
log4j.appender.ROLLINGFILE.layout=org.apache.log4j.PatternLayout
log4j.appender.ROLLINGFILE.layout.ConversionPattern=%d{ISO8601} [myid:%X{myid}] - %-5p [%t:%C{1}@%L] - %m%n
log4j.appender.ROLLINGFILE=org.apache.log4j.RollingFileAppender
log4j.appender.ROLLINGFILE.Threshold=DEBUG
log4j.appender.ROLLINGFILE.File=""" + (
                self.working_path + os.sep + "zookeeper.log\n"))

        self.process = subprocess.Popen(
            args=["java",
                  "-cp", self.classpath,
                  "-Dreadonlymode.enabled=true",
                  "-Dzookeeper.log.dir=%s" % log_path,
                  "-Dzookeeper.root.logger=INFO,CONSOLE",
                  "-Dlog4j.configuration=file:%s" % log4j_path,
                  # "-Dlog4j.debug",
                  "org.apache.zookeeper.server.quorum.QuorumPeerMain",
                  config_path],
            )
        self._running = True

    @property
    def classpath(self):
        """Get the classpath necessary to run ZooKeeper."""

        # Two possibilities, as seen in zkEnv.sh:
        # Check for a release - top-level zookeeper-*.jar?
        jars = glob((os.path.join(
            self.install_path, 'zookeeper-*.jar')))
        if jars:
            # Release build (`ant package`)
            jars.extend(glob(os.path.join(
                self.install_path,
                "lib/*.jar")))
        else:
            # Development build (plain `ant`)
            jars = glob((os.path.join(
                self.install_path, 'build/zookeeper-*.jar')))
            jars.extend(glob(os.path.join(
                self.install_path,
                "build/lib/*.jar")))
        return ":".join(jars)

    @property
    def address(self):
        """Get the address of the ZooKeeper instance."""
        return "%s:%s" % (self.host, self.client_port)

    @property
    def running(self):
        return self._running

    @property
    def client_port(self):
        return self.server_info.client_port

    def reset(self):
        """Stop the zookeeper instance, cleaning out its on disk-data."""
        self.stop()
        shutil.rmtree(os.path.join(self.working_path, "data"))
        os.mkdir(os.path.join(self.working_path, "data"))
        with open(os.path.join(self.working_path, "data", "myid"), "w") as fh:
            fh.write(str(self.server_info.server_id))

    def stop(self):
        """Stop the Zookeeper instance, retaining on disk state."""
        if not self.running:
            return
        self.process.terminate()
        self.process.wait()
        self._running = False

    def destroy(self):
        """Stop the ZooKeeper instance and destroy its on disk-state"""
        # called by at exit handler, reimport to avoid cleanup race.
        import shutil
        self.stop()

        shutil.rmtree(self.working_path)

class ZookeeperCluster(object):

    def __init__(self, install_path, size=3, port_offset=20000):
        self._install_path = install_path
        self._servers = []

        # Calculate ports and peer group
        port = port_offset
        peers = []

        for i in range(size):
            port += i * 10
            info = ServerInfo(i + 1, port, port + 1, port + 2)
            peers.append(info)

        # Instantiate Managed ZK Servers
        for i in range(size):
            server_peers = list(peers)
            server_info = server_peers.pop(i)
            self._servers.append(
                ManagedZooKeeper(
                    self._install_path, server_info, server_peers))

    def __getitem__(self, k):
        return self._servers[k]

    def __iter__(self):
        return iter(self._servers)

    def start(self):
        # Zookeeper client expresses a preference for either lower ports or
        # lexicographical ordering of hosts, to ensure that all servers have a
        # chance to startup, start them in reverse order.
        for server in reversed(list(self)):
            server.run()
        # Giving the servers a moment to start, decreases the overall time
        # required for a client to successfully connect (2s vs. 4s without
        # the sleep).
        import time
        time.sleep(2)

    def stop(self):
        for server in self:
            server.stop()
        self._servers = []

    def terminate(self):
        for server in self:
            server.destroy()

    def reset(self):
        for server in self:
            server.reset()
<commit_msg>Added path modifiers for java to pick up on Windows<commit_after>#
#  Copyright (C) 2010-2011, 2011 Canonical Ltd. All Rights Reserved
#
#  This file is part of txzookeeper.
#
#  Authors:
#   Kapil Thangavelu
#
#  txzookeeper is free software: you can redistribute it and/or modify
#  it under the terms of the GNU Lesser General Public License as published by
#  the Free Software Foundation, either version 3 of the License, or
#  (at your option) any later version.
#
#  txzookeeper 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 Lesser General Public License for more details.
#
#  You should have received a copy of the GNU Lesser General Public License
#  along with txzookeeper.  If not, see <http://www.gnu.org/licenses/>.
#

import os
import os.path
import shutil
import subprocess
import tempfile

from itertools import chain
from collections import namedtuple
from glob import glob

if os.name == 'nt':
    class_path_separator = ';'
else:
    class_path_separator = ':'

def to_java_compatible_path(path):
    if os.name == 'nt':
        path = path.replace('\\', '/')
    return path

ServerInfo = namedtuple(
    "ServerInfo", "server_id client_port election_port leader_port")

class ManagedZooKeeper(object):
    """Class to manage the running of a ZooKeeper instance for testing.

    Note: no attempt is made to probe the ZooKeeper instance is
    actually available, or that the selected port is free. In the
    future, we may want to do that, especially when run in a
    Hudson/Buildbot context, to ensure more test robustness."""

    def __init__(self, software_path, server_info, peers=()):
        """Define the ZooKeeper test instance.

        @param install_path: The path to the install for ZK
        @param port: The port to run the managed ZK instance
        """
        self.install_path = software_path
        self.server_info = server_info
        self.host = "127.0.0.1"
        self.peers = peers
        self.working_path = tempfile.mkdtemp()
        self._running = False

    def run(self):
        """Run the ZooKeeper instance under a temporary directory.

        Writes ZK log messages to zookeeper.log in the current directory.
        """
        if self.running:
            return
        config_path = os.path.join(self.working_path, "zoo.cfg")
        log_path = os.path.join(self.working_path, "log")
        log4j_path = os.path.join(self.working_path, "log4j.properties")
        data_path = os.path.join(self.working_path, "data")

        # various setup steps
        if not os.path.exists(self.working_path):
            os.mkdir(self.working_path)
        if not os.path.exists(log_path):
            os.mkdir(log_path)
        if not os.path.exists(data_path):
            os.mkdir(data_path)

        with open(config_path, "wb") as config:
            config.write("""
tickTime=2000
dataDir=%s
clientPort=%s
maxClientCnxns=0
""" % (to_java_compatible_path(data_path), self.server_info.client_port))

        # setup a replicated setup if peers are specified
        if self.peers:
            servers_cfg = []
            for p in chain((self.server_info,), self.peers):
                servers_cfg.append("server.%s=localhost:%s:%s" % (
                    p.server_id, p.leader_port, p.election_port))

            with open(config_path, "a") as config:
                config.write("""
initLimit=4
syncLimit=2
%s
""" % ("\n".join(servers_cfg)))

        # Write server ids into datadir
        with open(os.path.join(data_path, "myid"), "w") as myid_file:
            myid_file.write(str(self.server_info.server_id))

        with open(log4j_path, "wb") as log4j:
            log4j.write("""
# DEFAULT: console appender only
log4j.rootLogger=INFO, ROLLINGFILE
log4j.appender.ROLLINGFILE.layout=org.apache.log4j.PatternLayout
log4j.appender.ROLLINGFILE.layout.ConversionPattern=%d{ISO8601} [myid:%X{myid}] - %-5p [%t:%C{1}@%L] - %m%n
log4j.appender.ROLLINGFILE=org.apache.log4j.RollingFileAppender
log4j.appender.ROLLINGFILE.Threshold=DEBUG
log4j.appender.ROLLINGFILE.File=""" + to_java_compatible_path(
                self.working_path + os.sep + "zookeeper.log\n"))

        self.process = subprocess.Popen(
            args=["java",
                  "-cp", self.classpath,
                  "-Dreadonlymode.enabled=true",
                  "-Dzookeeper.log.dir=%s" % log_path,
                  "-Dzookeeper.root.logger=INFO,CONSOLE",
                  "-Dlog4j.configuration=file:%s" % log4j_path,
                  # "-Dlog4j.debug",
                  "org.apache.zookeeper.server.quorum.QuorumPeerMain",
                  config_path])
        self._running = True

    @property
    def classpath(self):
        """Get the classpath necessary to run ZooKeeper."""

        # Two possibilities, as seen in zkEnv.sh:
        # Check for a release - top-level zookeeper-*.jar?
        jars = glob((os.path.join(
            self.install_path, 'zookeeper-*.jar')))
        if jars:
            # Release build (`ant package`)
            jars.extend(glob(os.path.join(
                self.install_path,
                "lib/*.jar")))
        else:
            # Development build (plain `ant`)
            jars = glob((os.path.join(
                self.install_path, 'build/zookeeper-*.jar')))
            jars.extend(glob(os.path.join(
                self.install_path,
                "build/lib/*.jar")))

        return class_path_separator.join(jars)

    @property
    def address(self):
        """Get the address of the ZooKeeper instance."""
        return "%s:%s" % (self.host, self.client_port)

    @property
    def running(self):
        return self._running

    @property
    def client_port(self):
        return self.server_info.client_port

    def reset(self):
        """Stop the zookeeper instance, cleaning out its on disk-data."""
        self.stop()
        shutil.rmtree(os.path.join(self.working_path, "data"))
        os.mkdir(os.path.join(self.working_path, "data"))
        with open(os.path.join(self.working_path, "data", "myid"), "w") as fh:
            fh.write(str(self.server_info.server_id))

    def stop(self):
        """Stop the Zookeeper instance, retaining on disk state."""
        if not self.running:
            return
        self.process.terminate()
        self.process.wait()
        self._running = False

    def destroy(self):
        """Stop the ZooKeeper instance and destroy its on disk-state"""
        # called by at exit handler, reimport to avoid cleanup race.
        import shutil
        self.stop()

        shutil.rmtree(self.working_path)

class ZookeeperCluster(object):

    def __init__(self, install_path, size=3, port_offset=20000):
        self._install_path = install_path
        self._servers = []

        # Calculate ports and peer group
        port = port_offset
        peers = []

        for i in range(size):
            port += i * 10
            info = ServerInfo(i + 1, port, port + 1, port + 2)
            peers.append(info)

        # Instantiate Managed ZK Servers
        for i in range(size):
            server_peers = list(peers)
            server_info = server_peers.pop(i)
            self._servers.append(
                ManagedZooKeeper(
                    self._install_path, server_info, server_peers))

    def __getitem__(self, k):
        return self._servers[k]

    def __iter__(self):
        return iter(self._servers)

    def start(self):
        # Zookeeper client expresses a preference for either lower ports or
        # lexicographical ordering of hosts, to ensure that all servers have a
        # chance to startup, start them in reverse order.
        for server in reversed(list(self)):
            server.run()
        # Giving the servers a moment to start, decreases the overall time
        # required for a client to successfully connect (2s vs. 4s without
        # the sleep).
        import time
        time.sleep(2)

    def stop(self):
        for server in self:
            server.stop()
        self._servers = []

    def terminate(self):
        for server in self:
            server.destroy()

    def reset(self):
        for server in self:
            server.reset()
<|endoftext|>


<commit_before>#!/usr/bin/python

# Comments and/or additions are welcome (send e-mail to:
# robert.paton@chem.ox.ac.uk

#####################################
#          Hess_to_Freq.py          #
#####################################
###  Written by:  Rob Paton #########
###  Last modified:  Mar 27, 2017 ###
#####################################

import sys, math
import numpy as np

def read_hess(file, iso):
   # The force constant matrix is read from g09 ouptut
   # The matrix values are mass-weighted according to the isotopic masses
   # Vibrational scaling factors are not applied to matrix elements at this stage:
   # the resulting frequencies can be scaled after diagonalization

   # Read gaussian output
   g09_output = open(file, 'r')

   inlines = g09_output.readlines()

   mass_list = []
   start_force = len(inlines) < 100
   for i in range(0,len(inlines)):
      if inlines[i].strip().find('and mass') > -1:
         mass_list.append(float(inlines[i].strip().split()[8]))

      if inlines[i].strip().startswith('NAtoms='):
         d_o_f = int(inlines[i].strip().split()[1])*3

      if inlines[i].strip().find('NImag') > -1: start_force = i

   # Hessian and mass-weighted Hessian (3N x 3N matrices, N = no. atoms)
   hess_mat = np.ndarray(shape=(d_o_f,d_o_f))
   mw_hess_mat = np.ndarray(shape=(d_o_f,d_o_f))

   # Isotopic substitution will consider 1H/2H, 12C/13C and 16O/17O. More can be
   # added, but it hasn't been necessary so far...
   isolist = iso.split()
   for i in range(0,len(mass_list)):
      for atom in isolist:
         if i == int(atom)-1:
            if mass_list[i] == 1.00783: mass_list[i] = 2.0141 # protium - deuterium
            if mass_list[i] == 12.00000: mass_list[i] = 13.00335 # 12C - 13C
            if mass_list[i] == 15.99491: mass_list[i] = 16.9991 # 16O - 17O
   longline = ""
   for i in range(start_force,len(inlines)):
      longline = longline + inlines[i].rstrip().lstrip()

   forces = longline.split("NImag")[1].split('\\')[2].split(',')
   n = 0; l = -1
   for m in range (0, d_o_f):
      for n in range(0,m+1):
         l = l + 1
         Hmn = forces[l]
         sqrt_Mmn = (mass_list[m/3] * mass_list[n/3]) ** 0.5
         hess_mat[n,m] = float(Hmn)
         mw_hess_mat[m,n] = ((float(Hmn) / sqrt_Mmn))
         mw_hess_mat[n,m] = ((float(Hmn) / sqrt_Mmn))

   return mw_hess_mat

if __name__ == "__main__":
   # Takes arguments: g09_output_files and optional temperature and vibrational scaling factor
   files = []
   freq_scale_factor = 1.0

   if len(sys.argv) > 1:
      for i in range(1,len(sys.argv)):
         if sys.argv[i] == "-s": freq_scale_factor = float(sys.argv[i+1])
         else:
            if len(sys.argv[i].split(".")) > 1:
               if sys.argv[i].split(".")[1] == "out" or sys.argv[i].split(".")[1] == "log": files.append(sys.argv[i])

      print "   Frequency scale factor =", freq_scale_factor

   else:
      print "\nWrong number of arguments used. Correct format: Hess_to_Freq.py (-s scalefactor) g09_output_file(s)\n"
      sys.exit()

   for file in files:
      mw_hessmat = read_hess(file, 0)
      eigs = np.linalg.eigvalsh(mw_hessmat)
      for eig in eigs:
         if eig > 0: print eig ** 0.5
         else: print -1 * -eig ** 0.5
<commit_msg>fixed old python2 print statements<commit_after>#!/usr/bin/python

# Comments and/or additions are welcome (send e-mail to:
# robert.paton@chem.ox.ac.uk

#####################################
#          Hess_to_Freq.py          #
#####################################
###  Written by:  Rob Paton #########
###  Last modified:  Mar 27, 2017 ###
#####################################

import sys, math
import numpy as np

def read_hess(file, iso):
   # The force constant matrix is read from g09 ouptut
   # The matrix values are mass-weighted according to the isotopic masses
   # Vibrational scaling factors are not applied to matrix elements at this stage:
   # the resulting frequencies can be scaled after diagonalization

   # Read gaussian output
   g09_output = open(file, 'r')

   inlines = g09_output.readlines()

   mass_list = []
   start_force = len(inlines) < 100
   for i in range(0,len(inlines)):
      if inlines[i].strip().find('and mass') > -1:
         mass_list.append(float(inlines[i].strip().split()[8]))

      if inlines[i].strip().startswith('NAtoms='):
         d_o_f = int(inlines[i].strip().split()[1])*3

      if inlines[i].strip().find('NImag') > -1: start_force = i

   # Hessian and mass-weighted Hessian (3N x 3N matrices, N = no. atoms)
   hess_mat = np.ndarray(shape=(d_o_f,d_o_f))
   mw_hess_mat = np.ndarray(shape=(d_o_f,d_o_f))

   # Isotopic substitution will consider 1H/2H, 12C/13C and 16O/17O. More can be
   # added, but it hasn't been necessary so far...
   isolist = iso.split()
   for i in range(0,len(mass_list)):
      for atom in isolist:
         if i == int(atom)-1:
            if mass_list[i] == 1.00783: mass_list[i] = 2.0141 # protium - deuterium
            if mass_list[i] == 12.00000: mass_list[i] = 13.00335 # 12C - 13C
            if mass_list[i] == 15.99491: mass_list[i] = 16.9991 # 16O - 17O
   longline = ""
   for i in range(start_force,len(inlines)):
      longline = longline + inlines[i].rstrip().lstrip()

   forces = longline.split("NImag")[1].split('\\')[2].split(',')
   n = 0; l = -1
   for m in range (0, d_o_f):
      for n in range(0,m+1):
         l = l + 1
         Hmn = forces[l]
         sqrt_Mmn = (mass_list[m/3] * mass_list[n/3]) ** 0.5
         hess_mat[n,m] = float(Hmn)
         mw_hess_mat[m,n] = ((float(Hmn) / sqrt_Mmn))
         mw_hess_mat[n,m] = ((float(Hmn) / sqrt_Mmn))

   return mw_hess_mat

if __name__ == "__main__":
   # Takes arguments: g09_output_files and optional temperature and vibrational scaling factor
   files = []
   freq_scale_factor = 1.0

   if len(sys.argv) > 1:
      for i in range(1,len(sys.argv)):
         if sys.argv[i] == "-s": freq_scale_factor = float(sys.argv[i+1])
         else:
            if len(sys.argv[i].split(".")) > 1:
               if sys.argv[i].split(".")[1] == "out" or sys.argv[i].split(".")[1] == "log": files.append(sys.argv[i])

      print("   Frequency scale factor =", freq_scale_factor)

   else:
      print("\nWrong number of arguments used. Correct format: Hess_to_Freq.py (-s scalefactor) g09_output_file(s)\n")
      sys.exit()

   for file in files:
      mw_hessmat = read_hess(file, 0)
      eigs = np.linalg.eigvalsh(mw_hessmat)
      for eig in eigs:
         if eig > 0: print(eig ** 0.5)
         else: print(-1 * -eig ** 0.5)
<|endoftext|>


<commit_before>import asyncio
import shutil
from contextlib import suppress
from functools import cached_property
from pathlib import Path

import aiopath
from datafiles import datafile, field, frozen
from furl import furl
from sanic import Request
from sanic.log import logger

from .. import settings, utils
from ..types import Dimensions
from .overlay import Overlay
from .text import Text

@datafile("../../templates/{self.id}/config.yml", defaults=True)
class Template:

    id: str
    name: str = ""
    source: str | None = None

    text: list[Text] = field(
        default_factory=lambda: [Text(), Text(anchor_x=0.0, anchor_y=0.8)]
    )

    example: list[str] = field(default_factory=lambda: ["Top Line", "Bottom Line"])

    overlay: list[Overlay] = field(default_factory=lambda: [Overlay()])

    def __str__(self):
        return str(self.directory)

    def __lt__(self, other):
        return self.id < other.id

    def __hash__(self):
        return hash(self.id)

    @cached_property
    def valid(self) -> bool:
        if not settings.DEPLOYED:
            self._update_example()
            self.datafile.save()
        return (
            not self.id.startswith("_")
            and self.image.suffix != settings.PLACEHOLDER_SUFFIX
        )

    def _update_example(self):
        for line in self.example:
            if line and not line.isupper():
                return
        self.example = [line.lower() for line in self.example]

    @cached_property
    def styles(self):
        styles = []
        for path in self.directory.iterdir():
            if path.stem[0] in {".", "_"}:
                continue
            if path.stem not in {"config", "default"}:
                styles.append(path.stem)
            elif path.name == "default.gif":
                styles.append("animated")
        if styles or self.overlay != [Overlay()]:
            styles.append("default")
        styles.sort()
        return styles

    @cached_property
    def animated_image(self):
        return "animated" in self.styles

    @cached_property
    def animated_text(self):
        return any(text.animated for text in self.text)

    @cached_property
    def directory(self) -> Path:
        return self.datafile.path.parent

    @cached_property
    def image(self) -> Path:
        return self.get_image()

    def get_image(self, style: str = "default", *, animated=False) -> Path:
        level = 20 if (style != "default" or animated is True) else 10

        style = style or "default"
        logger.log(level, f"Getting background image: {self.id=} {style=} {animated=}")
        if style == "animated":
            style = "default"
            animated = True

        url = ""
        if utils.urls.schema(style):
            url = style
            style = utils.text.fingerprint(url)

        self.directory.mkdir(exist_ok=True)
        paths: list[Path] = []
        for path in self.directory.iterdir():
            if path.stem == style and path.suffix != settings.PLACEHOLDER_SUFFIX:
                paths.append(path)

        if animated:
            for path in paths:
                if path.suffix == ".gif":
                    logger.log(level, f"Matched path by suffix: {path}")
                    return path

            for path in paths:
                if path.stem == style:
                    logger.log(level, f"Matched path by stem: {path}")
                    return path
        else:
            for path in paths:
                if path.suffix != ".gif":
                    logger.log(level, f"Matched path by suffix: {path}")
                    return path

        path = self.directory / "default.gif"
        if path.exists():
            logger.log(level, f"Matched path by default: {path}")
            return path

        if style == "default" and not animated:
            logger.info(f"No default background image for template: {self.id}")
            return self.directory / ("default" + settings.PLACEHOLDER_SUFFIX)

        if animated:
            logger.info(f"Using static image to animate template: {self.id}")
        else:
            logger.warning(
                f"Style {url or style!r} not available for template: {self.id}"
            )

        return self.get_image()

    def jsonify(self, request: Request) -> dict:
        return {
            "id": self.id,
            "name": self.name,
            "lines": len(self.text),
            "overlays": len(self.overlay) if self.styles else 0,
            "styles": self.styles,
            "blank": request.app.url_for(
                "images.detail_blank",
                template_filename=self.id + "." + settings.DEFAULT_EXTENSION,
                _external=True,
                _scheme=settings.SCHEME,
            ),
            "example": {
                "text": self.example if any(self.example) else [],
                "url": self.build_example_url(request),
            },
            "source": self.source,
            "_self": self.build_self_url(request),
        }

    def build_self_url(self, request: Request) -> str:
        return request.app.url_for(
            "templates.detail",
            id=self.id,
            _external=True,
            _scheme=settings.SCHEME,
        )

    def build_example_url(
        self,
        request: Request,
        *,
        extension: str = "",
        external: bool = True,
    ) -> str:
        extension = extension or self._extension
        kwargs = {
            "template_id": self.id,
            "text_filepath": utils.text.encode(self.example) + "." + extension,
            "_external": external,
        }
        if external:
            kwargs["_scheme"] = settings.SCHEME
        url = request.app.url_for("images.detail_text", **kwargs)
        return utils.urls.clean(url)

    def build_custom_url(
        self,
        request: Request,
        text_lines: list[str],
        *,
        extension: str = "",
        background: str = "",
        style: str = "",
        font: str = "",
    ):
        with suppress(IndexError):
            for index, text in enumerate(self.text):
                text_lines[index] = text.normalize(text_lines[index])
        if not extension and background:
            url = furl(background)
            extension = url.path.segments[-1].split(".")[-1]
        if extension not in settings.ALLOWED_EXTENSIONS:
            extension = self._extension
        if style == "default":
            style = ""
        elif style == "animated":
            extension = "gif"
            style = ""
        url = request.app.url_for(
            "images.detail_text",
            template_id="custom" if self.id == "_custom" else self.id,
            text_filepath=utils.text.encode(text_lines) + "." + extension,
            _external=True,
            _scheme=settings.SCHEME,
            **utils.urls.params(background=background, style=style, font=font),
        )
        return utils.urls.clean(url)

    @property
    def _extension(self) -> str:
        return "gif" if self.animated_image else settings.DEFAULT_EXTENSION

    def build_path(
        self,
        text_lines: list[str],
        font_name: str,
        style: str,
        size: Dimensions,
        watermark: str,
        extension: str,
        frames: int = 0,
    ) -> Path:
        slug = utils.text.encode(text_lines)
        variant = str(self.text) + font_name + style + str(size) + watermark
        if frames:
            variant += str(frames)
        fingerprint = utils.text.fingerprint(variant, prefix="")
        filename = f"{slug}.{fingerprint}.{extension}"
        return Path(self.id) / filename

    @classmethod
    async def create(cls, url: str, *, force=False) -> "Template":
        try:
            parsed = furl(url)
        except ValueError as e:
            logger.error(e)
            return cls.objects.get("_error")

        if parsed.netloc and "memegen.link" in parsed.netloc:
            logger.info(f"Handling template URL: {url}")
            if len(parsed.path.segments) > 1:
                id = Path(parsed.path.segments[1]).stem
                if id != "custom":
                    return cls.objects.get_or_none(id) or cls.objects.get("_error")
                background = parsed.args.get("background")
                if not background:
                    return cls.objects.get("_error")
                url = background
                try:
                    parsed = furl(url)
                except ValueError as e:
                    logger.error(e)
                    return cls.objects.get("_error")

        id = utils.text.fingerprint(url)
        template = cls.objects.get_or_create(id, url)

        suffix = Path(str(parsed.path)).suffix
        if not suffix or len(suffix) > 10:
            logger.warning(f"Unable to determine image extension: {url}")
            suffix = settings.PLACEHOLDER_SUFFIX

        filename = "default" + suffix
        path = aiopath.AsyncPath(template.directory) / filename

        if await path.exists() and not settings.DEBUG and not force:
            logger.info(f"Found background {url} at {path}")
            return template

        logger.info(f"Saving background {url} to {path}")
        if not await utils.http.download(url, path):
            return template

        try:
            await asyncio.to_thread(utils.images.load, Path(path))
        except utils.images.EXCEPTIONS as e:
            logger.error(e)
            await path.unlink(missing_ok=True)

        return template

    async def check(self, style: str, *, animated=False, force=False) -> bool:
        if style in {None, "", "default"}:
            return True

        if style in self.styles:
            return True

        if not utils.urls.schema(style):
            logger.error(f"Invalid style for {self.id} template: {style}")
            return False

        image = self.get_image(animated=animated)
        filename = utils.text.fingerprint(f"{style}{self.overlay}", suffix=image.suffix)
        path = aiopath.AsyncPath(self.directory) / filename
        if await path.exists() and not settings.DEBUG and not force:
            logger.info(f"Found overlay {style} at {path}")
            return True

        urls = style.split(",")
        logger.info(f"Embedding {len(urls)} overlay image(s) onto {path}")
        await asyncio.to_thread(shutil.copy, image, path)

        embedded = 0
        for index, url in enumerate(urls):
            success = await self._embed(index, url, path, force)
            if success:
                embedded += 1

        if len(urls) == 1 and not embedded:
            await path.unlink()

        return embedded == len(urls)

    async def _embed(
        self, index: int, url: str, background: aiopath.AsyncPath, force: bool
    ) -> bool:
        if url.strip() in {"", "default"}:
            return True

        suffix = Path(str(furl(url).path)).suffix
        if not suffix:
            logger.warning(f"Unable to determine image extension: {url}")
            suffix = ".png"

        filename = utils.text.fingerprint(url, prefix="_embed-", suffix=suffix)
        foreground = aiopath.AsyncPath(self.directory) / filename

        if await foreground.exists() and not settings.DEBUG and not force:
            logger.info(f"Found overlay {url} at {foreground}")
        else:
            logger.info(f"Saving overlay {url} to {foreground}")
            await utils.http.download(url, foreground)

        try:
            # TODO: Can 'merge' and 'embed' be combined?
            if background.suffix == ".gif":
                await asyncio.to_thread(
                    utils.images.merge, self, index, Path(foreground), Path(background)
                )
            else:
                await asyncio.to_thread(
                    utils.images.embed, self, index, Path(foreground), Path(background)
                )
        except utils.images.EXCEPTIONS as e:
            logger.error(e)
            await foreground.unlink(missing_ok=True)

        return await foreground.exists()

    def animate(self, start: str = "0.2,0.6", stop: str = "1.0,1.0"):
        try:
            starts = [float(value) for value in start.split(",") if value]
            stops = [float(value) for value in stop.split(",") if value]
        except ValueError:
            logger.error(f"Invalid text animation: {start=} {stop=}")
            return

        with frozen():
            with suppress(IndexError):
                for index, value in enumerate(starts):
                    self.text[index].start = value
            with suppress(IndexError):
                for index, value in enumerate(stops):
                    self.text[index].stop = value

        if starts or stops:
            logger.info(f"Updated {self} with: {starts=} {stops=}")

    def customize(self, scale: float | None):
        with frozen(self):
            if scale is not None:
                self.overlay[0].scale = scale

    def clean(self):
        for path in self.directory.iterdir():
            if path.stem not in {"config", "default"}:
                path.unlink()

    def delete(self):
        if self.directory.exists():
            shutil.rmtree(self.directory)

    def matches(self, query: str) -> bool:
        example = " ".join(line.lower() for line in self.example)
        return any((query in self.id, query in self.name.lower(), query in example))
<commit_msg>Ensure correct overlay is used<commit_after>import asyncio
import shutil
from contextlib import suppress
from functools import cached_property
from pathlib import Path

import aiopath
from datafiles import datafile, field, frozen
from furl import furl
from sanic import Request
from sanic.log import logger

from .. import settings, utils
from ..types import Dimensions
from .overlay import Overlay
from .text import Text

@datafile("../../templates/{self.id}/config.yml", defaults=True)
class Template:

    id: str
    name: str = ""
    source: str | None = None

    text: list[Text] = field(
        default_factory=lambda: [Text(), Text(anchor_x=0.0, anchor_y=0.8)]
    )

    example: list[str] = field(default_factory=lambda: ["Top Line", "Bottom Line"])

    overlay: list[Overlay] = field(default_factory=lambda: [Overlay()])

    def __str__(self):
        return str(self.directory)

    def __lt__(self, other):
        return self.id < other.id

    def __hash__(self):
        return hash(self.id)

    @cached_property
    def valid(self) -> bool:
        if not settings.DEPLOYED:
            self._update_example()
            self.datafile.save()
        return (
            not self.id.startswith("_")
            and self.image.suffix != settings.PLACEHOLDER_SUFFIX
        )

    def _update_example(self):
        for line in self.example:
            if line and not line.isupper():
                return
        self.example = [line.lower() for line in self.example]

    @cached_property
    def styles(self):
        styles = []
        for path in self.directory.iterdir():
            if path.stem[0] in {".", "_"}:
                continue
            if path.stem not in {"config", "default"}:
                styles.append(path.stem)
            elif path.name == "default.gif":
                styles.append("animated")
        if styles or self.overlay != [Overlay()]:
            styles.append("default")
        styles.sort()
        return styles

    @cached_property
    def animated_image(self):
        return "animated" in self.styles

    @cached_property
    def animated_text(self):
        return any(text.animated for text in self.text)

    @cached_property
    def directory(self) -> Path:
        return self.datafile.path.parent

    @cached_property
    def image(self) -> Path:
        return self.get_image()

    def get_image(self, style: str = "default", *, animated=False) -> Path:
        level = 20 if (style != "default" or animated is True) else 10

        style = style or "default"
        logger.log(level, f"Getting background image: {self.id=} {style=} {animated=}")
        if style == "animated":
            style = "default"
            animated = True

        url = ""
        if utils.urls.schema(style):
            url = style
            style = utils.text.fingerprint(f"{url}{self.overlay}")

        self.directory.mkdir(exist_ok=True)
        paths: list[Path] = []
        for path in self.directory.iterdir():
            if path.stem == style and path.suffix != settings.PLACEHOLDER_SUFFIX:
                paths.append(path)

        if animated:
            for path in paths:
                if path.suffix == ".gif":
                    logger.log(level, f"Matched path by suffix: {path}")
                    return path

            for path in paths:
                if path.stem == style:
                    logger.log(level, f"Matched path by stem: {path}")
                    return path
        else:
            for path in paths:
                if path.suffix != ".gif":
                    logger.log(level, f"Matched path by suffix: {path}")
                    return path

        path = self.directory / "default.gif"
        if path.exists():
            logger.log(level, f"Matched path by default: {path}")
            return path

        if style == "default" and not animated:
            logger.info(f"No default background image for template: {self.id}")
            return self.directory / ("default" + settings.PLACEHOLDER_SUFFIX)

        if animated:
            logger.info(f"Using static image to animate template: {self.id}")
        else:
            logger.warning(
                f"Style {url or style!r} not available for template: {self.id}"
            )

        return self.get_image()

    def jsonify(self, request: Request) -> dict:
        return {
            "id": self.id,
            "name": self.name,
            "lines": len(self.text),
            "overlays": len(self.overlay) if self.styles else 0,
            "styles": self.styles,
            "blank": request.app.url_for(
                "images.detail_blank",
                template_filename=self.id + "." + settings.DEFAULT_EXTENSION,
                _external=True,
                _scheme=settings.SCHEME,
            ),
            "example": {
                "text": self.example if any(self.example) else [],
                "url": self.build_example_url(request),
            },
            "source": self.source,
            "_self": self.build_self_url(request),
        }

    def build_self_url(self, request: Request) -> str:
        return request.app.url_for(
            "templates.detail",
            id=self.id,
            _external=True,
            _scheme=settings.SCHEME,
        )

    def build_example_url(
        self,
        request: Request,
        *,
        extension: str = "",
        external: bool = True,
    ) -> str:
        extension = extension or self._extension
        kwargs = {
            "template_id": self.id,
            "text_filepath": utils.text.encode(self.example) + "." + extension,
            "_external": external,
        }
        if external:
            kwargs["_scheme"] = settings.SCHEME
        url = request.app.url_for("images.detail_text", **kwargs)
        return utils.urls.clean(url)

    def build_custom_url(
        self,
        request: Request,
        text_lines: list[str],
        *,
        extension: str = "",
        background: str = "",
        style: str = "",
        font: str = "",
    ):
        with suppress(IndexError):
            for index, text in enumerate(self.text):
                text_lines[index] = text.normalize(text_lines[index])
        if not extension and background:
            url = furl(background)
            extension = url.path.segments[-1].split(".")[-1]
        if extension not in settings.ALLOWED_EXTENSIONS:
            extension = self._extension
        if style == "default":
            style = ""
        elif style == "animated":
            extension = "gif"
            style = ""
        url = request.app.url_for(
            "images.detail_text",
            template_id="custom" if self.id == "_custom" else self.id,
            text_filepath=utils.text.encode(text_lines) + "." + extension,
            _external=True,
            _scheme=settings.SCHEME,
            **utils.urls.params(background=background, style=style, font=font),
        )
        return utils.urls.clean(url)

    @property
    def _extension(self) -> str:
        return "gif" if self.animated_image else settings.DEFAULT_EXTENSION

    def build_path(
        self,
        text_lines: list[str],
        font_name: str,
        style: str,
        size: Dimensions,
        watermark: str,
        extension: str,
        frames: int = 0,
    ) -> Path:
        slug = utils.text.encode(text_lines)
        variant = str(self.text) + font_name + style + str(size) + watermark
        if frames:
            variant += str(frames)
        fingerprint = utils.text.fingerprint(variant, prefix="")
        filename = f"{slug}.{fingerprint}.{extension}"
        return Path(self.id) / filename

    @classmethod
    async def create(cls, url: str, *, force=False) -> "Template":
        try:
            parsed = furl(url)
        except ValueError as e:
            logger.error(e)
            return cls.objects.get("_error")

        if parsed.netloc and "memegen.link" in parsed.netloc:
            logger.info(f"Handling template URL: {url}")
            if len(parsed.path.segments) > 1:
                id = Path(parsed.path.segments[1]).stem
                if id != "custom":
                    return cls.objects.get_or_none(id) or cls.objects.get("_error")
                background = parsed.args.get("background")
                if not background:
                    return cls.objects.get("_error")
                url = background
                try:
                    parsed = furl(url)
                except ValueError as e:
                    logger.error(e)
                    return cls.objects.get("_error")

        id = utils.text.fingerprint(url)
        template = cls.objects.get_or_create(id, url)

        suffix = Path(str(parsed.path)).suffix
        if not suffix or len(suffix) > 10:
            logger.warning(f"Unable to determine image extension: {url}")
            suffix = settings.PLACEHOLDER_SUFFIX

        filename = "default" + suffix
        path = aiopath.AsyncPath(template.directory) / filename

        if await path.exists() and not settings.DEBUG and not force:
            logger.info(f"Found background {url} at {path}")
            return template

        logger.info(f"Saving background {url} to {path}")
        if not await utils.http.download(url, path):
            return template

        try:
            await asyncio.to_thread(utils.images.load, Path(path))
        except utils.images.EXCEPTIONS as e:
            logger.error(e)
            await path.unlink(missing_ok=True)

        return template

    async def check(self, style: str, *, animated=False, force=False) -> bool:
        if style in {None, "", "default"}:
            return True

        if style in self.styles:
            return True

        if not utils.urls.schema(style):
            logger.error(f"Invalid style for {self.id} template: {style}")
            return False

        image = self.get_image(animated=animated)
        filename = utils.text.fingerprint(f"{style}{self.overlay}", suffix=image.suffix)
        path = aiopath.AsyncPath(self.directory) / filename
        if await path.exists() and not settings.DEBUG and not force:
            logger.info(f"Found overlay {style} at {path}")
            return True

        urls = style.split(",")
        logger.info(f"Embedding {len(urls)} overlay image(s) onto {path}")
        await asyncio.to_thread(shutil.copy, image, path)

        embedded = 0
        for index, url in enumerate(urls):
            success = await self._embed(index, url, path, force)
            if success:
                embedded += 1

        if len(urls) == 1 and not embedded:
            await path.unlink()

        return embedded == len(urls)

    async def _embed(
        self, index: int, url: str, background: aiopath.AsyncPath, force: bool
    ) -> bool:
        if url.strip() in {"", "default"}:
            return True

        suffix = Path(str(furl(url).path)).suffix
        if not suffix:
            logger.warning(f"Unable to determine image extension: {url}")
            suffix = ".png"

        filename = utils.text.fingerprint(url, prefix="_embed-", suffix=suffix)
        foreground = aiopath.AsyncPath(self.directory) / filename

        if await foreground.exists() and not settings.DEBUG and not force:
            logger.info(f"Found overlay {url} at {foreground}")
        else:
            logger.info(f"Saving overlay {url} to {foreground}")
            await utils.http.download(url, foreground)

        try:
            # TODO: Can 'merge' and 'embed' be combined?
            if background.suffix == ".gif":
                await asyncio.to_thread(
                    utils.images.merge, self, index, Path(foreground), Path(background)
                )
            else:
                await asyncio.to_thread(
                    utils.images.embed, self, index, Path(foreground), Path(background)
                )
        except utils.images.EXCEPTIONS as e:
            logger.error(e)
            await foreground.unlink(missing_ok=True)

        return await foreground.exists()

    def animate(self, start: str = "0.2,0.6", stop: str = "1.0,1.0"):
        try:
            starts = [float(value) for value in start.split(",") if value]
            stops = [float(value) for value in stop.split(",") if value]
        except ValueError:
            logger.error(f"Invalid text animation: {start=} {stop=}")
            return

        with frozen():
            with suppress(IndexError):
                for index, value in enumerate(starts):
                    self.text[index].start = value
            with suppress(IndexError):
                for index, value in enumerate(stops):
                    self.text[index].stop = value

        if starts or stops:
            logger.info(f"Updated {self} with: {starts=} {stops=}")

    def customize(self, scale: float | None):
        with frozen(self):
            if scale is not None:
                self.overlay[0].scale = scale

    def clean(self):
        for path in self.directory.iterdir():
            if path.stem not in {"config", "default"}:
                path.unlink()

    def delete(self):
        if self.directory.exists():
            shutil.rmtree(self.directory)

    def matches(self, query: str) -> bool:
        example = " ".join(line.lower() for line in self.example)
        return any((query in self.id, query in self.name.lower(), query in example))
<|endoftext|>


<commit_before>#!/usr/bin/env python
#
# (C) 2006 British Broadcasting Corporation and Kamaelia Contributors(1)
#     All Rights Reserved.
#
# You may only modify and redistribute this under the terms of any of the
# following licenses(2): Mozilla Public License, V1.1, GNU General
# Public License, V2.0, GNU Lesser General Public License, V2.1
#
# (1) Kamaelia Contributors are listed in the AUTHORS file and at
#     http://kamaelia.sourceforge.net/AUTHORS - please extend this file,
#     not this notice.
# (2) Reproduced in the COPYING file, and at:
#     http://kamaelia.sourceforge.net/COPYING
# Under section 3.5 of the MPL, we are using this text since we deem the MPL
# notice inappropriate for this file. As per MPL/GPL/LGPL removal of this
# notice is prohibited.
#
# Please contact us via: kamaelia-list-owner@lists.sourceforge.net
# to discuss alternative licensing.
# -------------------------------------------------------------------------
# Licensed to the BBC under a Contributor Agreement: RJL

"""\
===================
Icecast/SHOUTcast MP3 streaming client
===================

This component uses HTTP to stream MP3 audio from a SHOUTcast/Icecast
server.

Example Usage
-------------
IcecastClient fetches the combined audio and metadata stream from the
HTTP server hosting the stream. IcecastDemux separates the audio data
from the metadata in stream and IcecastStreamWriter writes the audio
data to disk (discarding metadata).

pipeline(
    IcecastClient("http://64.236.34.97:80/stream/1049"),
    IcecastDemux(),
    IcecastStreamWriter("stream.mp3"),
).run()

How does it work?
-----------------
The SHOUTcast/Icecast protocol is virtually identical to HTTP. As such,
IcecastClient subclasses SingleShotHTTPClient modifying the request
slightly to ask for stream metadata(e.g. track name) to be included
(by adding the icy-metadata header).
It is otherwise identical to its parent class.
"""

import string, time

from Axon.Component import component
from Axon.Ipc import producerFinished, shutdown

from Kamaelia.Protocol.HTTP.HTTPParser import *
from Kamaelia.Protocol.HTTP.HTTPClient import *
from Kamaelia.Util.PureTransformer import PureTransformer
from Kamaelia.IPC.BaseIPC import IPC

def intval(mystring):
    try:
        retval = int(mystring)
    except ValueError:
        retval = None
    except TypeError:
        retval = None
    return retval

class IceIPCHeader(IPC):
    "Icecast header - content type: %(contenttype)s"
    Parameters = ["contenttype"]

class IceIPCMetadata(IPC):
    "Icecast stream metadata"
    Parameters = ["metadata"]

class IceIPCDataChunk(IPC):
    "An audio/video stream data chunk (typically MP3)"
    Parameters = ["data"]

class IceIPCDisconnected(IPC):
    "Icecast stream disconnected"
    Parameters = []

class IcecastDemux(component):
    """Splits a raw Icecast stream into A/V data and metadata"""

    def dictizeMetadata(self, metadata):
        "Convert metadata that was embedded in the stream into a dictionary."

        #format of metadata:
        #"StreamUrl='www.example.com';\nStreamTitle='singer, title';"
        lines = metadata.split(";")
        metadict = {}
        for line in lines:
            splitline = line.split("=", 1)
            if len(splitline) > 1:
                key = splitline[0]
                val = splitline[1]
                if key[:1] == "\n":
                    key = key[1:]
                if val[0:1] == "'" and val[-1:] == "'":
                    val = val[1:-1]
                metadict[key] = val
        return metadict

    def main(self):
        metadatamode = False
        readbuffer = ""
        while 1:
            yield 1
            while self.dataReady("inbox"):
                msg = self.recv("inbox")

                if isinstance(msg, ParsedHTTPHeader):
                    # metadata is inserted into the stream data every metadatainterval bytes
                    # if an icy-metaint header was in the response
                    metadatainterval = intval(msg.header["headers"].get("icy-metaint", 0))
                    if metadatainterval == None:
                        metadatainterval = 0
                    bytesUntilMetadata = metadatainterval
                    self.send(IceIPCHeader(contenttype=msg.header["headers"].get("content-type")), "outbox")

                    print "Metadata interval is " + str(metadatainterval)

                elif isinstance(msg, ParsedHTTPBodyChunk):
                    readbuffer += msg.bodychunk # NOTE: this is inefficient, consider using a character queue instead

                elif isinstance(msg, ParsedHTTPEnd):
                    self.send(IceIPCDisconnected(), "outbox")

            while len(readbuffer) > 0:
                if metadatainterval == 0: # if the stream does not include metadata
                    # send all the data we have on
                    self.send(IceIPCDataChunk(data=readbuffer), "outbox")
                    readbuffer = ""
                else:
                    chunkdata = readbuffer[0:bytesUntilMetadata]
                    if len(chunkdata) > 0:
                        self.send(IceIPCDataChunk(data=chunkdata), "outbox")

                    readbuffer = readbuffer[bytesUntilMetadata:]
                    bytesUntilMetadata -= len(chunkdata)
                    if len(readbuffer) > 0: #we must have some metadata (perhaps only partially complete) at the start
                        metadatalength = ord(readbuffer[0]) * 16 # they encode it as bytes / 16
                        if len(readbuffer) >= metadatalength + 1: # +1 for the length byte we just read. if we have all the metadata chunk
                            metadictionary = self.dictizeMetadata(readbuffer[1:metadatalength + 1])
                            self.send(IceIPCMetadata(metadata=metadictionary), "outbox")

                            bytesUntilMetadata = metadatainterval
                            readbuffer = readbuffer[metadatalength + 1:]
                        else:
                            break #we need more data before we can do anything

            while self.dataReady("control"):
                msg = self.recv("control")
                if isinstance(msg, producerFinished) or isinstance(msg, shutdown):
                    self.send(producerFinished(self), "signal")
                    return

            self.pause()

class IcecastClient(SingleShotHTTPClient):
    """\
    IcecastClient(starturl) -> Icecast/SHOUTcast MP3 streaming component

    Arguments:
    - starturl    -- the URL of the stream
    """

    def formRequest(self, url):
        """Overrides the standard HTTP request with an Icecast/SHOUTcast variant
        which includes the icy-metadata header required to get metadata with the
        stream"""

        self.send("IcecastClient.formRequest()", "debug")

        splituri = splitUri(url)

        host = splituri["uri-server"]
        if splituri.has_key("uri-port"):
            host += ":" + splituri["uri-port"]

        splituri["request"] =  "GET " + splituri["raw-uri"] + " HTTP/1.1\r\n"
        splituri["request"] += "Host: " + host + "\r\n"
        splituri["request"] += "User-agent: Kamaelia Icecast Client 0.3 (RJL)\r\n"
        splituri["request"] += "Connection: Keep-Alive\r\n"
        splituri["request"] += "icy-metadata: 1\r\n"
        splituri["request"] += "\r\n"
        return splituri

    def main(self):
        while 1: # keep reconnecting
            self.requestqueue.append(HTTPRequest(self.formRequest(self.starturl), 0))
            while self.mainBody():
                yield 1

def IcecastStreamRemoveMetadata():
    def extraDataChunks(msg):
        if isinstance(msg, IceIPCDataChunk):
            return msg.data
        else:
            return None

    return PureTransformer(extraDataChunks)

class IcecastStreamWriter(component):
    Inboxes = {
        "inbox"   : "Icecast stream",
        "control" : "UNUSED"
    }
    Outboxes = {
        "outbox"  : "UNUSED",
        "signal"  : "UNUSED"
    }

    def __init__(self, filename):
        super(IcecastStreamWriter, self).__init__()
        self.filename = filename
    def main(self):
        f = open(self.filename, "wb")
        while 1:
            yield 1
            while self.dataReady("inbox"):
                msg = self.recv("inbox")
                if isinstance(msg, IceIPCDataChunk):
                    f.write(msg.data)

            self.pause()

__kamaelia_components__  = ( IcecastDemux, IcecastClient, IcecastStreamWriter )
__kamaelia_prefabs__  = ( IcecastStreamRemoveMetadata, )

if __name__ == '__main__':
    from Kamaelia.Chassis.Pipeline import pipeline

    # Save a SHOUTcast/Icecast stream to disk
    # (you can then use an MP3 player program to listen to it while it downloads).
    streamurl = raw_input("Stream URL: ") # e.g. "http://a.stream.url.example.com:1234/"
    pipeline(
        IcecastClient(streamurl),
        IcecastDemux(),
        IcecastStreamWriter("stream.mp3"),
    ).run()
<commit_msg>Fixed broken import.<commit_after>#!/usr/bin/env python
#
# (C) 2006 British Broadcasting Corporation and Kamaelia Contributors(1)
#     All Rights Reserved.
#
# You may only modify and redistribute this under the terms of any of the
# following licenses(2): Mozilla Public License, V1.1, GNU General
# Public License, V2.0, GNU Lesser General Public License, V2.1
#
# (1) Kamaelia Contributors are listed in the AUTHORS file and at
#     http://kamaelia.sourceforge.net/AUTHORS - please extend this file,
#     not this notice.
# (2) Reproduced in the COPYING file, and at:
#     http://kamaelia.sourceforge.net/COPYING
# Under section 3.5 of the MPL, we are using this text since we deem the MPL
# notice inappropriate for this file. As per MPL/GPL/LGPL removal of this
# notice is prohibited.
#
# Please contact us via: kamaelia-list-owner@lists.sourceforge.net
# to discuss alternative licensing.
# -------------------------------------------------------------------------
# Licensed to the BBC under a Contributor Agreement: RJL

"""\
===================
Icecast/SHOUTcast MP3 streaming client
===================

This component uses HTTP to stream MP3 audio from a SHOUTcast/Icecast
server.

Example Usage
-------------
IcecastClient fetches the combined audio and metadata stream from the
HTTP server hosting the stream. IcecastDemux separates the audio data
from the metadata in stream and IcecastStreamWriter writes the audio
data to disk (discarding metadata).

pipeline(
    IcecastClient("http://64.236.34.97:80/stream/1049"),
    IcecastDemux(),
    IcecastStreamWriter("stream.mp3"),
).run()

How does it work?
-----------------
The SHOUTcast/Icecast protocol is virtually identical to HTTP. As such,
IcecastClient subclasses SingleShotHTTPClient modifying the request
slightly to ask for stream metadata(e.g. track name) to be included
(by adding the icy-metadata header).
It is otherwise identical to its parent class.
"""

import string, time

from Axon.Component import component
from Axon.Ipc import producerFinished, shutdown

from Kamaelia.Protocol.HTTP.HTTPParser import *
from Kamaelia.Protocol.HTTP.HTTPClient import *
from Kamaelia.Util.PureTransformer import PureTransformer
from Kamaelia.BaseIPC import IPC

def intval(mystring):
    try:
        retval = int(mystring)
    except ValueError:
        retval = None
    except TypeError:
        retval = None
    return retval

class IceIPCHeader(IPC):
    "Icecast header - content type: %(contenttype)s"
    Parameters = ["contenttype"]

class IceIPCMetadata(IPC):
    "Icecast stream metadata"
    Parameters = ["metadata"]

class IceIPCDataChunk(IPC):
    "An audio/video stream data chunk (typically MP3)"
    Parameters = ["data"]

class IceIPCDisconnected(IPC):
    "Icecast stream disconnected"
    Parameters = []

class IcecastDemux(component):
    """Splits a raw Icecast stream into A/V data and metadata"""

    def dictizeMetadata(self, metadata):
        "Convert metadata that was embedded in the stream into a dictionary."

        #format of metadata:
        #"StreamUrl='www.example.com';\nStreamTitle='singer, title';"
        lines = metadata.split(";")
        metadict = {}
        for line in lines:
            splitline = line.split("=", 1)
            if len(splitline) > 1:
                key = splitline[0]
                val = splitline[1]
                if key[:1] == "\n":
                    key = key[1:]
                if val[0:1] == "'" and val[-1:] == "'":
                    val = val[1:-1]
                metadict[key] = val
        return metadict

    def main(self):
        metadatamode = False
        readbuffer = ""
        while 1:
            yield 1
            while self.dataReady("inbox"):
                msg = self.recv("inbox")

                if isinstance(msg, ParsedHTTPHeader):
                    # metadata is inserted into the stream data every metadatainterval bytes
                    # if an icy-metaint header was in the response
                    metadatainterval = intval(msg.header["headers"].get("icy-metaint", 0))
                    if metadatainterval == None:
                        metadatainterval = 0
                    bytesUntilMetadata = metadatainterval
                    self.send(IceIPCHeader(contenttype=msg.header["headers"].get("content-type")), "outbox")

                    print "Metadata interval is " + str(metadatainterval)

                elif isinstance(msg, ParsedHTTPBodyChunk):
                    readbuffer += msg.bodychunk # NOTE: this is inefficient, consider using a character queue instead

                elif isinstance(msg, ParsedHTTPEnd):
                    self.send(IceIPCDisconnected(), "outbox")

            while len(readbuffer) > 0:
                if metadatainterval == 0: # if the stream does not include metadata
                    # send all the data we have on
                    self.send(IceIPCDataChunk(data=readbuffer), "outbox")
                    readbuffer = ""
                else:
                    chunkdata = readbuffer[0:bytesUntilMetadata]
                    if len(chunkdata) > 0:
                        self.send(IceIPCDataChunk(data=chunkdata), "outbox")

                    readbuffer = readbuffer[bytesUntilMetadata:]
                    bytesUntilMetadata -= len(chunkdata)
                    if len(readbuffer) > 0: #we must have some metadata (perhaps only partially complete) at the start
                        metadatalength = ord(readbuffer[0]) * 16 # they encode it as bytes / 16
                        if len(readbuffer) >= metadatalength + 1: # +1 for the length byte we just read. if we have all the metadata chunk
                            metadictionary = self.dictizeMetadata(readbuffer[1:metadatalength + 1])
                            self.send(IceIPCMetadata(metadata=metadictionary), "outbox")

                            bytesUntilMetadata = metadatainterval
                            readbuffer = readbuffer[metadatalength + 1:]
                        else:
                            break #we need more data before we can do anything

            while self.dataReady("control"):
                msg = self.recv("control")
                if isinstance(msg, producerFinished) or isinstance(msg, shutdown):
                    self.send(producerFinished(self), "signal")
                    return

            self.pause()

class IcecastClient(SingleShotHTTPClient):
    """\
    IcecastClient(starturl) -> Icecast/SHOUTcast MP3 streaming component

    Arguments:
    - starturl    -- the URL of the stream
    """

    def formRequest(self, url):
        """Overrides the standard HTTP request with an Icecast/SHOUTcast variant
        which includes the icy-metadata header required to get metadata with the
        stream"""

        self.send("IcecastClient.formRequest()", "debug")

        splituri = splitUri(url)

        host = splituri["uri-server"]
        if splituri.has_key("uri-port"):
            host += ":" + splituri["uri-port"]

        splituri["request"] =  "GET " + splituri["raw-uri"] + " HTTP/1.1\r\n"
        splituri["request"] += "Host: " + host + "\r\n"
        splituri["request"] += "User-agent: Kamaelia Icecast Client 0.3 (RJL)\r\n"
        splituri["request"] += "Connection: Keep-Alive\r\n"
        splituri["request"] += "icy-metadata: 1\r\n"
        splituri["request"] += "\r\n"
        return splituri

    def main(self):
        while 1: # keep reconnecting
            self.requestqueue.append(HTTPRequest(self.formRequest(self.starturl), 0))
            while self.mainBody():
                yield 1

def IcecastStreamRemoveMetadata():
    def extraDataChunks(msg):
        if isinstance(msg, IceIPCDataChunk):
            return msg.data
        else:
            return None

    return PureTransformer(extraDataChunks)

class IcecastStreamWriter(component):
    Inboxes = {
        "inbox"   : "Icecast stream",
        "control" : "UNUSED"
    }
    Outboxes = {
        "outbox"  : "UNUSED",
        "signal"  : "UNUSED"
    }

    def __init__(self, filename):
        super(IcecastStreamWriter, self).__init__()
        self.filename = filename
    def main(self):
        f = open(self.filename, "wb")
        while 1:
            yield 1
            while self.dataReady("inbox"):
                msg = self.recv("inbox")
                if isinstance(msg, IceIPCDataChunk):
                    f.write(msg.data)

            self.pause()

__kamaelia_components__  = ( IcecastDemux, IcecastClient, IcecastStreamWriter )
__kamaelia_prefabs__  = ( IcecastStreamRemoveMetadata, )

if __name__ == '__main__':
    from Kamaelia.Chassis.Pipeline import pipeline

    # Save a SHOUTcast/Icecast stream to disk
    # (you can then use an MP3 player program to listen to it while it downloads).
    streamurl = raw_input("Stream URL: ") # e.g. "http://a.stream.url.example.com:1234/"
    pipeline(
        IcecastClient(streamurl),
        IcecastDemux(),
        IcecastStreamWriter("stream.mp3"),
    ).run()
<|endoftext|>