<commit_before>"""This defines classes for working with geometry file formats"""

########################################################################
#                                                                      #
# This script was written by Thomas Heavey in 2017.                    #
#        theavey@bu.edu     thomasjheavey@gmail.com                    #
#                                                                      #
# Copyright 2017 Thomas J. Heavey IV                                   #
#                                                                      #
# Licensed under the Apache License, Version 2.0 (the "License");      #
# you may not use this file except in compliance with the License.     #
# You may obtain a copy of the License at                              #
#                                                                      #
#    http://www.apache.org/licenses/LICENSE-2.0                        #
#                                                                      #
# Unless required by applicable law or agreed to in writing, software  #
# distributed under the License is distributed on an "AS IS" BASIS,    #
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or      #
# implied.                                                             #
# See the License for the specific language governing permissions and  #
# limitations under the License.                                       #
#                                                                      #
########################################################################

from __future__ import absolute_import

import re
import numpy as np
from numpy.linalg import norm
from .exceptions import UnknownEnergyError, InputError
# TODO add tests for these

__all__ = ['rotation_matrix', 'Vector', 'XYZ', 'COM']

def rotation_matrix(axis, theta):
    """
    Return the rotation matrix associated with counterclockwise rotation about
    the given axis by theta radians.

    copied from
    https://stackoverflow.com/questions/6802577/python-rotation-of-3d-vector
    """
    axis = np.asarray(axis)
    axis = axis/np.sqrt(np.dot(axis, axis))
    a = np.cos(theta/2.0)
    b, c, d = -axis * np.sin(theta/2.0)
    aa, bb, cc, dd = a*a, b*b, c*c, d*d
    bc, ad, ac, ab, bd, cd = b*c, a*d, a*c, a*b, b*d, c*d
    return np.array([[aa+bb-cc-dd, 2*(bc+ad), 2*(bd-ac)],
                     [2*(bc-ad), aa+cc-bb-dd, 2*(cd+ab)],
                     [2*(bd+ac), 2*(cd-ab), aa+dd-bb-cc]])

class Vector(np.ndarray):

    def __new__(cls, *xyz):
        if len(xyz) != 3:
            try:
                xyz = xyz[0]
            except IndexError:
                raise InputError(xyz, '3 values are required to make a vector')
        if len(xyz) != 3:
            raise InputError(xyz, 'Length of vector must be 3.')
        obj = super(Vector, cls).__new__(cls, shape=(3,),
                                         buffer=np.array(xyz, dtype=float))
        return obj

    def cross(self, vec):
        return np.cross(self, vec)

    def diff_angle(self, vec):
        return np.arccos(self.dot(vec) /
                         (norm(self) * norm(vec)))

    def rotate(self, axis, angle):
        r_mat = rotation_matrix(axis, angle)
        return np.dot(r_mat, self)

    @property
    def x(self):
        return self[0]

    @property
    def y(self):
        return self[1]

    @property
    def z(self):
        return self[2]

    @property
    def mag(self):
        return norm(self)

class XYZ(object):
    def __init__(self, f_name):
        self.file = f_name
        with open(f_name, 'r') as f_file:
            f_lines = f_file.readlines()
        self._header = f_lines[0:2]
        if 'Energy' in self._header[1]:
            energy_match = re.search('(?:Energy:\s+)(-\d+\.\d+)',
                                     self._header[1])
            self._energy = float(energy_match.group(1))
        else:
            self._energy = None
        self._original_energy = self._energy
        data = [line.split() for line in f_lines[2:]]
        if not data[-1]:
            data = data[:-1]  # Ignore blank last line
        try:
            self.atoms = [atom[0] for atom in data]
        except IndexError:
            raise ValueError('invalid line in xyz file: {}'.format(atom))
        self.coords = [Vector([float(coord) for coord in atom[1:4]]) for
                       atom in data]

    def center_on(self, index):
        center = self.coords[index]
        self.coords = [coord - center for coord in self.coords]

    def rotate_to_x_axis_on(self, index):
        vec_x = Vector(1, 0, 0)
        angle = self.coords[index].diff_angle(vec_x)
        axis = self.coords[index].cross(vec_x)
        self.coords = [coord.rotate(angle, axis) for coord in self.coords]

    def center_and_rotate_on(self, index1, index2):
        self.center_on(index1)
        self.rotate_to_x_axis_on(index2)

    def __str__(self):
        f_string = ('   {0: <10s} {1.x: > 10.5f} {1.y: > 10.5f} '
                    '{1.z: > 10.5f}\n')
        output_list = list(self._header)
        output_list += [f_string.format(self.atoms[i], self.coords[i]) for i
                        in range(len(self.atoms))]
        return ''.join(output_list)

    @property
    def n_atoms(self):
        _n_atoms = len(self.atoms)
        _n_coords = len(self.coords)
        if _n_atoms != _n_coords:
            print('!!n atoms != n coords!! ({} != {})'.format(_n_atoms,
                                                              _n_coords))
        else:
            return _n_atoms

    @property
    def energy(self):
        if self._energy is None:
            raise UnknownEnergyError()
        return self._energy

    @property
    def original_energy(self):
        return self._original_energy

    def replace_coords(self, arg):
        if isinstance(arg, str):
            self.coords = XYZ(arg).coords.copy()
        else:
            self.coords = arg.coords.copy()
        self._energy = None  # Moved atoms, don't know energy

    def move_subset(self, movement, indices):
        if not isinstance(movement, Vector):
            movement = Vector(*movement)
        for index in indices:
            self.coords[index] = self.coords[index] + movement
        self._energy = None  # Moved atoms, don't know energy

    def write(self, f_name):
        with open(f_name, 'w') as f_file:
            f_file.write(str(self))

    def average_loc(self, *args):
        if len(args) == 1:  # if an Iterable is passed in
            args = args[0]
        total_vec = Vector(0, 0, 0)
        for i in args:
            total_vec = total_vec + self.coords[i]
        return total_vec / len(args)

class COM(XYZ):
    def __init__(self, f_name):
        self.file = f_name
        self._header = []
        self._title = []
        self._cm = []
        self.atoms = []
        self.coords = []
        self._footer = []
        with open(f_name, 'r') as f_file:
            f_lines = f_file.readlines()
        self._parser(f_lines)

    def _parser(self, lines):
        section = 'header'
        data = []
        for line in lines:
            if section == 'header':
                self._header.append(line)
                if line.strip() == '':
                    section = 'title'
                continue
            elif section == 'title':
                self._title.append(line)
                if line.strip() == '':
                    section = 'charge_mult'
                continue
            elif section == 'charge_mult':
                self._cm.append(line)
                section = 'geom'
                continue
            elif section == 'geom':
                if line.strip() == '':
                    section = 'opt_input'
                    continue
                data.append(line.split())
                continue
            elif section == 'opt_input':
                self._footer.append(line)
                continue
        self.atoms = [atom[0] for atom in data]
        self.coords = [Vector([float(coord) for
                               coord in atom[1:4]]) for atom in data]

    def __str__(self):
        f_string = ('   {0: <10s} {1.x: > 10.5f} {1.y: > 10.5f} '
                    '{1.z: > 10.5f}\n')
        output_list = list(self._header)
        output_list += list(self._title)
        output_list += list(self._cm)
        output_list += [f_string.format(self.atoms[i], self.coords[i]) for i
                        in range(len(self.atoms))]
        output_list += ['\n']
        output_list += list(self._footer)
        return ''.join(output_list)
<commit_msg>Fix atom names if they contain indexes with the atomic symbols<commit_after>"""This defines classes for working with geometry file formats"""

########################################################################
#                                                                      #
# This script was written by Thomas Heavey in 2017.                    #
#        theavey@bu.edu     thomasjheavey@gmail.com                    #
#                                                                      #
# Copyright 2017 Thomas J. Heavey IV                                   #
#                                                                      #
# Licensed under the Apache License, Version 2.0 (the "License");      #
# you may not use this file except in compliance with the License.     #
# You may obtain a copy of the License at                              #
#                                                                      #
#    http://www.apache.org/licenses/LICENSE-2.0                        #
#                                                                      #
# Unless required by applicable law or agreed to in writing, software  #
# distributed under the License is distributed on an "AS IS" BASIS,    #
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or      #
# implied.                                                             #
# See the License for the specific language governing permissions and  #
# limitations under the License.                                       #
#                                                                      #
########################################################################

from __future__ import absolute_import

import re
import numpy as np
from numpy.linalg import norm
from .exceptions import UnknownEnergyError, InputError
# TODO add tests for these

__all__ = ['rotation_matrix', 'Vector', 'XYZ', 'COM']

def rotation_matrix(axis, theta):
    """
    Return the rotation matrix associated with counterclockwise rotation about
    the given axis by theta radians.

    copied from
    https://stackoverflow.com/questions/6802577/python-rotation-of-3d-vector
    """
    axis = np.asarray(axis)
    axis = axis/np.sqrt(np.dot(axis, axis))
    a = np.cos(theta/2.0)
    b, c, d = -axis * np.sin(theta/2.0)
    aa, bb, cc, dd = a*a, b*b, c*c, d*d
    bc, ad, ac, ab, bd, cd = b*c, a*d, a*c, a*b, b*d, c*d
    return np.array([[aa+bb-cc-dd, 2*(bc+ad), 2*(bd-ac)],
                     [2*(bc-ad), aa+cc-bb-dd, 2*(cd+ab)],
                     [2*(bd+ac), 2*(cd-ab), aa+dd-bb-cc]])

class Vector(np.ndarray):

    def __new__(cls, *xyz):
        if len(xyz) != 3:
            try:
                xyz = xyz[0]
            except IndexError:
                raise InputError(xyz, '3 values are required to make a vector')
        if len(xyz) != 3:
            raise InputError(xyz, 'Length of vector must be 3.')
        obj = super(Vector, cls).__new__(cls, shape=(3,),
                                         buffer=np.array(xyz, dtype=float))
        return obj

    def cross(self, vec):
        return np.cross(self, vec)

    def diff_angle(self, vec):
        return np.arccos(self.dot(vec) /
                         (norm(self) * norm(vec)))

    def rotate(self, axis, angle):
        r_mat = rotation_matrix(axis, angle)
        return np.dot(r_mat, self)

    @property
    def x(self):
        return self[0]

    @property
    def y(self):
        return self[1]

    @property
    def z(self):
        return self[2]

    @property
    def mag(self):
        return norm(self)

class XYZ(object):
    def __init__(self, f_name):
        self.file = f_name
        with open(f_name, 'r') as f_file:
            f_lines = f_file.readlines()
        self._header = f_lines[0:2]
        if 'Energy' in self._header[1]:
            energy_match = re.search('(?:Energy:\s+)(-\d+\.\d+)',
                                     self._header[1])
            self._energy = float(energy_match.group(1))
        else:
            self._energy = None
        self._original_energy = self._energy
        data = [line.split() for line in f_lines[2:]]
        if not data[-1]:
            data = data[:-1]  # Ignore blank last line
        try:
            self.atoms = [atom[0] for atom in data]
        except IndexError:
            raise ValueError('invalid line in xyz file: {}'.format(atom))
        if re.search(r'([a-z]|[A-Z])+\d+', self.atoms[0]):
            self._fix_atom_names()
        self.coords = [Vector([float(coord) for coord in atom[1:4]]) for
                       atom in data]

    def _fix_atom_names(self):
        atoms_new = []
        for atom in self.atoms:
            atoms_new.append(re.match(r'([a-z]|[A-Z])+', atom).group(0))
        self.atoms = atoms_new

    def center_on(self, index):
        center = self.coords[index]
        self.coords = [coord - center for coord in self.coords]

    def rotate_to_x_axis_on(self, index):
        vec_x = Vector(1, 0, 0)
        angle = self.coords[index].diff_angle(vec_x)
        axis = self.coords[index].cross(vec_x)
        self.coords = [coord.rotate(angle, axis) for coord in self.coords]

    def center_and_rotate_on(self, index1, index2):
        self.center_on(index1)
        self.rotate_to_x_axis_on(index2)

    def __str__(self):
        f_string = ('   {0: <10s} {1.x: > 10.5f} {1.y: > 10.5f} '
                    '{1.z: > 10.5f}\n')
        output_list = list(self._header)
        output_list += [f_string.format(self.atoms[i], self.coords[i]) for i
                        in range(len(self.atoms))]
        return ''.join(output_list)

    @property
    def n_atoms(self):
        _n_atoms = len(self.atoms)
        _n_coords = len(self.coords)
        if _n_atoms != _n_coords:
            print('!!n atoms != n coords!! ({} != {})'.format(_n_atoms,
                                                              _n_coords))
        else:
            return _n_atoms

    @property
    def energy(self):
        if self._energy is None:
            raise UnknownEnergyError()
        return self._energy

    @property
    def original_energy(self):
        return self._original_energy

    def replace_coords(self, arg):
        if isinstance(arg, str):
            self.coords = XYZ(arg).coords.copy()
        else:
            self.coords = arg.coords.copy()
        self._energy = None  # Moved atoms, don't know energy

    def move_subset(self, movement, indices):
        if not isinstance(movement, Vector):
            movement = Vector(*movement)
        for index in indices:
            self.coords[index] = self.coords[index] + movement
        self._energy = None  # Moved atoms, don't know energy

    def write(self, f_name):
        with open(f_name, 'w') as f_file:
            f_file.write(str(self))

    def average_loc(self, *args):
        if len(args) == 1:  # if an Iterable is passed in
            args = args[0]
        total_vec = Vector(0, 0, 0)
        for i in args:
            total_vec = total_vec + self.coords[i]
        return total_vec / len(args)

class COM(XYZ):
    def __init__(self, f_name):
        self.file = f_name
        self._header = []
        self._title = []
        self._cm = []
        self.atoms = []
        self.coords = []
        self._footer = []
        with open(f_name, 'r') as f_file:
            f_lines = f_file.readlines()
        self._parser(f_lines)

    def _parser(self, lines):
        section = 'header'
        data = []
        for line in lines:
            if section == 'header':
                self._header.append(line)
                if line.strip() == '':
                    section = 'title'
                continue
            elif section == 'title':
                self._title.append(line)
                if line.strip() == '':
                    section = 'charge_mult'
                continue
            elif section == 'charge_mult':
                self._cm.append(line)
                section = 'geom'
                continue
            elif section == 'geom':
                if line.strip() == '':
                    section = 'opt_input'
                    continue
                data.append(line.split())
                continue
            elif section == 'opt_input':
                self._footer.append(line)
                continue
        self.atoms = [atom[0] for atom in data]
        self.coords = [Vector([float(coord) for
                               coord in atom[1:4]]) for atom in data]

    def __str__(self):
        f_string = ('   {0: <10s} {1.x: > 10.5f} {1.y: > 10.5f} '
                    '{1.z: > 10.5f}\n')
        output_list = list(self._header)
        output_list += list(self._title)
        output_list += list(self._cm)
        output_list += [f_string.format(self.atoms[i], self.coords[i]) for i
                        in range(len(self.atoms))]
        output_list += ['\n']
        output_list += list(self._footer)
        return ''.join(output_list)
<|endoftext|>


<commit_before>from .version import __version__
from .network import Network, NodeScanner
from .node import RemoteNode, LocalNode
from .sdo import SdoCommunicationError, SdoAbortedError
from .objectdictionary import import_od, ObjectDictionary, ObjectDictionaryError
from .profiles.p402 import Node402

Node = RemoteNode
<commit_msg>fix broken import<commit_after>from .version import __version__
from .network import Network, NodeScanner
from .node import RemoteNode, LocalNode
from .sdo import SdoCommunicationError, SdoAbortedError
from .objectdictionary import import_od, ObjectDictionary, ObjectDictionaryError
from .profiles.p402 import BaseNode402

Node = RemoteNode
<|endoftext|>


<commit_before>#!/usr/bin/env python
# -*- coding: utf-8 -*-

import logging
import re

from utils import remove_null_strings

logger = logging.getLogger(__name__)

class GameParser():

    # defining time zone information
    TZINFO = {'CET':   3600, 'CEST':    7200,
              'EET':   7200, 'EETDST': 10800,
              'EDT': -14400, 'EST': -18000,
              'CDT': -18000, 'CST': -21600,
              'MDT': -21600, 'MST': -25200,
              'PDT': -25200, 'PST': -28800,
              'BST':   3600}

    # regular expression retrieve attendance figure
    ATTENDANCE_AT_VENUE_REGEX = re.compile("\s(@|at)\s")

    def __init__(self, raw_data, raw_so_data):
        self.raw_data = raw_data
        self.raw_so_data = raw_so_data

    def load_data(self):
        """
        Loads raw data from html and pre-processes it.
        """
        # finding content of html element with *GameInfo*-id
        game_data_str = self.raw_data.xpath(
            "//table[@id='GameInfo']/tr/td/text()")
        game_data_str = [re.sub("\s+", " ", s) for s in game_data_str]
        self.game_data = remove_null_strings(game_data_str)[-5:]
<commit_msg>Extend game parser functions<commit_after>#!/usr/bin/env python
# -*- coding: utf-8 -*-

import logging
import re

from dateutil import parser

from utils import remove_null_strings, retrieve_season

logger = logging.getLogger(__name__)

class GameParser():

    # defining time zone information
    TZINFO = {'CET':   3600, 'CEST': 7200,
              'EET':   7200, 'EETDST': 10800,
              'EDT': -14400, 'EST': -18000,
              'CDT': -18000, 'CST': -21600,
              'MDT': -21600, 'MST': -25200,
              'PDT': -25200, 'PST': -28800,
              'BST':   3600}

    # regular expression retrieve attendance figure
    ATTENDANCE_AT_VENUE_REGEX = re.compile("\s(@|at)\s")

    def __init__(self, game_id, raw_data, raw_so_data):
        self.game_id = game_id
        self.raw_data = raw_data
        self.raw_so_data = raw_so_data

    def create_game(self, teams):
        # loading and pre-processing raw data
        self.load_data()

        # retrieving game date from raw data
        game_date = parser.parse(self.game_data[0]).date()
        # retrieving season for current game date
        season = retrieve_season(game_date)
        # setting up full game id, containing season, game type
        # and partial game id
        full_game_id = "%d%s" % (season, self.game_id)

        attendance, venue = self.retrieve_game_attendance_venue()

    def retrieve_game_attendance_venue(self):
        """
        Retrieves attendance and venue information for current game.
        """
        # retrieving attendance and venue from string,
        # i.e. *Ass./Att. 21,273 @ Centre Bell*
        if any(s in self.game_data[1] for s in ['@', 'at']):
            attendance_venue = self.game_data[1].split(" ", 3)
            # print(attendance_venue)
            # attendance_venue = re.split(
            #     self.ATTENDANCE_AT_VENUE_REGEX, self.game_data[1])
            # print(attendance_venue)
        else:
            attendance_venue = self.game_data[1]

        try:
            attendance = int(attendance_venue[1].replace(",", ""))
        except:
            attendance = 0
        venue = attendance_venue[-1]

        return attendance, venue

    def load_data(self):
        """
        Loads raw data from html and pre-processes it.
        """
        # finding content of html element with *GameInfo*-id
        game_data_str = self.raw_data.xpath(
            "//table[@id='GameInfo']/tr/td/text()")
        game_data_str = [re.sub("\s+", " ", s) for s in game_data_str]
        self.game_data = remove_null_strings(game_data_str)[-5:]
<|endoftext|>


<commit_before>import os
import time
import logging
from ipyparallel import Client

from parsl.executors.base import ParslExecutor
from parsl.executors.errors import *

logger = logging.getLogger(__name__)

class IPyParallelExecutor(ParslExecutor):
    """The IPython Parallel executor.

    This executor allows us to take advantage of multiple processes running locally
    or remotely via  IPythonParallel's pilot execution system.

    .. note::
           Some deficiencies with this executor are:

               1. Ipengine's execute one task at a time. This means one engine per core
                  is necessary to exploit the full parallelism of a node.
               2. No notion of remaining walltime.
               3. Lack of throttling means tasks could be queued up on a worker.

    """

    def compose_launch_cmd(self, filepath, engine_dir):
        """Reads the json contents from filepath and uses that to compose the engine launch command.

        Args:
            filepath: Path to the engine file
            engine_dir : CWD for the engines .

        """
        self.engine_file = os.path.expanduser(filepath)

        engine_json = None
        try:
            with open(self.engine_file, 'r') as f:
                engine_json = f.read()

        except OSError as e:
            logger.error("Could not open engine_json : ", self.engine_file)
            raise e

        return """cd {0}
cat <<EOF > ipengine.json
{1}
EOF

mkdir -p '.ipengine_logs'
ipengine --file=ipengine.json &>> .ipengine_logs/$JOBNAME.log
""".format(engine_dir, engine_json)

    def __init__(self, execution_provider=None,
                 reuse_controller=True,
                 engine_json_file='~/.ipython/profile_default/security/ipcontroller-engine.json',
                 engine_dir='.',
                 controller=None,
                 config=None):
        """Initialize the IPyParallel pool. The initialization takes all relevant parameters via KWargs.

        .. note::

              If initBlocks > 0, and a scalable execution_provider is attached, then the provider
              will be initialized here.

        Args:
             - self

        KWargs:
             - execution_provider (ExecutionProvider object)
             - reuse_controller (Bool) : If True ipp executor will attempt to connect to an available
               controller. Default: True
             - engine_json_file (str): Path to json engine file that will be used to compose ipp launch
               commands at scaling events. Default : '~/.ipython/profile_default/security/ipcontroller-engine.json'
             - engine_dir (str) : Alternative to above, specify the engine_dir
             - config (dict). Default: '.'
        """
        self.controller = controller
        self.engine_file = engine_json_file
        self.client_file = None

        if self.controller:
            # Find the Client json
            self.client_file = self.controller.client_file
            self.engine_file = self.controller.engine_file

            if not os.path.exists(self.client_file):
                logger.debug("Waiting for {0}".format(self.client_file))

            sleep_dur = 20  # 20 seconds
            for i in range(0, int(sleep_dur / 0.2)):
                time.sleep(0.2)
                if os.path.exists(self.client_file):
                    break

            if not os.path.exists(self.client_file):
                raise Exception("Controller client file is missing at {0}".format(self.client_file))

        self.executor = Client(url_file=self.client_file)
        self.config = config
        self.sitename = config['site'] if config else 'Static_IPP'
        # NOTE: Copying the config here only partially fixes the issue. There needs to be
        # multiple controllers launched by the factory, and each must have different jsons.
        # There could be timing issues here,
        # local_engine_json = "{0}.{1}.engine.json".format(self.config["site"], int(time.time()))
        # copyfile(engine_json_file, local_engine_json)
        # if not os.path.exists(self.config["execution"]["script_dir"]):
        #    os.makedirs(self.config["execution"]["script_dir"])

        self.launch_cmd = self.compose_launch_cmd(self.engine_file, engine_dir)
        self.execution_provider = execution_provider
        self.engines = []

        if reuse_controller:
            # Reuse existing controller if one is available
            pass

        if execution_provider:
            self._scaling_enabled = True
            logger.debug("Starting IpyParallelExecutor with provider:%s", execution_provider)
            try:
                for i in range(self.config["execution"]["block"].get("initBlocks", 1)):
                    eng = self.execution_provider.submit(self.launch_cmd, 1)
                    logger.debug("Launched block : {0}:{1}".format(i, eng))
                    if not eng:
                        raise(ScalingFailed(self.execution_provider.sitename,
                                            "Ipp executor failed to scale via execution_provider"))
                    self.engines.extend([eng])

            except Exception as e:
                logger.error("Scaling out failed : %s", e)
                raise e

        else:
            self._scaling_enabled = False
            logger.debug("Starting IpyParallelExecutor with no provider")

        self.lb_view = self.executor.load_balanced_view()
        logger.debug("Starting executor")

    @property
    def scaling_enabled(self):
        return self._scaling_enabled

    def submit(self, *args, **kwargs):
        """Submits work to the thread pool.

        This method is simply pass through and behaves like a submit call as described
        here `Python docs: <https://docs.python.org/3/library/concurrent.futures.html#concurrent.futures.ThreadPoolExecutor>`_

        Returns:
              Future
        """
        # logger.debug("Got args : %s,", args)
        # logger.debug("Got kwargs : %s,", kwargs)
        return self.lb_view.apply_async(*args, **kwargs)

    def scale_out(self, *args, **kwargs):
        """Scales out the number of active workers by 1.

        This method is notImplemented for threads and will raise the error if called.

        """
        if self.execution_provider:
            r = self.execution_provider.submit(self.launch_cmd, *args, **kwargs)
            self.engines.extend([r])
        else:
            logger.error("No execution provider available")
            r = None

        return r

    def scale_in(self, blocks, *args, **kwargs):
        """Scale in the number of active workers by 1.

        This method is notImplemented for threads and will raise the error if called.

        Raises:
             NotImplemented exception
        """
        status = dict(zip(self.engines, self.execution_provider.status(self.engines)))

        # This works for blocks=0
        to_kill = [engine for engine in status if status[engine] == "RUNNING"][:blocks]

        if self.execution_provider:
            r = self.execution_provider.cancel(to_kill, *args, **kwargs)
        else:
            logger.error("No execution provider available")
            r = None

        return r

    def status(self):
        """Returns the status of the executor via probing the execution providers."""
        if self.execution_provider:
            status = self.execution_provider.status(self.engines)

        else:
            status = []

        return status

    def shutdown(self, hub=True, targets='all', block=False):
        """Shutdown the executor, including all workers and controllers.

        The interface documentation for IPP is `here <http://ipyparallel.readthedocs.io/en/latest/api/ipyparallel.html#ipyparallel.Client.shutdown>`_

        Kwargs:
            - hub (Bool): Whether the hub should be shutdown, Default:True,
            - targets (list of ints| 'all'): List of engine id's to kill, Default:'all'
            - block (Bool): To block for confirmations or not

        Raises:
             NotImplementedError
        """
        if self.controller:
            logger.debug("IPP:Shutdown sequence: Attempting controller kill")
            self.controller.close()

        # We do not actually do executor.shutdown because
        # this blocks even when requested to not block, killing the
        # controller is more effective although impolite.
        # x = self.executor.shutdown(targets=targets,
        #                           hub=hub,
        #                           block=block)

        logger.debug("Done with executor shutdown")
        return True

    def __repr__(self):
        return "<IPP Executor for site:{0}>".format(self.sitename)

if __name__ == "__main__":

    pool1_config = {"poolname": "pool1",
                    "queue": "foo"}
<commit_msg>Adding support for specifying ipengine launch in containers<commit_after>import os
import time
import logging
from ipyparallel import Client

from parsl.executors.base import ParslExecutor
from parsl.executors.errors import *

logger = logging.getLogger(__name__)

class IPyParallelExecutor(ParslExecutor):
    """The IPython Parallel executor.

    This executor allows us to take advantage of multiple processes running locally
    or remotely via  IPythonParallel's pilot execution system.

    .. note::
           Some deficiencies with this executor are:

               1. Ipengine's execute one task at a time. This means one engine per core
                  is necessary to exploit the full parallelism of a node.
               2. No notion of remaining walltime.
               3. Lack of throttling means tasks could be queued up on a worker.

    """

    def compose_launch_cmd(self, filepath, engine_dir, container_image):
        """Reads the json contents from filepath and uses that to compose the engine launch command.

        Args:
            filepath: Path to the engine file
            engine_dir : CWD for the engines .

        """
        self.engine_file = os.path.expanduser(filepath)

        engine_json = None
        try:
            with open(self.engine_file, 'r') as f:
                engine_json = f.read()

        except OSError as e:
            logger.error("Could not open engine_json : ", self.engine_file)
            raise e

        return """cd {0}
cat <<EOF > ipengine.json
{1}
EOF

mkdir -p '.ipengine_logs'
ipengine --file=ipengine.json &>> .ipengine_logs/$JOBNAME.log
""".format(engine_dir, engine_json)

    def compose_containerized_launch_cmd(self, filepath, engine_dir, container_image):
        """Reads the json contents from filepath and uses that to compose the engine launch command.

        Args:
            filepath (str): Path to the engine file
            engine_dir (str): CWD for the engines .
            container_image (str): The container to be used to launch workers
        """
        self.engine_file = os.path.expanduser(filepath)

        engine_json = None
        try:
            with open(self.engine_file, 'r') as f:
                engine_json = f.read()

        except OSError as e:
            logger.error("Could not open engine_json : ", self.engine_file)
            raise e

        return """cd {0}
cat <<EOF > ipengine.json
{1}
EOF

DOCKER_ID=$(docker create --network host {2} ipengine --file=ipengine.json)
docker cp ipengine.json $DOCKER_ID:/home/ipengine.json
docker start $DOCKER_ID
""".format(engine_dir, engine_json, container_image)

    def __init__(self, execution_provider=None,
                 reuse_controller=True,
                 engine_json_file='~/.ipython/profile_default/security/ipcontroller-engine.json',
                 engine_dir='.',
                 controller=None,
                 config=None):
        """Initialize the IPyParallel pool. The initialization takes all relevant parameters via KWargs.

        .. note::

              If initBlocks > 0, and a scalable execution_provider is attached, then the provider
              will be initialized here.

        Args:
             - self

        KWargs:
             - execution_provider (ExecutionProvider object)
             - reuse_controller (Bool) : If True ipp executor will attempt to connect to an available
               controller. Default: True
             - engine_json_file (str): Path to json engine file that will be used to compose ipp launch
               commands at scaling events. Default : '~/.ipython/profile_default/security/ipcontroller-engine.json'
             - engine_dir (str) : Alternative to above, specify the engine_dir
             - config (dict). Default: '.'
        """
        self.controller = controller
        self.engine_file = engine_json_file
        self.client_file = None

        if self.controller:
            # Find the Client json
            self.client_file = self.controller.client_file
            self.engine_file = self.controller.engine_file

            if not os.path.exists(self.client_file):
                logger.debug("Waiting for {0}".format(self.client_file))

            sleep_dur = 20  # 20 seconds
            for i in range(0, int(sleep_dur / 0.2)):
                time.sleep(0.2)
                if os.path.exists(self.client_file):
                    break

            if not os.path.exists(self.client_file):
                raise Exception("Controller client file is missing at {0}".format(self.client_file))

        self.executor = Client(url_file=self.client_file)
        self.config = config
        self.container_type = None
        self.container_image = None
        self.sitename = config['site'] if config else 'Static_IPP'
        # NOTE: Copying the config here only partially fixes the issue. There needs to be
        # multiple controllers launched by the factory, and each must have different jsons.
        # There could be timing issues here,
        # local_engine_json = "{0}.{1}.engine.json".format(self.config["site"], int(time.time()))
        # copyfile(engine_json_file, local_engine_json)
        # if not os.path.exists(self.config["execution"]["script_dir"]):
        #    os.makedirs(self.config["execution"]["script_dir"])

        command_composer = self.compose_launch_cmd

        if config and config["execution"]["container"]["type"]:
            self.container_type = config["execution"]["container"]["type"]
            self.container_image = config["execution"]["container"]["image"]
            command_composer = self.compose_containerized_launch_cmd
            logger.info("Launching IPP with Docker:{0}".format(self.container_image))

        self.launch_cmd = command_composer(self.engine_file, engine_dir, self.container_image)
        logger.debug("Launch cmd : {}".format(self.launch_cmd))
        self.execution_provider = execution_provider
        self.engines = []

        if reuse_controller:
            # Reuse existing controller if one is available
            pass

        if execution_provider:
            self._scaling_enabled = True
            logger.debug("Starting IpyParallelExecutor with provider:%s", execution_provider)
            try:
                for i in range(self.config["execution"]["block"].get("initBlocks", 1)):
                    eng = self.execution_provider.submit(self.launch_cmd, 1)
                    logger.debug("Launched block : {0}:{1}".format(i, eng))
                    if not eng:
                        raise(ScalingFailed(self.execution_provider.sitename,
                                            "Ipp executor failed to scale via execution_provider"))
                    self.engines.extend([eng])

            except Exception as e:
                logger.error("Scaling out failed : %s", e)
                raise e

        else:
            self._scaling_enabled = False
            logger.debug("Starting IpyParallelExecutor with no provider")

        self.lb_view = self.executor.load_balanced_view()
        logger.debug("Starting executor")

    @property
    def scaling_enabled(self):
        return self._scaling_enabled

    def submit(self, *args, **kwargs):
        """Submits work to the thread pool.

        This method is simply pass through and behaves like a submit call as described
        here `Python docs: <https://docs.python.org/3/library/concurrent.futures.html#concurrent.futures.ThreadPoolExecutor>`_

        Returns:
              Future
        """
        # logger.debug("Got args : %s,", args)
        # logger.debug("Got kwargs : %s,", kwargs)
        return self.lb_view.apply_async(*args, **kwargs)

    def scale_out(self, *args, **kwargs):
        """Scales out the number of active workers by 1.

        This method is notImplemented for threads and will raise the error if called.

        """
        if self.execution_provider:
            r = self.execution_provider.submit(self.launch_cmd, *args, **kwargs)
            self.engines.extend([r])
        else:
            logger.error("No execution provider available")
            r = None

        return r

    def scale_in(self, blocks, *args, **kwargs):
        """Scale in the number of active workers by 1.

        This method is notImplemented for threads and will raise the error if called.

        Raises:
             NotImplemented exception
        """
        status = dict(zip(self.engines, self.execution_provider.status(self.engines)))

        # This works for blocks=0
        to_kill = [engine for engine in status if status[engine] == "RUNNING"][:blocks]

        if self.execution_provider:
            r = self.execution_provider.cancel(to_kill, *args, **kwargs)
        else:
            logger.error("No execution provider available")
            r = None

        return r

    def status(self):
        """Returns the status of the executor via probing the execution providers."""
        if self.execution_provider:
            status = self.execution_provider.status(self.engines)

        else:
            status = []

        return status

    def shutdown(self, hub=True, targets='all', block=False):
        """Shutdown the executor, including all workers and controllers.

        The interface documentation for IPP is `here <http://ipyparallel.readthedocs.io/en/latest/api/ipyparallel.html#ipyparallel.Client.shutdown>`_

        Kwargs:
            - hub (Bool): Whether the hub should be shutdown, Default:True,
            - targets (list of ints| 'all'): List of engine id's to kill, Default:'all'
            - block (Bool): To block for confirmations or not

        Raises:
             NotImplementedError
        """
        if self.controller:
            logger.debug("IPP:Shutdown sequence: Attempting controller kill")
            self.controller.close()

        # We do not actually do executor.shutdown because
        # this blocks even when requested to not block, killing the
        # controller is more effective although impolite.
        # x = self.executor.shutdown(targets=targets,
        #                           hub=hub,
        #                           block=block)

        logger.debug("Done with executor shutdown")
        return True

    def __repr__(self):
        return "<IPP Executor for site:{0}>".format(self.sitename)

if __name__ == "__main__":

    pool1_config = {"poolname": "pool1",
                    "queue": "foo"}
<|endoftext|>


<commit_before>#!/usr/bin/env python
# Copyright 2009 Google Inc. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#      http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""Provides data sources to use for benchmarking."""

import ConfigParser
import glob
import os
import os.path
import random
import re
import subprocess
import sys
import time

# relative
import addr_util
import selectors
import util

# Pick the most accurate timer for a platform. Stolen from timeit.py:
if sys.platform[:3] == 'win':
  DEFAULT_TIMER = time.clock
else:
  DEFAULT_TIMER = time.time

GLOBAL_DATA_CACHE = {}

DEFAULT_CONFIG_PATH = 'config/data_sources.cfg'
MAX_NON_UNIQUE_RECORD_COUNT = 500000
MAX_FILE_MTIME_AGE_DAYS = 45
MIN_FILE_SIZE = 10000
MIN_RECOMMENDED_RECORD_COUNT = 200
MAX_FQDN_SYNTHESIZE_PERCENT = 4

class DataSources(object):
  """A collection of methods related to available hostname data sources."""

  def __init__(self, config_path=DEFAULT_CONFIG_PATH, status_callback=None):
    self.source_cache = GLOBAL_DATA_CACHE
    self.source_config = {}
    self.status_callback = status_callback
    self._LoadConfigFromPath(config_path)

  def msg(self, msg, **kwargs):
    if self.status_callback:
      self.status_callback(msg, **kwargs)
    else:
      print '- %s' % msg

  def _LoadConfigFromPath(self, path):
    """Load a configuration file describing data sources that may be available."""
    conf_file = util.FindDataFile(path)
    config = ConfigParser.ConfigParser()
    config.read(conf_file)
    for section in config.sections():
      if section not in self.source_config:
        self.source_config[section] = {
            'name': None,
            'search_paths': set(),
            # Store whether or not this data source contains personal data
            'full_hostnames': True,
            'synthetic': False,
            'include_duplicates': False,
            'max_mtime_days': MAX_FILE_MTIME_AGE_DAYS
        }

      for (key, value) in config.items(section):
        if key == 'name':
          self.source_config[section]['name'] = value
        elif key == 'full_hostnames' and int(value) == 0:
          self.source_config[section]['full_hostnames'] = False
        elif key == 'max_mtime_days':
          self.source_config[section]['max_mtime_days'] = int(value)
        elif key == 'include_duplicates':
          self.source_config[section]['include_duplicates'] = bool(value)
        elif key == 'synthetic':
          self.source_config[section]['synthetic'] = bool(value)
        else:
          self.source_config[section]['search_paths'].add(value)

  def ListSourceTypes(self):
    """Get a list of all data sources we know about."""
    return sorted(self.source_config.keys())

  def ListSourcesWithDetails(self):
    """Get a list of all data sources found with total counts.

    Returns:
      List of tuples in form of (short_name, full_name, full_hosts, # of entries)
    """
    for source in self.ListSourceTypes():
      max_mtime = self.source_config[source]['max_mtime_days']
      self._GetHostsFromSource(source, min_file_size=MIN_FILE_SIZE,
                               max_mtime_age_days=max_mtime)

    details = []
    for source in self.source_cache:
      details.append((source,
                      self.source_config[source]['name'],
                      self.source_config[source]['synthetic'],
                      len(self.source_cache[source])))
    return sorted(details, key=lambda x: (x[2], x[3] * -1))

  def ListSourceTitles(self):
    """Return a list of sources in title + count format."""
    titles = []
    seen_synthetic = False
    for (unused_type, name, is_synthetic, count) in self.ListSourcesWithDetails():
      if is_synthetic and not seen_synthetic:
        titles.append('-' * 36)
        seen_synthetic = True
      titles.append('%s (%s)' % (name, count))
    return titles

  def ConvertSourceTitleToType(self, detail):
    """Convert a detail name to a source type."""
    for source_type in self.source_config:
      if detail.startswith(self.source_config[source_type]['name']):
        return source_type

  def GetBestSourceDetails(self):
    return self.ListSourcesWithDetails()[0]

  def GetNameForSource(self, source):
    if source in self.source_config:
      return self.source_config[source]['name']
    else:
      # Most likely a custom file path
      return source

  def GetCachedRecordCountForSource(self, source):
    return len(self.source_cache[source])

  def _CreateRecordsFromHostEntries(self, entries, include_duplicates=False):
    """Create records from hosts, removing duplicate entries and IP's.

    Args:
      entries: A list of test-data entries.
      include_duplicates: Whether or not to filter duplicates (optional: False)

    Returns:
      A tuple of (filtered records, full_host_names (Boolean)

    Raises:
      ValueError: If no records could be grokked from the input.
    """
    last_entry = None

    records = []
    full_host_count = 0
    for entry in entries:
      if entry == last_entry and not include_duplicates:
        continue
      else:
        last_entry = entry

      if ' ' in entry:
        (record_type, host) = entry.split(' ')
      else:
        record_type = 'A'
        host = entry

      if not addr_util.IP_RE.match(host) and not addr_util.INTERNAL_RE.search(host):
        if not host.endswith('.'):
          host += '.'
        records.append((record_type, host))

        if addr_util.FQDN_RE.match(host):
          full_host_count += 1

    if not records:
      raise ValueError('No records could be created from: %s' % entries)

    # Now that we've read everything, are we dealing with domains or full hostnames?
    full_host_percent = full_host_count / float(len(records)) * 100
    if full_host_percent < MAX_FQDN_SYNTHESIZE_PERCENT:
      full_host_names = True
    else:
      full_host_names = False
    return (records, full_host_names)

  def GetTestsFromSource(self, source, count=50, select_mode=None):
    """Parse records from source, and return tuples to use for testing.

    Args:
      source: A source name (str) that has been configured.
      count: Number of tests to generate from the source (int)
      select_mode: automatic, weighted, random, chunk (str)

    Returns:
      A list of record tuples in the form of (req_type, hostname)

    Raises:
      ValueError: If no usable records are found from the data source.

    This is tricky because we support 3 types of input data:

    - List of domains
    - List of hosts
    - List of record_type + hosts
    """
    records = []

    if source in self.source_config:
      include_duplicates = self.source_config[source].get('include_duplicates', False)
    else:
      include_duplicates = False

    records = self._GetHostsFromSource(source)
    if not records:
      raise ValueError('Unable to generate records from %s (nothing found)' % source)

    self.msg('Generating tests from %s (%s records, selecting %s %s)'
             % (self.GetNameForSource(source), len(records), count, select_mode))
    (records, are_records_fqdn) = self._CreateRecordsFromHostEntries(records,
                                                                     include_duplicates=include_duplicates)
    # First try to resolve whether to use weighted or random.
    if select_mode in ('weighted', 'automatic', None):
      # If we are in include_duplicates mode (cachemiss, cachehit, etc.), we have different rules.
      if include_duplicates:
        if count > len(records):
          select_mode = 'random'
        else:
          select_mode = 'chunk'
      elif len(records) != len(set(records)):
        if select_mode == 'weighted':
          self.msg('%s data contains duplicates, switching select_mode to random' % source)
        select_mode = 'random'
      else:
        select_mode = 'weighted'

    self.msg('Selecting %s out of %s sanitized records (%s mode).' %
             (count, len(records), select_mode))
    if select_mode == 'weighted':
      records = selectors.WeightedDistribution(records, count)
    elif select_mode == 'chunk':
      records = selectors.ChunkSelect(records, count)
    elif select_mode == 'random':
      records = selectors.RandomSelect(records, count, include_duplicates=include_duplicates)
    else:
      raise ValueError('No such final selection mode: %s' % select_mode)

    # For custom filenames
    if source not in self.source_config:
      self.source_config[source] = {'synthetic': True}

    if are_records_fqdn:
      self.source_config[source]['full_hostnames'] = False
      self.msg('%s input appears to be predominantly domain names. Synthesizing FQDNs' % source)
      synthesized = []
      for (req_type, hostname) in records:
        if not addr_util.FQDN_RE.match(hostname):
          hostname = self._GenerateRandomHostname(hostname)
        synthesized.append((req_type, hostname))
      return synthesized
    else:
      return records

  def _GenerateRandomHostname(self, domain):
    """Generate a random hostname f or a given domain."""
    oracle = random.randint(0, 100)
    if oracle < 70:
      return 'www.%s' % domain
    elif oracle < 95:
      return domain
    elif oracle < 98:
      return 'static.%s' % domain
    else:
      return 'cache-%s.%s' % (random.randint(0, 10), domain)

  def _GetHostsFromSource(self, source, min_file_size=None, max_mtime_age_days=None):
    """Get data for a particular source. This needs to be fast.

    Args:
      source: A configured source type (str)
      min_file_size: What the minimum allowable file size is for this source (int)
      max_mtime_age_days: Maximum days old the file can be for this source (int)

    Returns:
      list of hostnames gathered from data source.

    The results of this function are cached by source type.
    """
    if source in self.source_cache:
      return self.source_cache[source]
    filename = self._FindBestFileForSource(source, min_file_size=min_file_size,
                                           max_mtime_age_days=max_mtime_age_days)
    if not filename:
      return None

    size_mb = os.path.getsize(filename) / 1024.0 / 1024.0
    self.msg('Reading %s: %s (%0.1fMB)' % (self.GetNameForSource(source), filename, size_mb))
    start_clock = DEFAULT_TIMER()
    if filename.endswith('.pcap') or filename.endswith('.tcp'):
      hosts = self._ExtractHostsFromPcapFile(filename)
    else:
      hosts = self._ExtractHostsFromHistoryFile(filename)

    if not hosts:
      hosts = self._ReadDataFile(filename)
    duration = DEFAULT_TIMER() - start_clock
    if duration > 5:
      self.msg('%s data took %1.1fs to read!' % (self.GetNameForSource(source), duration))
    self.source_cache[source] = hosts
    return hosts

  def _ExtractHostsFromHistoryFile(self, path):
    """Get a list of sanitized records from a history file containing URLs."""
    # This regexp is fairly general (no ip filtering), since we need speed more
    # than precision at this stage.
    parse_re = re.compile('https*://([\-\w]+\.[\-\w\.]+)')
    return parse_re.findall(open(path, 'rb').read())

  def _ExtractHostsFromPcapFile(self, path):
    """Get a list of requests out of a pcap file - requires tcpdump."""
    self.msg('Extracting requests from pcap file using tcpdump')
    cmd = 'tcpdump -r %s -n port 53' % path
    pipe = subprocess.Popen(cmd, shell=True, stdout=subprocess.PIPE).stdout
    parse_re = re.compile(' ([A-Z]+)\? ([\-\w\.]+)')
    requests = []
    for line in pipe:
      if '?' not in line:
        continue
      match = parse_re.search(line)
      if match:
        requests.append(' '.join(match.groups()))
    return requests

  def _ReadDataFile(self, path):
    """Read a line-based datafile."""
    records = []
    domains_re = re.compile('^\w[\w\.-]+[a-zA-Z]$')
    requests_re = re.compile('^[A-Z]{1,4} \w[\w\.-]+\.$')
    for line in open(path):
      if domains_re.match(line) or requests_re.match(line):
        records.append(line.rstrip())
    return records

  def _GetSourceSearchPaths(self, source):
    """Get a list of possible search paths (globs) for a given source."""

    # This is likely a custom file path
    if source not in self.source_config:
      return [source]

    search_paths = []
    environment_re = re.compile('%(\w+)%')

    # First get through resolving environment variables
    for path in self.source_config[source]['search_paths']:
      env_vars = set(environment_re.findall(path))
      if env_vars:
        for variable in env_vars:
          env_var = os.getenv(variable, False)
          if env_var:
            path = path.replace('%%%s%%' % variable, env_var)
          else:
            path = None

      # If everything is good, replace all '/'  chars with the os path variable.
      if path:
        path = path.replace('/', os.sep)
        search_paths.append(path)

        # This moment of weirdness brought to you by Windows XP(tm). If we find
        # a Local or Roaming keyword in path, add the other forms to the search
        # path.
        if sys.platform[:3] == 'win':
          keywords = ('Local', 'Roaming')
          for keyword in keywords:
            if keyword in path:
              replacement = keywords[keywords.index(keyword)-1]
              search_paths.append(path.replace('\\%s' % keyword, '\\%s' % replacement))
              search_paths.append(path.replace('\\%s' % keyword, ''))

    return search_paths

  def _FindBestFileForSource(self, source, min_file_size=None, max_mtime_age_days=None):
    """Find the best file (newest over X size) to use for a given source type.

    Args:
      source: source type
      min_file_size: What the minimum allowable file size is for this source (int)
      max_mtime_age_days: Maximum days old the file can be for this source (int)

    Returns:
      A file path.
    """
    found = []
    for path in self._GetSourceSearchPaths(source):
      if not os.path.isabs(path):
        path = util.FindDataFile(path)

      for filename in glob.glob(path):
        if min_file_size and os.path.getsize(filename) < min_file_size:
          self.msg('Skipping %s (only %sb)' % (filename, os.path.getsize(filename)))
        else:
          found.append(filename)

    if found:
      newest = sorted(found, key=os.path.getmtime)[-1]
      age_days = (time.time() - os.path.getmtime(newest)) / 86400
      if max_mtime_age_days and age_days > max_mtime_age_days:
        self.msg('Skipping %s (%2.0fd old)' % (newest, age_days))
      else:
        return newest
    else:
      return None

if __name__ == '__main__':
  parser = DataSources()
  print parser.ListSourceTypes()
  print parser.ListSourcesWithDetails()
  best = parser.ListSourcesWithDetails()[0][0]
  print len(parser.GetRecordsFromSource(best))
<commit_msg>Only add a separator in the data sources if we found an organic source<commit_after>#!/usr/bin/env python
# Copyright 2009 Google Inc. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#      http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""Provides data sources to use for benchmarking."""

import ConfigParser
import glob
import os
import os.path
import random
import re
import subprocess
import sys
import time

# relative
import addr_util
import selectors
import util

# Pick the most accurate timer for a platform. Stolen from timeit.py:
if sys.platform[:3] == 'win':
  DEFAULT_TIMER = time.clock
else:
  DEFAULT_TIMER = time.time

GLOBAL_DATA_CACHE = {}

DEFAULT_CONFIG_PATH = 'config/data_sources.cfg'
MAX_NON_UNIQUE_RECORD_COUNT = 500000
MAX_FILE_MTIME_AGE_DAYS = 45
MIN_FILE_SIZE = 10000
MIN_RECOMMENDED_RECORD_COUNT = 200
MAX_FQDN_SYNTHESIZE_PERCENT = 4

class DataSources(object):
  """A collection of methods related to available hostname data sources."""

  def __init__(self, config_path=DEFAULT_CONFIG_PATH, status_callback=None):
    self.source_cache = GLOBAL_DATA_CACHE
    self.source_config = {}
    self.status_callback = status_callback
    self._LoadConfigFromPath(config_path)

  def msg(self, msg, **kwargs):
    if self.status_callback:
      self.status_callback(msg, **kwargs)
    else:
      print '- %s' % msg

  def _LoadConfigFromPath(self, path):
    """Load a configuration file describing data sources that may be available."""
    conf_file = util.FindDataFile(path)
    config = ConfigParser.ConfigParser()
    config.read(conf_file)
    for section in config.sections():
      if section not in self.source_config:
        self.source_config[section] = {
            'name': None,
            'search_paths': set(),
            # Store whether or not this data source contains personal data
            'full_hostnames': True,
            'synthetic': False,
            'include_duplicates': False,
            'max_mtime_days': MAX_FILE_MTIME_AGE_DAYS
        }

      for (key, value) in config.items(section):
        if key == 'name':
          self.source_config[section]['name'] = value
        elif key == 'full_hostnames' and int(value) == 0:
          self.source_config[section]['full_hostnames'] = False
        elif key == 'max_mtime_days':
          self.source_config[section]['max_mtime_days'] = int(value)
        elif key == 'include_duplicates':
          self.source_config[section]['include_duplicates'] = bool(value)
        elif key == 'synthetic':
          self.source_config[section]['synthetic'] = bool(value)
        else:
          self.source_config[section]['search_paths'].add(value)

  def ListSourceTypes(self):
    """Get a list of all data sources we know about."""
    return sorted(self.source_config.keys())

  def ListSourcesWithDetails(self):
    """Get a list of all data sources found with total counts.

    Returns:
      List of tuples in form of (short_name, full_name, full_hosts, # of entries)
    """
    for source in self.ListSourceTypes():
      max_mtime = self.source_config[source]['max_mtime_days']
      self._GetHostsFromSource(source, min_file_size=MIN_FILE_SIZE,
                               max_mtime_age_days=max_mtime)

    details = []
    for source in self.source_cache:
      details.append((source,
                      self.source_config[source]['name'],
                      self.source_config[source]['synthetic'],
                      len(self.source_cache[source])))
    return sorted(details, key=lambda x: (x[2], x[3] * -1))

  def ListSourceTitles(self):
    """Return a list of sources in title + count format."""
    titles = []
    seen_synthetic = False
    seen_organic = False
    for (unused_type, name, is_synthetic, count) in self.ListSourcesWithDetails():
      if not is_synthetic:
        seen_organic = True

      if is_synthetic and seen_organic and not seen_synthetic:
        titles.append('-' * 36)
        seen_synthetic = True
      titles.append('%s (%s)' % (name, count))
    return titles

  def ConvertSourceTitleToType(self, detail):
    """Convert a detail name to a source type."""
    for source_type in self.source_config:
      if detail.startswith(self.source_config[source_type]['name']):
        return source_type

  def GetBestSourceDetails(self):
    return self.ListSourcesWithDetails()[0]

  def GetNameForSource(self, source):
    if source in self.source_config:
      return self.source_config[source]['name']
    else:
      # Most likely a custom file path
      return source

  def GetCachedRecordCountForSource(self, source):
    return len(self.source_cache[source])

  def _CreateRecordsFromHostEntries(self, entries, include_duplicates=False):
    """Create records from hosts, removing duplicate entries and IP's.

    Args:
      entries: A list of test-data entries.
      include_duplicates: Whether or not to filter duplicates (optional: False)

    Returns:
      A tuple of (filtered records, full_host_names (Boolean)

    Raises:
      ValueError: If no records could be grokked from the input.
    """
    last_entry = None

    records = []
    full_host_count = 0
    for entry in entries:
      if entry == last_entry and not include_duplicates:
        continue
      else:
        last_entry = entry

      if ' ' in entry:
        (record_type, host) = entry.split(' ')
      else:
        record_type = 'A'
        host = entry

      if not addr_util.IP_RE.match(host) and not addr_util.INTERNAL_RE.search(host):
        if not host.endswith('.'):
          host += '.'
        records.append((record_type, host))

        if addr_util.FQDN_RE.match(host):
          full_host_count += 1

    if not records:
      raise ValueError('No records could be created from: %s' % entries)

    # Now that we've read everything, are we dealing with domains or full hostnames?
    full_host_percent = full_host_count / float(len(records)) * 100
    if full_host_percent < MAX_FQDN_SYNTHESIZE_PERCENT:
      full_host_names = True
    else:
      full_host_names = False
    return (records, full_host_names)

  def GetTestsFromSource(self, source, count=50, select_mode=None):
    """Parse records from source, and return tuples to use for testing.

    Args:
      source: A source name (str) that has been configured.
      count: Number of tests to generate from the source (int)
      select_mode: automatic, weighted, random, chunk (str)

    Returns:
      A list of record tuples in the form of (req_type, hostname)

    Raises:
      ValueError: If no usable records are found from the data source.

    This is tricky because we support 3 types of input data:

    - List of domains
    - List of hosts
    - List of record_type + hosts
    """
    records = []

    if source in self.source_config:
      include_duplicates = self.source_config[source].get('include_duplicates', False)
    else:
      include_duplicates = False

    records = self._GetHostsFromSource(source)
    if not records:
      raise ValueError('Unable to generate records from %s (nothing found)' % source)

    self.msg('Generating tests from %s (%s records, selecting %s %s)'
             % (self.GetNameForSource(source), len(records), count, select_mode))
    (records, are_records_fqdn) = self._CreateRecordsFromHostEntries(records,
                                                                     include_duplicates=include_duplicates)
    # First try to resolve whether to use weighted or random.
    if select_mode in ('weighted', 'automatic', None):
      # If we are in include_duplicates mode (cachemiss, cachehit, etc.), we have different rules.
      if include_duplicates:
        if count > len(records):
          select_mode = 'random'
        else:
          select_mode = 'chunk'
      elif len(records) != len(set(records)):
        if select_mode == 'weighted':
          self.msg('%s data contains duplicates, switching select_mode to random' % source)
        select_mode = 'random'
      else:
        select_mode = 'weighted'

    self.msg('Selecting %s out of %s sanitized records (%s mode).' %
             (count, len(records), select_mode))
    if select_mode == 'weighted':
      records = selectors.WeightedDistribution(records, count)
    elif select_mode == 'chunk':
      records = selectors.ChunkSelect(records, count)
    elif select_mode == 'random':
      records = selectors.RandomSelect(records, count, include_duplicates=include_duplicates)
    else:
      raise ValueError('No such final selection mode: %s' % select_mode)

    # For custom filenames
    if source not in self.source_config:
      self.source_config[source] = {'synthetic': True}

    if are_records_fqdn:
      self.source_config[source]['full_hostnames'] = False
      self.msg('%s input appears to be predominantly domain names. Synthesizing FQDNs' % source)
      synthesized = []
      for (req_type, hostname) in records:
        if not addr_util.FQDN_RE.match(hostname):
          hostname = self._GenerateRandomHostname(hostname)
        synthesized.append((req_type, hostname))
      return synthesized
    else:
      return records

  def _GenerateRandomHostname(self, domain):
    """Generate a random hostname f or a given domain."""
    oracle = random.randint(0, 100)
    if oracle < 70:
      return 'www.%s' % domain
    elif oracle < 95:
      return domain
    elif oracle < 98:
      return 'static.%s' % domain
    else:
      return 'cache-%s.%s' % (random.randint(0, 10), domain)

  def _GetHostsFromSource(self, source, min_file_size=None, max_mtime_age_days=None):
    """Get data for a particular source. This needs to be fast.

    Args:
      source: A configured source type (str)
      min_file_size: What the minimum allowable file size is for this source (int)
      max_mtime_age_days: Maximum days old the file can be for this source (int)

    Returns:
      list of hostnames gathered from data source.

    The results of this function are cached by source type.
    """
    if source in self.source_cache:
      return self.source_cache[source]
    filename = self._FindBestFileForSource(source, min_file_size=min_file_size,
                                           max_mtime_age_days=max_mtime_age_days)
    if not filename:
      return None

    size_mb = os.path.getsize(filename) / 1024.0 / 1024.0
    self.msg('Reading %s: %s (%0.1fMB)' % (self.GetNameForSource(source), filename, size_mb))
    start_clock = DEFAULT_TIMER()
    if filename.endswith('.pcap') or filename.endswith('.tcp'):
      hosts = self._ExtractHostsFromPcapFile(filename)
    else:
      hosts = self._ExtractHostsFromHistoryFile(filename)

    if not hosts:
      hosts = self._ReadDataFile(filename)
    duration = DEFAULT_TIMER() - start_clock
    if duration > 5:
      self.msg('%s data took %1.1fs to read!' % (self.GetNameForSource(source), duration))
    self.source_cache[source] = hosts
    return hosts

  def _ExtractHostsFromHistoryFile(self, path):
    """Get a list of sanitized records from a history file containing URLs."""
    # This regexp is fairly general (no ip filtering), since we need speed more
    # than precision at this stage.
    parse_re = re.compile('https*://([\-\w]+\.[\-\w\.]+)')
    return parse_re.findall(open(path, 'rb').read())

  def _ExtractHostsFromPcapFile(self, path):
    """Get a list of requests out of a pcap file - requires tcpdump."""
    self.msg('Extracting requests from pcap file using tcpdump')
    cmd = 'tcpdump -r %s -n port 53' % path
    pipe = subprocess.Popen(cmd, shell=True, stdout=subprocess.PIPE).stdout
    parse_re = re.compile(' ([A-Z]+)\? ([\-\w\.]+)')
    requests = []
    for line in pipe:
      if '?' not in line:
        continue
      match = parse_re.search(line)
      if match:
        requests.append(' '.join(match.groups()))
    return requests

  def _ReadDataFile(self, path):
    """Read a line-based datafile."""
    records = []
    domains_re = re.compile('^\w[\w\.-]+[a-zA-Z]$')
    requests_re = re.compile('^[A-Z]{1,4} \w[\w\.-]+\.$')
    for line in open(path):
      if domains_re.match(line) or requests_re.match(line):
        records.append(line.rstrip())
    return records

  def _GetSourceSearchPaths(self, source):
    """Get a list of possible search paths (globs) for a given source."""

    # This is likely a custom file path
    if source not in self.source_config:
      return [source]

    search_paths = []
    environment_re = re.compile('%(\w+)%')

    # First get through resolving environment variables
    for path in self.source_config[source]['search_paths']:
      env_vars = set(environment_re.findall(path))
      if env_vars:
        for variable in env_vars:
          env_var = os.getenv(variable, False)
          if env_var:
            path = path.replace('%%%s%%' % variable, env_var)
          else:
            path = None

      # If everything is good, replace all '/'  chars with the os path variable.
      if path:
        path = path.replace('/', os.sep)
        search_paths.append(path)

        # This moment of weirdness brought to you by Windows XP(tm). If we find
        # a Local or Roaming keyword in path, add the other forms to the search
        # path.
        if sys.platform[:3] == 'win':
          keywords = ('Local', 'Roaming')
          for keyword in keywords:
            if keyword in path:
              replacement = keywords[keywords.index(keyword)-1]
              search_paths.append(path.replace('\\%s' % keyword, '\\%s' % replacement))
              search_paths.append(path.replace('\\%s' % keyword, ''))

    return search_paths

  def _FindBestFileForSource(self, source, min_file_size=None, max_mtime_age_days=None):
    """Find the best file (newest over X size) to use for a given source type.

    Args:
      source: source type
      min_file_size: What the minimum allowable file size is for this source (int)
      max_mtime_age_days: Maximum days old the file can be for this source (int)

    Returns:
      A file path.
    """
    found = []
    for path in self._GetSourceSearchPaths(source):
      if not os.path.isabs(path):
        path = util.FindDataFile(path)

      for filename in glob.glob(path):
        if min_file_size and os.path.getsize(filename) < min_file_size:
          self.msg('Skipping %s (only %sb)' % (filename, os.path.getsize(filename)))
        else:
          found.append(filename)

    if found:
      newest = sorted(found, key=os.path.getmtime)[-1]
      age_days = (time.time() - os.path.getmtime(newest)) / 86400
      if max_mtime_age_days and age_days > max_mtime_age_days:
        self.msg('Skipping %s (%2.0fd old)' % (newest, age_days))
      else:
        return newest
    else:
      return None

if __name__ == '__main__':
  parser = DataSources()
  print parser.ListSourceTypes()
  print parser.ListSourcesWithDetails()
  best = parser.ListSourcesWithDetails()[0][0]
  print len(parser.GetRecordsFromSource(best))
<|endoftext|>


<commit_before># Copyright (c) 2021 Graeme Holliday
#
# Permission is hereby granted, free of charge, to any person obtaining a copy
# of this software and associated documentation files (the "Software"), to deal
# in the Software without restriction, including without limitation the rights
# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
# copies of the Software, and to permit persons to whom the Software is
# furnished to do so, subject to the following conditions:
#
# The above copyright notice and this permission notice shall be included in
# all copies or substantial portions of the Software.
#
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
# SOFTWARE.

from dbus_next.aio import MessageBus
from dbus_next.constants import BusType

from libqtile.widget import base

BLUEZ = 'org.bluez'
BLUEZ_PATH = '/org/bluez/hci0'
BLUEZ_ADAPTER = 'org.bluez.Adapter1'
BLUEZ_DEVICE = 'org.bluez.Device1'
BLUEZ_PROPERTIES = 'org.freedesktop.DBus.Properties'

class Bluetooth(base._TextBox):
    """
    Displays bluetooth status or connected device.

    Uses dbus to communicate with the system bus.
    """

    orientations = base.ORIENTATION_HORIZONTAL
    defaults = [
        ('hci', '/dev_XX_XX_XX_XX_XX_XX', 'hci0 device path, can be found with d-feet or similar dbus explorer.')
    ]

    def __init__(self, **config):
        base._TextBox.__init__(self, '', **config)
        self.add_defaults(Bluetooth.defaults)

    async def _config_async(self):
        # set initial values
        self.powered = await self._init_adapter()
        self.connected, self.device = await self._init_device()

        self.update_text()

    async def _init_adapter(self):
        # set up interface to adapter properties using high-level api
        bus = await MessageBus(bus_type=BusType.SYSTEM).connect()
        introspect = await bus.introspect(BLUEZ, BLUEZ_PATH)
        obj = bus.get_proxy_object(BLUEZ, BLUEZ_PATH, introspect)
        iface = obj.get_interface(BLUEZ_ADAPTER)
        props = obj.get_interface(BLUEZ_PROPERTIES)

        powered = await iface.get_powered()
        # subscribe receiver to property changed
        props.on_properties_changed(self._signal_received_adapter)
        return powered

    async def _init_device(self):
        # set up interface to device properties using high-level api
        bus = await MessageBus(bus_type=BusType.SYSTEM).connect()
        introspect = await bus.introspect(BLUEZ, BLUEZ_PATH + self.hci)
        obj = bus.get_proxy_object(BLUEZ, BLUEZ_PATH + self.hci, introspect)
        iface = obj.get_interface(BLUEZ_DEVICE)
        props = obj.get_interface(BLUEZ_PROPERTIES)

        connected = await iface.get_connected()
        name = await iface.get_name()
        # subscribe receiver to property changed
        props.on_properties_changed(self._signal_received_device)
        return connected, name

    def _signal_received_adapter(self, interface_name, changed_properties, invalidated_properties):
        powered = changed_properties.get('Powered', None)
        if powered is not None:
            self.powered = powered.value
            self.update_text()

    def _signal_received_device(self, interface_name, changed_properties, invalidated_properties):
        connected = changed_properties.get('Connected', None)
        if connected is not None:
            self.connected = connected.value
            self.update_text()

        device = changed_properties.get('Name', None)
        if device is not None:
            self.device = device.value
            self.update_text()

    def update_text(self):
        if not self.powered:
            self.text = 'off'
        else:
            if not self.connected:
                self.text = 'on'
            else:
                self.text = self.device

        self.bar.draw()
<commit_msg>Handle all signals in one function; mypy fix<commit_after># Copyright (c) 2021 Graeme Holliday
#
# Permission is hereby granted, free of charge, to any person obtaining a copy
# of this software and associated documentation files (the "Software"), to deal
# in the Software without restriction, including without limitation the rights
# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
# copies of the Software, and to permit persons to whom the Software is
# furnished to do so, subject to the following conditions:
#
# The above copyright notice and this permission notice shall be included in
# all copies or substantial portions of the Software.
#
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
# SOFTWARE.

from dbus_next.aio import MessageBus # type: ignore
from dbus_next.constants import BusType # type: ignore

from libqtile.widget import base

BLUEZ = 'org.bluez'
BLUEZ_PATH = '/org/bluez/hci0'
BLUEZ_ADAPTER = 'org.bluez.Adapter1'
BLUEZ_DEVICE = 'org.bluez.Device1'
BLUEZ_PROPERTIES = 'org.freedesktop.DBus.Properties'

class Bluetooth(base._TextBox):
    """
    Displays bluetooth status or connected device.

    Uses dbus to communicate with the system bus.
    """

    orientations = base.ORIENTATION_HORIZONTAL
    defaults = [
        ('hci', '/dev_XX_XX_XX_XX_XX_XX', 'hci0 device path, can be found with d-feet or similar dbus explorer.')
    ]

    def __init__(self, **config):
        base._TextBox.__init__(self, '', **config)
        self.add_defaults(Bluetooth.defaults)

    async def _config_async(self):
        # set initial values
        self.powered = await self._init_adapter()
        self.connected, self.device = await self._init_device()

        self.update_text()

    async def _init_adapter(self):
        # set up interface to adapter properties using high-level api
        bus = await MessageBus(bus_type=BusType.SYSTEM).connect()
        introspect = await bus.introspect(BLUEZ, BLUEZ_PATH)
        obj = bus.get_proxy_object(BLUEZ, BLUEZ_PATH, introspect)
        iface = obj.get_interface(BLUEZ_ADAPTER)
        props = obj.get_interface(BLUEZ_PROPERTIES)

        powered = await iface.get_powered()
        # subscribe receiver to property changed
        props.on_properties_changed(self._signal_received)
        return powered

    async def _init_device(self):
        # set up interface to device properties using high-level api
        bus = await MessageBus(bus_type=BusType.SYSTEM).connect()
        introspect = await bus.introspect(BLUEZ, BLUEZ_PATH + self.hci)
        obj = bus.get_proxy_object(BLUEZ, BLUEZ_PATH + self.hci, introspect)
        iface = obj.get_interface(BLUEZ_DEVICE)
        props = obj.get_interface(BLUEZ_PROPERTIES)

        connected = await iface.get_connected()
        name = await iface.get_name()
        # subscribe receiver to property changed
        props.on_properties_changed(self._signal_received)
        return connected, name

    def _signal_received(self, interface_name, changed_properties, invalidated_properties):
        powered = changed_properties.get('Powered', None)
        if powered is not None:
            self.powered = powered.value
            self.update_text()

        connected = changed_properties.get('Connected', None)
        if connected is not None:
            self.connected = connected.value
            self.update_text()

        device = changed_properties.get('Name', None)
        if device is not None:
            self.device = device.value
            self.update_text()

    def update_text(self):
        if not self.powered:
            self.text = 'off'
        else:
            if not self.connected:
                self.text = 'on'
            else:
                self.text = self.device

        self.bar.draw()
<|endoftext|>


<commit_before># coding=utf-8
import re

from subzero.modification.mods import SubtitleTextModification, empty_line_post_processors, EmptyEntryError
from subzero.modification.processors.re_processor import NReProcessor
from subzero.modification import registry

class FullBracketEntryProcessor(NReProcessor):
    def process(self, content, debug=False, **kwargs):
        entry = kwargs.get("entry")
        if entry:
            rep_content = super(FullBracketEntryProcessor, self).process(entry, debug=debug, **kwargs)
            if not rep_content.strip():
                raise EmptyEntryError()
        return content

class HearingImpaired(SubtitleTextModification):
    identifier = "remove_HI"
    description = "Remove Hearing Impaired tags"
    exclusive = True
    order = 20

    long_description = """\
    Removes tags, text and characters from subtitles that are meant for hearing impaired people
    """

    processors = [
        # full bracket entry, single or multiline; starting with brackets and ending with brackets
        FullBracketEntryProcessor(re.compile(ur'(?sux)^-?\s?[([].+(?=[^)\]]{3,}).+[)\]]$'), "",
                                  name="HI_brackets_full"),

        # brackets (only remove if at least 3 chars in brackets)
        NReProcessor(re.compile(ur'(?sux)-?\s*[([].+?(?=[A-zÀ-ž]{3,}).+?[)\]][\s:]*'), "", name="HI_brackets"),

        # text before colon (and possible dash in front), max 11 chars after the first whitespace (if any)
        # NReProcessor(re.compile(r'(?u)(^[A-z\-\'"_]+[\w\s]{0,11}:[^0-9{2}][\s]*)'), "", name="HI_before_colon"),

        # starting text before colon (at least 2 uppercase chars)
        NReProcessor(re.compile(ur'(?u)(^[\s-]*(?=[A-ZÀ-Ž]{2})[A-ZÀ-ŽA-zÀ-ž-_\s0-9"\']+:\s*)'), "", name="HI_before_colon"),

        # text in brackets at start, after optional dash, before colon or at end of line
        # fixme: may be too aggressive
        #NReProcessor(re.compile(ur'(?um)(^-?\s?[([][A-zÀ-ž-_\s]{3,}[)\]](?:(?=$)|:\s*))'), "",
        #             name="HI_brackets_special"),

        # all caps line (at least 4 consecutive uppercase chars)
        NReProcessor(re.compile(ur'(?u)(^(?=.*[A-ZÀ-Ž]{4,})[A-ZÀ-Ž-_\s]+$)'), "", name="HI_all_caps"),

        # dash in front
        # NReProcessor(re.compile(r'(?u)^\s*-\s*'), "", name="HI_starting_dash"),

        # all caps at start before new sentence
        NReProcessor(re.compile(ur'(?u)^(?=[A-ZÀ-Ž]{4,})[A-ZÀ-Ž-_\s]+\s([A-ZÀ-Ž][a-zà-ž].+)'), r"\1",
                     name="HI_starting_upper_then_sentence"),
    ]

    post_processors = empty_line_post_processors

registry.register(HearingImpaired)
<commit_msg>submod: HI: HI_before_colon: remove redundant regex data<commit_after># coding=utf-8
import re

from subzero.modification.mods import SubtitleTextModification, empty_line_post_processors, EmptyEntryError
from subzero.modification.processors.re_processor import NReProcessor
from subzero.modification import registry

class FullBracketEntryProcessor(NReProcessor):
    def process(self, content, debug=False, **kwargs):
        entry = kwargs.get("entry")
        if entry:
            rep_content = super(FullBracketEntryProcessor, self).process(entry, debug=debug, **kwargs)
            if not rep_content.strip():
                raise EmptyEntryError()
        return content

class HearingImpaired(SubtitleTextModification):
    identifier = "remove_HI"
    description = "Remove Hearing Impaired tags"
    exclusive = True
    order = 20

    long_description = """\
    Removes tags, text and characters from subtitles that are meant for hearing impaired people
    """

    processors = [
        # full bracket entry, single or multiline; starting with brackets and ending with brackets
        FullBracketEntryProcessor(re.compile(ur'(?sux)^-?\s?[([].+(?=[^)\]]{3,}).+[)\]]$'), "",
                                  name="HI_brackets_full"),

        # brackets (only remove if at least 3 chars in brackets)
        NReProcessor(re.compile(ur'(?sux)-?\s*[([].+?(?=[A-zÀ-ž]{3,}).+?[)\]][\s:]*'), "", name="HI_brackets"),

        # text before colon (and possible dash in front), max 11 chars after the first whitespace (if any)
        # NReProcessor(re.compile(r'(?u)(^[A-z\-\'"_]+[\w\s]{0,11}:[^0-9{2}][\s]*)'), "", name="HI_before_colon"),

        # starting text before colon (at least 2 uppercase chars)
        NReProcessor(re.compile(ur'(?u)(^[\s-]*(?=[A-ZÀ-Ž]{2})[A-zÀ-ž-_\s0-9"\']+:\s*)'), "", name="HI_before_colon"),

        # text in brackets at start, after optional dash, before colon or at end of line
        # fixme: may be too aggressive
        #NReProcessor(re.compile(ur'(?um)(^-?\s?[([][A-zÀ-ž-_\s]{3,}[)\]](?:(?=$)|:\s*))'), "",
        #             name="HI_brackets_special"),

        # all caps line (at least 4 consecutive uppercase chars)
        NReProcessor(re.compile(ur'(?u)(^(?=.*[A-ZÀ-Ž]{4,})[A-ZÀ-Ž-_\s]+$)'), "", name="HI_all_caps"),

        # dash in front
        # NReProcessor(re.compile(r'(?u)^\s*-\s*'), "", name="HI_starting_dash"),

        # all caps at start before new sentence
        NReProcessor(re.compile(ur'(?u)^(?=[A-ZÀ-Ž]{4,})[A-ZÀ-Ž-_\s]+\s([A-ZÀ-Ž][a-zà-ž].+)'), r"\1",
                     name="HI_starting_upper_then_sentence"),
    ]

    post_processors = empty_line_post_processors

registry.register(HearingImpaired)
<|endoftext|>