<commit_before>#    Copyright 2015-2017 ARM Limited
#
# 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.
#

"""Base class to parse trace.dat dumps"""
from __future__ import unicode_literals
from __future__ import division
from __future__ import print_function

from builtins import zip
from builtins import range
from builtins import object
from past.builtins import basestring
import re
import pandas as pd
import warnings

from resource import getrusage, RUSAGE_SELF

from trappy.exception import TrappyParseError

def _get_free_memory_kb():
    try:
        with open("/proc/meminfo") as f:
            memfree_line = [l for l in f.readlines() if "MemFree" in l][0]
            _, num_kb, _ = memfree_line.split()
            return int(num_kb)
    except:
        # Probably either not running on Linux (no /proc/meminfo), or format has
        # changed (we didn't find num_kb).
        return None

def trace_parser_explode_array(string, array_lengths):
    """Explode an array in the trace into individual elements for easy parsing

    Basically, turn :code:`load={1 1 2 2}` into :code:`load0=1 load1=1 load2=2
    load3=2`.

    :param string: Input string from the trace
    :type string: str

    :param array_lengths: A dictionary of array names and their
        expected length.  If we get array that's shorter than the expected
        length, additional keys have to be introduced with value 0 to
        compensate.
    :type array_lengths: dict

    For example:
    ::

        trace_parser_explode_array(string="load={1 2}",
                                   array_lengths={"load": 4})
        "load0=1 load1=2 load2=0 load3=0"
    """

    while True:
        match = re.search(r"[^ ]+={[^}]+}", string)
        if match is None:
            break

        to_explode = match.group()
        col_basename = re.match(r"([^=]+)=", to_explode).groups()[0]
        vals_str = re.search(r"{(.+)}", to_explode).groups()[0]
        vals_array = vals_str.split(' ')

        exploded_str = ""
        for (idx, val) in enumerate(vals_array):
            exploded_str += "{}{}={} ".format(col_basename, idx, val)

        vals_added = len(vals_array)
        if vals_added < array_lengths[col_basename]:
            for idx in range(vals_added, array_lengths[col_basename]):
                exploded_str += "{}{}=0 ".format(col_basename, idx)

        exploded_str = exploded_str[:-1]
        begin_idx = match.start()
        end_idx = match.end()

        string = string[:begin_idx] + exploded_str + string[end_idx:]

    return string

class Base(object):
    """Base class to parse trace.dat dumps.

    Don't use directly, create a subclass that has a unique_word class
    variable.  unique_word is a string that can uniquely identify
    lines in the trace that correspond to this event.  This is usually
    the trace_name (optionally followed by a semicolong,
    e.g. "sched_switch:") but it can be anything else for trace points
    generated using trace_printk().

    :param parse_raw: If :code:`True`, raw trace data (-r option) to
        trace-cmd will be used

    :param fallback: If :code:`True`, the parsing class will be used
        only if no other candidate class's unique_word matched. subclasses
        should override this (for ex. TracingMarkWrite uses it)

    This class acts as a base class for all TRAPpy events

    """

    def __init__(self, parse_raw=False, fallback=False):
        self.fallback = fallback
        self.tracer = None
        self.data_frame = pd.DataFrame()
        self.line_array = []
        self.data_array = []
        self.time_array = []
        self.comm_array = []
        self.pid_array = []
        self.cpu_array = []
        self.parse_raw = parse_raw
        self.cached = False

    def finalize_object(self):
        pass

    def __get_trace_array_lengths(self):
        """Calculate the lengths of all arrays in the trace

        Returns a dict with the name of each array found in the trace
        as keys and their corresponding length as value

        """
        from collections import defaultdict

        pat_array = re.compile(r"([A-Za-z0-9_]+)={([^}]+)}")

        ret = defaultdict(int)

        for line in self.data_array:
            while True:
                match = re.search(pat_array, line)
                if not match:
                    break

                (array_name, array_elements) = match.groups()

                array_len = len(array_elements.split(' '))

                if array_len > ret[array_name]:
                    ret[array_name] = array_len

                line = line[match.end():]

            # Stop scanning if the trace doesn't have arrays
            if len(ret) == 0:
                break

        return ret

    def append_data(self, time, comm, pid, cpu, line, data):
        """Append data parsed from a line to the corresponding arrays

        The :mod:`DataFrame` will be created from this when the whole trace
        has been parsed.

        :param time: The time for the line that was printed in the trace
        :type time: float

        :param comm: The command name or the execname from which the trace
            line originated
        :type comm: str

        :param pid: The PID of the process from which the trace
            line originated
        :type pid: int

        :param data: The data for matching line in the trace
        :type data: str
        """

        self.time_array.append(time)
        self.comm_array.append(comm)
        self.pid_array.append(pid)
        self.cpu_array.append(cpu)
        self.line_array.append(line)
        self.data_array.append(data)

    def string_cast(self, string, type):
        """ Attempt to convert string to another type

        Here we attempt to cast string to a type. Currently only
        integer conversion is supported with future expansion
        left open to other types.

        :param string: The value to convert.
        :type string: str

        :param type: The type to convert to.
        :type type: type
        """
        # Currently this function only supports int conversion
        if type != int:
            return
        # Handle false-positives for negative numbers
        if not string.lstrip("-").isdigit():
            return string
        return int(string)

    def generate_data_dict(self, data_str):
        data_dict = {}
        prev_key = None
        for field in data_str.split():
            if "=" not in field:
                if not prev_key:
                    if 'FAILED TO PARSE' in data_str:
                        warnings.warn(
                            'trace-cmd failed to parse the "{}" event. You may '
                            'need to compile the latest trace-cmd and put it in '
                            'your $PATH. Continuing, but some data may be missing'
                            .format(self.unique_word))
                        continue
                    else:
                        raise TrappyParseError(
                            "TRAPpy's parser for '{}' failed to parse the line:"
                            "\n{}".format(self.unique_word, data_str))
                # Concatenation is supported only for "string" values
                if not isinstance(data_dict[prev_key], basestring):
                    continue
                data_dict[prev_key] += ' ' + field
                continue
            (key, value) = field.split('=', 1)
            value = self.string_cast(value, int)
            data_dict[key] = value
            prev_key = key
        return data_dict

    def generate_parsed_data(self):

        # Get a rough idea of how much memory we have to play with
        CHECK_MEM_COUNT = 10000
        kb_free = _get_free_memory_kb()
        starting_maxrss = getrusage(RUSAGE_SELF).ru_maxrss
        check_memory_usage = True
        check_memory_count = 1

        for (comm, pid, cpu, line, data_str) in zip(self.comm_array, self.pid_array,
                                              self.cpu_array, self.line_array,
                                              self.data_array):
            data_dict = {"__comm": comm, "__pid": pid, "__cpu": cpu, "__line": line}
            data_dict.update(self.generate_data_dict(data_str))

            # When running out of memory, Pandas has been observed to segfault
            # rather than throwing a proper Python error.
            # Look at how much memory our process is using and warn if we seem
            # to be getting close to the system's limit, check it only once
            # in the beginning and then every CHECK_MEM_COUNT events
            check_memory_count -= 1
            if check_memory_usage and check_memory_count == 0:
                kb_used = (getrusage(RUSAGE_SELF).ru_maxrss - starting_maxrss)
                if kb_free and kb_used > kb_free * 0.9:
                    warnings.warn("TRAPpy: Appear to be low on memory. "
                                  "If errors arise, try providing more RAM")
                    check_memory_usage = False
                check_memory_count = CHECK_MEM_COUNT

            yield data_dict

    def create_dataframe(self):
        """Create the final :mod:`pandas.DataFrame`"""
        if not self.time_array:
            return

        trace_arr_lengths = self.__get_trace_array_lengths()

        if trace_arr_lengths:
            for (idx, val) in enumerate(self.data_array):
                expl_val = trace_parser_explode_array(val, trace_arr_lengths)
                self.data_array[idx] = expl_val

        time_idx = pd.Index(self.time_array, name="Time")
        self.data_frame = pd.DataFrame(self.generate_parsed_data(), index=time_idx)

        self.time_array = []
        self.line_array = []
        self.comm_array = []
        self.pid_array = []
        self.cpu_array = []
        self.data_array = []

    def write_csv(self, fname):
        """Write the csv info into a CSV file

        :param fname: The name of the CSV file
        :type fname: str
        """
        self.data_frame.to_csv(fname)

    def read_csv(self, fname):
        """Read the csv data into a DataFrame

        :param fname: The name of the CSV file
        :type fname: str
        """
        self.data_frame = pd.read_csv(
            fname,
            index_col=0,
            # This ensures cached vs parsed timestamps are converted using the
            # same method, aka python's float() and not numpy's
            converters={'Time' : float}
        )

    def normalize_time(self, basetime):
        """Substract basetime from the Time of the data frame

        :param basetime: The offset which needs to be subtracted from
            the time index
        :type basetime: float
        """
        if basetime and not self.data_frame.empty:
            self.data_frame.reset_index(inplace=True)
            self.data_frame["Time"] = self.data_frame["Time"] - basetime
            self.data_frame.set_index("Time", inplace=True)
<commit_msg>base: Fix string to int conversion<commit_after>#    Copyright 2015-2017 ARM Limited
#
# 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.
#

"""Base class to parse trace.dat dumps"""
from __future__ import unicode_literals
from __future__ import division
from __future__ import print_function

from builtins import zip
from builtins import range
from builtins import object
from past.builtins import basestring
import re
import pandas as pd
import warnings

from resource import getrusage, RUSAGE_SELF

from trappy.exception import TrappyParseError

def _get_free_memory_kb():
    try:
        with open("/proc/meminfo") as f:
            memfree_line = [l for l in f.readlines() if "MemFree" in l][0]
            _, num_kb, _ = memfree_line.split()
            return int(num_kb)
    except:
        # Probably either not running on Linux (no /proc/meminfo), or format has
        # changed (we didn't find num_kb).
        return None

def trace_parser_explode_array(string, array_lengths):
    """Explode an array in the trace into individual elements for easy parsing

    Basically, turn :code:`load={1 1 2 2}` into :code:`load0=1 load1=1 load2=2
    load3=2`.

    :param string: Input string from the trace
    :type string: str

    :param array_lengths: A dictionary of array names and their
        expected length.  If we get array that's shorter than the expected
        length, additional keys have to be introduced with value 0 to
        compensate.
    :type array_lengths: dict

    For example:
    ::

        trace_parser_explode_array(string="load={1 2}",
                                   array_lengths={"load": 4})
        "load0=1 load1=2 load2=0 load3=0"
    """

    while True:
        match = re.search(r"[^ ]+={[^}]+}", string)
        if match is None:
            break

        to_explode = match.group()
        col_basename = re.match(r"([^=]+)=", to_explode).groups()[0]
        vals_str = re.search(r"{(.+)}", to_explode).groups()[0]
        vals_array = vals_str.split(' ')

        exploded_str = ""
        for (idx, val) in enumerate(vals_array):
            exploded_str += "{}{}={} ".format(col_basename, idx, val)

        vals_added = len(vals_array)
        if vals_added < array_lengths[col_basename]:
            for idx in range(vals_added, array_lengths[col_basename]):
                exploded_str += "{}{}=0 ".format(col_basename, idx)

        exploded_str = exploded_str[:-1]
        begin_idx = match.start()
        end_idx = match.end()

        string = string[:begin_idx] + exploded_str + string[end_idx:]

    return string

class Base(object):
    """Base class to parse trace.dat dumps.

    Don't use directly, create a subclass that has a unique_word class
    variable.  unique_word is a string that can uniquely identify
    lines in the trace that correspond to this event.  This is usually
    the trace_name (optionally followed by a semicolong,
    e.g. "sched_switch:") but it can be anything else for trace points
    generated using trace_printk().

    :param parse_raw: If :code:`True`, raw trace data (-r option) to
        trace-cmd will be used

    :param fallback: If :code:`True`, the parsing class will be used
        only if no other candidate class's unique_word matched. subclasses
        should override this (for ex. TracingMarkWrite uses it)

    This class acts as a base class for all TRAPpy events

    """

    def __init__(self, parse_raw=False, fallback=False):
        self.fallback = fallback
        self.tracer = None
        self.data_frame = pd.DataFrame()
        self.line_array = []
        self.data_array = []
        self.time_array = []
        self.comm_array = []
        self.pid_array = []
        self.cpu_array = []
        self.parse_raw = parse_raw
        self.cached = False

    def finalize_object(self):
        pass

    def __get_trace_array_lengths(self):
        """Calculate the lengths of all arrays in the trace

        Returns a dict with the name of each array found in the trace
        as keys and their corresponding length as value

        """
        from collections import defaultdict

        pat_array = re.compile(r"([A-Za-z0-9_]+)={([^}]+)}")

        ret = defaultdict(int)

        for line in self.data_array:
            while True:
                match = re.search(pat_array, line)
                if not match:
                    break

                (array_name, array_elements) = match.groups()

                array_len = len(array_elements.split(' '))

                if array_len > ret[array_name]:
                    ret[array_name] = array_len

                line = line[match.end():]

            # Stop scanning if the trace doesn't have arrays
            if len(ret) == 0:
                break

        return ret

    def append_data(self, time, comm, pid, cpu, line, data):
        """Append data parsed from a line to the corresponding arrays

        The :mod:`DataFrame` will be created from this when the whole trace
        has been parsed.

        :param time: The time for the line that was printed in the trace
        :type time: float

        :param comm: The command name or the execname from which the trace
            line originated
        :type comm: str

        :param pid: The PID of the process from which the trace
            line originated
        :type pid: int

        :param data: The data for matching line in the trace
        :type data: str
        """

        self.time_array.append(time)
        self.comm_array.append(comm)
        self.pid_array.append(pid)
        self.cpu_array.append(cpu)
        self.line_array.append(line)
        self.data_array.append(data)

    def string_cast_int(self, string):
        """
        Attempt to convert string to an int

        :param string: The value to convert.
        :type string: str
        """

        try:
            return int(string)
        except ValueError:
            try:
                return int(string, base=16)
            except ValueError:
                return string

    def generate_data_dict(self, data_str):
        data_dict = {}
        prev_key = None
        for field in data_str.split():
            if "=" not in field:
                if not prev_key:
                    if 'FAILED TO PARSE' in data_str:
                        warnings.warn(
                            'trace-cmd failed to parse the "{}" event. You may '
                            'need to compile the latest trace-cmd and put it in '
                            'your $PATH. Continuing, but some data may be missing'
                            .format(self.unique_word))
                        continue
                    else:
                        raise TrappyParseError(
                            "TRAPpy's parser for '{}' failed to parse the line:"
                            "\n{}".format(self.unique_word, data_str))
                # Concatenation is supported only for "string" values
                if not isinstance(data_dict[prev_key], basestring):
                    continue
                data_dict[prev_key] += ' ' + field
                continue
            (key, value) = field.split('=', 1)
            value = self.string_cast_int(value)
            data_dict[key] = value
            prev_key = key
        return data_dict

    def generate_parsed_data(self):

        # Get a rough idea of how much memory we have to play with
        CHECK_MEM_COUNT = 10000
        kb_free = _get_free_memory_kb()
        starting_maxrss = getrusage(RUSAGE_SELF).ru_maxrss
        check_memory_usage = True
        check_memory_count = 1

        for (comm, pid, cpu, line, data_str) in zip(self.comm_array, self.pid_array,
                                              self.cpu_array, self.line_array,
                                              self.data_array):
            data_dict = {"__comm": comm, "__pid": pid, "__cpu": cpu, "__line": line}
            data_dict.update(self.generate_data_dict(data_str))

            # When running out of memory, Pandas has been observed to segfault
            # rather than throwing a proper Python error.
            # Look at how much memory our process is using and warn if we seem
            # to be getting close to the system's limit, check it only once
            # in the beginning and then every CHECK_MEM_COUNT events
            check_memory_count -= 1
            if check_memory_usage and check_memory_count == 0:
                kb_used = (getrusage(RUSAGE_SELF).ru_maxrss - starting_maxrss)
                if kb_free and kb_used > kb_free * 0.9:
                    warnings.warn("TRAPpy: Appear to be low on memory. "
                                  "If errors arise, try providing more RAM")
                    check_memory_usage = False
                check_memory_count = CHECK_MEM_COUNT

            yield data_dict

    def create_dataframe(self):
        """Create the final :mod:`pandas.DataFrame`"""
        if not self.time_array:
            return

        trace_arr_lengths = self.__get_trace_array_lengths()

        if trace_arr_lengths:
            for (idx, val) in enumerate(self.data_array):
                expl_val = trace_parser_explode_array(val, trace_arr_lengths)
                self.data_array[idx] = expl_val

        time_idx = pd.Index(self.time_array, name="Time")
        self.data_frame = pd.DataFrame(self.generate_parsed_data(), index=time_idx)

        self.time_array = []
        self.line_array = []
        self.comm_array = []
        self.pid_array = []
        self.cpu_array = []
        self.data_array = []

    def write_csv(self, fname):
        """Write the csv info into a CSV file

        :param fname: The name of the CSV file
        :type fname: str
        """
        self.data_frame.to_csv(fname)

    def read_csv(self, fname):
        """Read the csv data into a DataFrame

        :param fname: The name of the CSV file
        :type fname: str
        """
        self.data_frame = pd.read_csv(
            fname,
            index_col=0,
            # This ensures cached vs parsed timestamps are converted using the
            # same method, aka python's float() and not numpy's
            converters={'Time' : float}
        )

    def normalize_time(self, basetime):
        """Substract basetime from the Time of the data frame

        :param basetime: The offset which needs to be subtracted from
            the time index
        :type basetime: float
        """
        if basetime and not self.data_frame.empty:
            self.data_frame.reset_index(inplace=True)
            self.data_frame["Time"] = self.data_frame["Time"] - basetime
            self.data_frame.set_index("Time", inplace=True)
<|endoftext|>


<commit_before># -*- coding: utf-8 -*-
from __future__ import absolute_import, division, print_function, unicode_literals

import functools
import json
import os
import shlex
import sys

from argparse import ArgumentParser
from itertools import product
from collections import namedtuple
from os import getcwd
from os.path import abspath, dirname, isdir, isfile, join
from subprocess import CalledProcessError, check_call

from termcolor import colored
from yaml import safe_load

__version__ = '0.0.5'

def log(message=''):
	print(message)

def log_command(command):
	log('$ %s' % (colored(command, attrs=['bold']),))

def log_error(error):
	log(colored('%s' % (error,), 'red'))

def as_tuple(obj):
	return tuple(obj) if isinstance(obj, (list, tuple)) else (obj,)

try:
	unicode = unicode
	def native_str(s):
		if isinstance(s, unicode):
			s = s.encode('utf-8')
		assert isinstance(s, str)
		return s
except NameError:
	def native_str(s):
		assert isinstance(s, str)
		return s

def native_str_result(function):
	@functools.wraps(function)
	def wrapper(*args, **kwargs):
		result = function(*args, **kwargs)
		return native_str(result)
	return wrapper

class Structure(object):
	fields = ()

	def __hash__(self):
		return hash([getattr(self, str(f)) for f in self.fields])

	def __eq__(self, other):
		return type(self) == type(other) and \
			[getattr(self, str(f)) for f in self.fields] == [getattr(other, str(f)) for f in self.fields]

	def __repr__(self):
		return '%s(%s)' % (
			self.__class__.__name__,
			', '.join(('%s=%r' % (key, getattr(self, str(key))) for key in self.fields))
		)

	def __unicode__(self):
		return self.__str__().decode('utf-8')

class Step(Structure):
	fields = ('name', 'commands', 'can_fail',)

	def __init__(self, name, commands, can_fail=False, check_call=check_call):
		self.name = name
		self.commands = commands
		self.can_fail = can_fail
		self.check_call = check_call

	def perform(self):
		try:
			for command in self.commands:
				self.execute(command)
		except Exception as e:
			log_error('Error performing %r step' % (self,))
			if not self.can_fail:
				raise

	def execute(self, command):
		if command.startswith('sudo'):
			log(colored(
				'%r ignored because it contains sudo reference' % (command,), 'yellow'))
		else:
			log_command(command)
			self.check_call(command, shell=True)

	@native_str_result
	def __str__(self):
		return self.name

class Build(Structure):
	fields = ('steps',)

	def __init__(self, steps):
		self.steps = steps

	def run(self):
		for step in self.steps:
			step.perform()

class Configuration(Structure):
	fields = ('python', 'variables', 'can_fail', 'recreate')

	def __init__(self,
			python, variables, base_path='.travis-solo', can_fail=False, recreate=False,
			check_call=check_call, isdir=isdir, environ=os.environ):
		self.base_path = base_path
		self.python = python
		self.variables = variables
		self.can_fail = can_fail
		self.recreate = recreate
		self.check_call = check_call
		self.isdir = isdir
		self.environ = environ

	@property
	def virtualenv_path(self):
		return join(self.base_path, self.python)

	@property
	def full_python(self):
		return ('python' if not self.python.startswith('py') else '') + self.python

	def run_build(self, build):
		original_environ = self.environ.copy()
		try:
			log(colored('\n\nBuild configuration %s running' % (self,), attrs=['bold']))
			log('Preparing the environment')
			self.prepare_virtualenv()
			self.prepare_environment()
			build.run()
		except Exception as e:
			log_error(e)
			raise
		finally:
			self.environ.clear()
			self.environ.update(original_environ)

	def prepare_virtualenv(self):
		try:
			command = 'virtualenv --distribute --python=%s %s' % (
				self.full_python, self.virtualenv_path)
			log_command(command)
			self.check_call(command, shell=True)
		except OSError as e:
			log_error(e)
			raise Exception('No virtualenv executable found, please install virtualenv')
		except CalledProcessError as e:
			if e.returncode == 3:
				log_error(e)
				raise Exception('Interpreter not found')
			else:
				raise

	def prepare_environment(self):
		self.environ['PATH'] = ':'.join((join(self.virtualenv_path, 'bin'), self.environ['PATH']))

		try:
			del self.environ['__PYVENV_LAUNCHER__']
		except KeyError:
			pass

		for name in ('CI', 'TRAVIS', 'TRAVIS_SOLO'):
			self.environ[name] = 'true'

		self.environ['TRAVIS_PYTHON_VERSION'] = self.python

		self.environ.update(self.variables)

	@native_str_result
	def __str__(self):
		envs = ', '.join(('%s=%r' % (k, v) for (k, v) in self.variables.items()))
		return self.full_python + (' (%s)' % (envs,) if envs else '')

class Runner(Structure):
	fields = ('build', 'configurations')

	def __init__(self, build, configurations):
		self.build = build
		self.configurations = configurations

	def run(self):
		results = []
		for c in self.configurations:
			try:
				c.run_build(self.build)
				results.append((c, True, 'Build succeeded'))
			except Exception as e:
				results.append((c, False, e))

		log(colored('\n\nBuild summary:', attrs=['bold']))
		for conf, result, message in results:
			color = 'green' if result else 'red'
			log(colored('%s: %s' % (conf, message), color))

		success = all(result for (configuration, result, _) in results if not configuration.can_fail)
		return 0 if success else 1

class Loader(object):
	def load_steps(self, settings):
		lifecycle = (
			('before_install', False),
			('install', False),
			('before_script', False),
			('script', False),
			('after_script', True),
		)

		steps = []
		for name, can_fail in lifecycle:
			commands = as_tuple(settings.get(name, []))
			if commands:
				steps.append(Step(
					name=name,
					commands=commands,
					can_fail=can_fail
				))

		return tuple(steps)

	def load_configurations(self, settings):
		assert settings['language'] == 'python', 'Only Python projects are supported right now'

		versions = as_tuple(settings.get('python', '2.7'))
		env_sets = [self.parse_env_set(es) for es in as_tuple(settings.get('env', ''))]

		build_matrix = list(product(versions, env_sets))

		matrix = settings.get('matrix', {})
		include = matrix.get('include', [])
		for i in include:
			version = i['python']
			env_set = self.parse_env_set(i.get('env', ''))
			element = (version, env_set)
			if element not in build_matrix:
				build_matrix.append(element)

		exclude = matrix.get('exclude', [])
		for e in exclude:
			version = e['python']
			env_set = self.parse_env_set(e.get('env', ''))
			element = (version, env_set)
			if element in build_matrix:
				build_matrix.remove(element)

		configurations = tuple((
			Configuration(python=p, variables=dict(v), base_path=join(getcwd(), '.travis-solo'))
			for (p, v) in build_matrix
		))
		allow_failures = matrix.get('allow_failures', [])
		for af in allow_failures:
			python = af.get('python')
			variables = dict(self.parse_env_set(af.get('env', '')))
			for c in configurations:
				if c.python == python and c.variables == variables:
					c.can_fail = True

		return configurations

	def parse_env_set(self, env_set):
		return tuple(tuple(e.strip().split('=')) for e in env_set.split())

class Application(object):
	def __init__(self, getcwd=getcwd, isfile=isfile, open=open):
		self.getcwd = getcwd
		self.isfile = isfile
		self.open = open

	def run(self, argv):
		args = self.get_args(argv)
		settings = self.get_settings(args)

		loader = Loader()
		steps = loader.load_steps(settings)
		configurations = loader.load_configurations(settings)

		build = Build(steps)
		runner = Runner(build, configurations)
		sys.exit(runner.run())

	def get_args(self, argv):
		parser = ArgumentParser(description='Local Travis build runner')
		parser.add_argument(
			'--overwrite', dest='overwrite', action='store',
			help='Overwrite settings loaded from file with JSON-encoded dict. Usage:\n'
			     '''--overwrite '{"python": "2.7", "env": ["A=a", "A=b"]}' ''')
		parser.add_argument('--version', action='version', version=__version__)

		return parser.parse_args(argv[1:])

	def get_settings(self, args):
		travis_file = join(self.getcwd(), '.travis.yml')
		assert self.isfile(travis_file), 'File %s not found' % (travis_file,)

		with self.open(travis_file) as fd:
			settings = safe_load(fd)

		overwrite = args.overwrite or '{}'
		decoded = json.loads(overwrite)
		settings.update(decoded)

		return settings

def main():
	app = Application()
	app.run(sys.argv)

if __name__ == '__main__':
	main()
<commit_msg>Bump version<commit_after># -*- coding: utf-8 -*-
from __future__ import absolute_import, division, print_function, unicode_literals

import functools
import json
import os
import shlex
import sys

from argparse import ArgumentParser
from itertools import product
from collections import namedtuple
from os import getcwd
from os.path import abspath, dirname, isdir, isfile, join
from subprocess import CalledProcessError, check_call

from termcolor import colored
from yaml import safe_load

__version__ = '0.0.6'

def log(message=''):
	print(message)

def log_command(command):
	log('$ %s' % (colored(command, attrs=['bold']),))

def log_error(error):
	log(colored('%s' % (error,), 'red'))

def as_tuple(obj):
	return tuple(obj) if isinstance(obj, (list, tuple)) else (obj,)

try:
	unicode = unicode
	def native_str(s):
		if isinstance(s, unicode):
			s = s.encode('utf-8')
		assert isinstance(s, str)
		return s
except NameError:
	def native_str(s):
		assert isinstance(s, str)
		return s

def native_str_result(function):
	@functools.wraps(function)
	def wrapper(*args, **kwargs):
		result = function(*args, **kwargs)
		return native_str(result)
	return wrapper

class Structure(object):
	fields = ()

	def __hash__(self):
		return hash([getattr(self, str(f)) for f in self.fields])

	def __eq__(self, other):
		return type(self) == type(other) and \
			[getattr(self, str(f)) for f in self.fields] == [getattr(other, str(f)) for f in self.fields]

	def __repr__(self):
		return '%s(%s)' % (
			self.__class__.__name__,
			', '.join(('%s=%r' % (key, getattr(self, str(key))) for key in self.fields))
		)

	def __unicode__(self):
		return self.__str__().decode('utf-8')

class Step(Structure):
	fields = ('name', 'commands', 'can_fail',)

	def __init__(self, name, commands, can_fail=False, check_call=check_call):
		self.name = name
		self.commands = commands
		self.can_fail = can_fail
		self.check_call = check_call

	def perform(self):
		try:
			for command in self.commands:
				self.execute(command)
		except Exception as e:
			log_error('Error performing %r step' % (self,))
			if not self.can_fail:
				raise

	def execute(self, command):
		if command.startswith('sudo'):
			log(colored(
				'%r ignored because it contains sudo reference' % (command,), 'yellow'))
		else:
			log_command(command)
			self.check_call(command, shell=True)

	@native_str_result
	def __str__(self):
		return self.name

class Build(Structure):
	fields = ('steps',)

	def __init__(self, steps):
		self.steps = steps

	def run(self):
		for step in self.steps:
			step.perform()

class Configuration(Structure):
	fields = ('python', 'variables', 'can_fail', 'recreate')

	def __init__(self,
			python, variables, base_path='.travis-solo', can_fail=False, recreate=False,
			check_call=check_call, isdir=isdir, environ=os.environ):
		self.base_path = base_path
		self.python = python
		self.variables = variables
		self.can_fail = can_fail
		self.recreate = recreate
		self.check_call = check_call
		self.isdir = isdir
		self.environ = environ

	@property
	def virtualenv_path(self):
		return join(self.base_path, self.python)

	@property
	def full_python(self):
		return ('python' if not self.python.startswith('py') else '') + self.python

	def run_build(self, build):
		original_environ = self.environ.copy()
		try:
			log(colored('\n\nBuild configuration %s running' % (self,), attrs=['bold']))
			log('Preparing the environment')
			self.prepare_virtualenv()
			self.prepare_environment()
			build.run()
		except Exception as e:
			log_error(e)
			raise
		finally:
			self.environ.clear()
			self.environ.update(original_environ)

	def prepare_virtualenv(self):
		try:
			command = 'virtualenv --distribute --python=%s %s' % (
				self.full_python, self.virtualenv_path)
			log_command(command)
			self.check_call(command, shell=True)
		except OSError as e:
			log_error(e)
			raise Exception('No virtualenv executable found, please install virtualenv')
		except CalledProcessError as e:
			if e.returncode == 3:
				log_error(e)
				raise Exception('Interpreter not found')
			else:
				raise

	def prepare_environment(self):
		self.environ['PATH'] = ':'.join((join(self.virtualenv_path, 'bin'), self.environ['PATH']))

		try:
			del self.environ['__PYVENV_LAUNCHER__']
		except KeyError:
			pass

		for name in ('CI', 'TRAVIS', 'TRAVIS_SOLO'):
			self.environ[name] = 'true'

		self.environ['TRAVIS_PYTHON_VERSION'] = self.python

		self.environ.update(self.variables)

	@native_str_result
	def __str__(self):
		envs = ', '.join(('%s=%r' % (k, v) for (k, v) in self.variables.items()))
		return self.full_python + (' (%s)' % (envs,) if envs else '')

class Runner(Structure):
	fields = ('build', 'configurations')

	def __init__(self, build, configurations):
		self.build = build
		self.configurations = configurations

	def run(self):
		results = []
		for c in self.configurations:
			try:
				c.run_build(self.build)
				results.append((c, True, 'Build succeeded'))
			except Exception as e:
				results.append((c, False, e))

		log(colored('\n\nBuild summary:', attrs=['bold']))
		for conf, result, message in results:
			color = 'green' if result else 'red'
			log(colored('%s: %s' % (conf, message), color))

		success = all(result for (configuration, result, _) in results if not configuration.can_fail)
		return 0 if success else 1

class Loader(object):
	def load_steps(self, settings):
		lifecycle = (
			('before_install', False),
			('install', False),
			('before_script', False),
			('script', False),
			('after_script', True),
		)

		steps = []
		for name, can_fail in lifecycle:
			commands = as_tuple(settings.get(name, []))
			if commands:
				steps.append(Step(
					name=name,
					commands=commands,
					can_fail=can_fail
				))

		return tuple(steps)

	def load_configurations(self, settings):
		assert settings['language'] == 'python', 'Only Python projects are supported right now'

		versions = as_tuple(settings.get('python', '2.7'))
		env_sets = [self.parse_env_set(es) for es in as_tuple(settings.get('env', ''))]

		build_matrix = list(product(versions, env_sets))

		matrix = settings.get('matrix', {})
		include = matrix.get('include', [])
		for i in include:
			version = i['python']
			env_set = self.parse_env_set(i.get('env', ''))
			element = (version, env_set)
			if element not in build_matrix:
				build_matrix.append(element)

		exclude = matrix.get('exclude', [])
		for e in exclude:
			version = e['python']
			env_set = self.parse_env_set(e.get('env', ''))
			element = (version, env_set)
			if element in build_matrix:
				build_matrix.remove(element)

		configurations = tuple((
			Configuration(python=p, variables=dict(v), base_path=join(getcwd(), '.travis-solo'))
			for (p, v) in build_matrix
		))
		allow_failures = matrix.get('allow_failures', [])
		for af in allow_failures:
			python = af.get('python')
			variables = dict(self.parse_env_set(af.get('env', '')))
			for c in configurations:
				if c.python == python and c.variables == variables:
					c.can_fail = True

		return configurations

	def parse_env_set(self, env_set):
		return tuple(tuple(e.strip().split('=')) for e in env_set.split())

class Application(object):
	def __init__(self, getcwd=getcwd, isfile=isfile, open=open):
		self.getcwd = getcwd
		self.isfile = isfile
		self.open = open

	def run(self, argv):
		args = self.get_args(argv)
		settings = self.get_settings(args)

		loader = Loader()
		steps = loader.load_steps(settings)
		configurations = loader.load_configurations(settings)

		build = Build(steps)
		runner = Runner(build, configurations)
		sys.exit(runner.run())

	def get_args(self, argv):
		parser = ArgumentParser(description='Local Travis build runner')
		parser.add_argument(
			'--overwrite', dest='overwrite', action='store',
			help='Overwrite settings loaded from file with JSON-encoded dict. Usage:\n'
			     '''--overwrite '{"python": "2.7", "env": ["A=a", "A=b"]}' ''')
		parser.add_argument('--version', action='version', version=__version__)

		return parser.parse_args(argv[1:])

	def get_settings(self, args):
		travis_file = join(self.getcwd(), '.travis.yml')
		assert self.isfile(travis_file), 'File %s not found' % (travis_file,)

		with self.open(travis_file) as fd:
			settings = safe_load(fd)

		overwrite = args.overwrite or '{}'
		decoded = json.loads(overwrite)
		settings.update(decoded)

		return settings

def main():
	app = Application()
	app.run(sys.argv)

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


<commit_before>import _curses, curses
from bisect import bisect_left
import math
import os
from threading import Thread
from time import sleep

from tsparser.utils import Singleton, StatisticDataCollector

class UserInterface(metaclass=Singleton):
    """
    User Interface is a singleton. Once ran, it renders UI until exiting the app.
    """

    def __init__(self, refreshing_frequency=30):
        self.__REFRESHING_FREQUENCY = refreshing_frequency

    def run(self):
        Thread(target=self.__interface_thread, daemon=True).start()

    def __interface_thread(self):
        try:
            self.__init_curses()
            while True:
                sleep(1/self.__REFRESHING_FREQUENCY)
                self.__update_filter()
                self.__process_events()
                self.__render_frame()
        except Exception as err:
            self.__screen.clear()
            error_message = '{}: {}'.format(err.__class__.__name__, err)
            self.__screen.addstr(error_message)
            StatisticDataCollector().get_logger().log('ui', error_message)
            self.__screen.refresh()

    def __init_curses(self):
        self.__screen = curses.initscr()
        curses.start_color()
        curses.curs_set(0)
        self.__screen.nodelay(True)
        self.__screen.keypad(True)
        self.__SCREEN_MINIMAL_SIZE = 24, 80  # lines, cols

        self.__TIMESTAMP_COLOR = 1
        curses.init_pair(self.__TIMESTAMP_COLOR, curses.COLOR_RED, curses.COLOR_BLACK)
        self.__MODULE_NAME_COLOR = 2
        curses.init_pair(self.__MODULE_NAME_COLOR, curses.COLOR_BLUE, curses.COLOR_BLACK)
        self.__INFO_BAR_DESC_COLOR = 3
        curses.init_pair(self.__INFO_BAR_DESC_COLOR, curses.COLOR_BLACK, curses.COLOR_CYAN)
        self.__FILTER_WINDOW_BACKGROUND = 4
        curses.init_pair(self.__FILTER_WINDOW_BACKGROUND, curses.COLOR_WHITE, curses.COLOR_BLUE)
        self.__FILTER_WINDOW_SELECTION = 5
        curses.init_pair(self.__FILTER_WINDOW_SELECTION, curses.COLOR_BLACK, curses.COLOR_CYAN)

        self.__logs_auto_scrolling = True
        self.__scroll_position = 0
        self.__log_index_to_last_line_no = list()
        self.__cached_processed_logs = list()
        self.__filter_window_active = False
        self.__filter = dict()
        self.__filter_selected_index = int()
        self.__filter_selected_module = str()
        StatisticDataCollector().get_logger().log('ui', 'User interface initialized!')

    def __update_filter(self):
        for module_name in StatisticDataCollector().get_logger().get_all_modules():
            if module_name not in self.__filter:
                self.__filter[module_name] = True

    def __process_events(self):
        while True:
            key_code = self.__screen.getch()
            if key_code == curses.ERR:
                return
            if key_code == curses.KEY_RESIZE:
                self.__delete_cached_logs()
                continue

            if self.__filter_window_active:
                self.__filter_window_process_event(key_code)
            else:
                self.__main_window_process_event(key_code)

    def __delete_cached_logs(self):
        self.__cached_processed_logs.clear()
        self.__log_index_to_last_line_no.clear()

    def __filter_window_process_event(self, key_code):
        if key_code == 27:  # escape
            self.__filter_window_active = False
        elif key_code == ord(' '):
            if self.__filter:
                self.__filter[self.__filter_selected_module] = not self.__filter[self.__filter_selected_module]
                self.__auto_scroll_position = True
                self.__delete_cached_logs()
        elif key_code == curses.KEY_UP:
            self.__filter_selected_index -= 1
            if self.__filter_selected_index == -1:
                self.__filter_selected_index = len(self.__filter) - 1
        elif key_code == curses.KEY_DOWN:
            self.__filter_selected_index += 1
            if self.__filter_selected_index == len(self.__filter):
                self.__filter_selected_index = 0

    def __main_window_process_event(self, key_code):
        if key_code == curses.KEY_F2:
            self.__logs_auto_scrolling = not self.__logs_auto_scrolling
        elif key_code == curses.KEY_F3:
            self.__filter_window_active = True
            self.__filter_selected_index = 0
        elif key_code == curses.KEY_F4:
            StatisticDataCollector().get_logger().clear_logs()
            self.__delete_cached_logs()
        elif key_code == curses.KEY_F9:
            curses.endwin()
            os.kill(os.getpid(), 15)
        elif key_code == curses.KEY_UP:
            self.__logs_auto_scrolling = False
            self.__scroll_position -= 1  # renderer will increase value if it is too small
        elif key_code == curses.KEY_DOWN:
            if self.__cached_processed_logs:
                if self.__scroll_position >= self.__log_index_to_last_line_no[-1]:
                    self.__logs_auto_scrolling = True
                else:
                    self.__scroll_position += 1

    def __render_frame(self):
        lines, cols = self.__screen.getmaxyx()
        min_lines, min_cols = self.__SCREEN_MINIMAL_SIZE
        if lines < min_lines or cols < min_cols:
            self.__screen.clear()
            self.__screen.addstr('Terminal size should be at least {}x{}!\n'.format(min_cols, min_lines))
            self.__screen.refresh()
            return
        statistics_window_width = 40

        logs_windows = self.__screen.subwin(lines - 1, cols - statistics_window_width, 0, 0)
        self.__render_logs_window(logs_windows)
        statistics_window = self.__screen.subwin(lines - 1, statistics_window_width, 0, cols - statistics_window_width)
        self.__render_statistics_window(statistics_window)
        info_bar_window = self.__screen.subwin(1, cols, lines - 1, 0)
        self.__render_info_bar(info_bar_window)
        if self.__filter_window_active:
            width, height = 60, 20
            filter_window = self.__screen.subwin(height, width, (lines - height) // 2, (cols - width) // 2)
            self.__render_filter_window(filter_window)
        self.__screen.refresh()

    def __render_logs_window(self, window):
        window.clear()
        self.__draw_entitled_box(window, 'Logs')

        sub_win = self.__get_sub_window(window)
        lines, cols = sub_win.getmaxyx()
        selected_modules = [module_name for module_name in self.__filter if self.__filter[module_name]]
        logs = StatisticDataCollector().get_logger().get_logs(selected_modules)
        new_logs = logs[len(self.__cached_processed_logs):]
        self.__cache_new_log_entries(new_logs, cols)
        self.__render_visible_log_entries(sub_win)

    def __cache_new_log_entries(self, new_entries, line_width):
        for timestamp, module_name, message in new_entries:
            timestamp_str = '{:02}:{:02}:{:02}.{:06}'.format(timestamp.hour, timestamp.minute,
                                                             timestamp.second, timestamp.microsecond)
            module_name = module_name.replace('\n', '<nl> ')
            self.__cached_processed_logs.append((timestamp_str, module_name, message))
            whole_message = timestamp_str + ' ' + module_name + ' ' + message
            lines_needed = 0
            for pseudo_line in whole_message.split('\n')[:-1]:
                pseudo_line += '\n'
                lines_needed += math.ceil(len(pseudo_line) / line_width)
            lines_needed += math.ceil(len(whole_message.split('\n')[-1]) / line_width)
            previous_entry_last_line = self.__log_index_to_last_line_no[-1] if self.__log_index_to_last_line_no else -1
            self.__log_index_to_last_line_no.append(previous_entry_last_line + lines_needed)

    def __render_visible_log_entries(self, window):
        lines, cols = window.getmaxyx()
        lines -= 1  # last line should be empty (cursor will be there)
                    # otherwise cursor will land below the window (what causes curses error)
        if not self.__cached_processed_logs:
            window.addstr(lines, 0, '(no logs)', curses.A_DIM)
            return
        if self.__logs_auto_scrolling:
            self.__scroll_position = self.__log_index_to_last_line_no[-1]
        self.__scroll_position = max(self.__scroll_position, lines - 1)
        first_line_no = max(0, self.__scroll_position - lines + 1)
        first_log_entry_index = bisect_left(self.__log_index_to_last_line_no, first_line_no)
        log_entry_index = first_log_entry_index
        while log_entry_index < len(self.__cached_processed_logs):
            last_entry_line_no = self.__log_index_to_last_line_no[log_entry_index]
            first_entry_line_no = self.__log_index_to_last_line_no[log_entry_index-1] + 1 if log_entry_index > 0 else 0
            self.__render_log_entry(window, log_entry_index,
                                    max(0, first_line_no - first_entry_line_no),
                                    max(0, last_entry_line_no - self.__scroll_position))

            log_entry_index += 1
            if last_entry_line_no >= self.__scroll_position:
                break

        if self.__scroll_position >= self.__log_index_to_last_line_no[-1]:
            info_message = '(end)'
        else:
            left_lines_count = self.__log_index_to_last_line_no[-1] - self.__scroll_position
            left_lines_count = str(left_lines_count) if left_lines_count < 10**9 else '>=10e9'
            info_message = '({} more lines)'.format(left_lines_count)
        window.addstr(lines, 0, info_message, curses.A_DIM)

    def __render_log_entry(self, window, log_entry_index, omitted_first_lines, omitted_last_lines):
        lines, cols = window.getmaxyx()
        timestamp_str, module_name, message = self.__cached_processed_logs[log_entry_index]
        colored_prefix = timestamp_str + ' ' + module_name + ' '
        whole_message = colored_prefix + message
        entry_lines = list()
        while len(whole_message) > 0:
            next_new_line = whole_message.find('\n') + 1
            if next_new_line == 0:
                next_new_line = cols
            split_point = min(cols, len(whole_message), next_new_line)
            entry_lines.append(whole_message[:split_point])
            whole_message = whole_message[split_point:]
            if len(entry_lines[-1]) < cols and not entry_lines[-1].endswith('\n'):
                    entry_lines[-1] += '\n'
        if omitted_first_lines + omitted_last_lines >= len(entry_lines):
            return

        if omitted_first_lines == 0:
            if len(colored_prefix) <= cols:
                rest_of_first_line = entry_lines[0][len(colored_prefix):]
                window.addstr(timestamp_str + ' ', curses.color_pair(self.__TIMESTAMP_COLOR))
                window.addstr(module_name + ' ', curses.color_pair(self.__MODULE_NAME_COLOR))
                window.addstr(rest_of_first_line)
            else:
                window.addstr(entry_lines[0])
        entry_lines = entry_lines[max(1, omitted_first_lines):len(entry_lines)-omitted_last_lines]
        for line in entry_lines:
            window.addstr(line)

    def __render_statistics_window(self, window):
        window.clear()
        self.__draw_entitled_box(window, 'Statistics')
        sub_win = self.__get_sub_window(window)
        stats_to_display = self.__prepare_stats()
        for name, value in stats_to_display:
            sub_win.addstr('{}: {}\n'.format(name, value))

    @staticmethod
    def __prepare_stats():
        def timedelta_to_str(timedelta_obj):
            seconds = timedelta_obj.total_seconds()
            if seconds < 1:
                return '<1 sec'
            seconds = int(seconds)
            result = '{} sec'.format(seconds % 60)
            if seconds >= 60:
                result = '{} min {}'.format(seconds // 60, result)
            return result

        def data_amount_to_str(data_amount):
            steps = (
                (1, 'B'),
                (10**3, 'kB'),
                (10**6, 'MB'),
                (10**9, 'GB')
            )
            last_good_result = '0 B'
            for bound, unit in steps:
                if data_amount < bound:
                    return last_good_result
                last_good_result = '{:.3f} {}'.format(data_amount / bound, unit)
            return last_good_result

        sdc = StatisticDataCollector()
        stats_scheme = (
            ('Time since last receiving', timedelta_to_str(sdc.get_time_since_last_data_receiving())),
            ('Time since start', timedelta_to_str(sdc.get_time_since_start())),
            ('Data receiving speed', data_amount_to_str(sdc.get_data_receiving_speed())+'/s'),
            ('Total data received', data_amount_to_str(sdc.get_total_data_received())),
            ('Queued requests', str(sdc.get_count_of_queued_requests())),
            ('Total sent requests', str(sdc.get_total_count_of_sent_requests()))
        )
        return stats_scheme

    @staticmethod
    def __draw_entitled_box(window, title):
        window.box()
        window.addstr(0, 1, title)

    @staticmethod
    def __get_sub_window(window, margin=1):
        lines, cols = window.getmaxyx()
        beg_y, beg_x = window.getbegyx()
        return window.subwin(lines - margin * 2, cols - margin * 2, beg_y + margin, beg_x + margin)

    def __render_info_bar(self, window):
        window.clear()
        info_bar_scheme = (
            ('F2', '{} auto scrolling'.format('Disable' if self.__logs_auto_scrolling else 'Enable')),
            ('F3', 'Filter'),
            ('F4', 'Clear'),
            ('F9', 'Exit'),
            ('↑↓', 'Scroll')
        )
        for key, description in info_bar_scheme:
            window.addstr(key)
            window.addstr(description, curses.color_pair(self.__INFO_BAR_DESC_COLOR))

    def __render_filter_window(self, window):
        window.clear()
        window.bkgd(' ', curses.color_pair(self.__FILTER_WINDOW_BACKGROUND))
        self.__draw_entitled_box(window, 'Filter')
        sub_win = self.__get_sub_window(window)
        sub_win.addstr('Please use arrows, space and escape to navigate.\n\n')
        self.__render_filter_list(sub_win)

    def __render_filter_list(self, window):
        lines, cols = window.getmaxyx()
        loop_cnt = 0
        for module_name, is_checked in self.__filter.items():
            color = self.__FILTER_WINDOW_BACKGROUND
            is_entry_selected = self.__filter_selected_index == loop_cnt
            if is_entry_selected:
                color = self.__FILTER_WINDOW_SELECTION
                self.__filter_selected_module = module_name

            prefix = '[x] ' if is_checked else '[ ] '
            module_name_needed_length = cols - len(prefix)
            if len(module_name) > module_name_needed_length:
                module_name = module_name[:module_name_needed_length-3] + '...'
            else:
                module_name += ' ' * (module_name_needed_length - len(module_name))
            try:  # TODO if there are many modules (False on May 18th, 2015), implement scrolling for this window
                window.addstr(prefix + module_name, curses.color_pair(color))
            except _curses.error:
                break
            loop_cnt += 1
<commit_msg>On UI crash terminal state is correctly restored<commit_after>import _curses, curses
from bisect import bisect_left
import math
import os
from threading import Thread
from time import sleep

from tsparser.utils import Singleton, StatisticDataCollector

class UserInterface(metaclass=Singleton):
    """
    User Interface is a singleton. Once ran, it renders UI until exiting the app.
    """

    def __init__(self, refreshing_frequency=30):
        self.__REFRESHING_FREQUENCY = refreshing_frequency

    def run(self):
        Thread(target=self.__interface_thread, daemon=True).start()

    def __interface_thread(self):
        try:
            self.__init_curses()
            while True:
                sleep(1/self.__REFRESHING_FREQUENCY)
                self.__update_filter()
                self.__process_events()
                self.__render_frame()
        except Exception as err:
            error_message = '{}: {}'.format(err.__class__.__name__, err)
            StatisticDataCollector().get_logger().log('ui', error_message)
            curses.endwin()
            print(error_message)

    def __init_curses(self):
        self.__screen = curses.initscr()
        curses.start_color()
        curses.curs_set(0)
        self.__screen.nodelay(True)
        self.__screen.keypad(True)
        self.__SCREEN_MINIMAL_SIZE = 24, 80  # lines, cols

        self.__TIMESTAMP_COLOR = 1
        curses.init_pair(self.__TIMESTAMP_COLOR, curses.COLOR_RED, curses.COLOR_BLACK)
        self.__MODULE_NAME_COLOR = 2
        curses.init_pair(self.__MODULE_NAME_COLOR, curses.COLOR_BLUE, curses.COLOR_BLACK)
        self.__INFO_BAR_DESC_COLOR = 3
        curses.init_pair(self.__INFO_BAR_DESC_COLOR, curses.COLOR_BLACK, curses.COLOR_CYAN)
        self.__FILTER_WINDOW_BACKGROUND = 4
        curses.init_pair(self.__FILTER_WINDOW_BACKGROUND, curses.COLOR_WHITE, curses.COLOR_BLUE)
        self.__FILTER_WINDOW_SELECTION = 5
        curses.init_pair(self.__FILTER_WINDOW_SELECTION, curses.COLOR_BLACK, curses.COLOR_CYAN)

        self.__logs_auto_scrolling = True
        self.__scroll_position = 0
        self.__log_index_to_last_line_no = list()
        self.__cached_processed_logs = list()
        self.__filter_window_active = False
        self.__filter = dict()
        self.__filter_selected_index = int()
        self.__filter_selected_module = str()
        StatisticDataCollector().get_logger().log('ui', 'User interface initialized!')

    def __update_filter(self):
        for module_name in StatisticDataCollector().get_logger().get_all_modules():
            if module_name not in self.__filter:
                self.__filter[module_name] = True

    def __process_events(self):
        while True:
            key_code = self.__screen.getch()
            if key_code == curses.ERR:
                return
            if key_code == curses.KEY_RESIZE:
                self.__delete_cached_logs()
                continue

            if self.__filter_window_active:
                self.__filter_window_process_event(key_code)
            else:
                self.__main_window_process_event(key_code)

    def __delete_cached_logs(self):
        self.__cached_processed_logs.clear()
        self.__log_index_to_last_line_no.clear()

    def __filter_window_process_event(self, key_code):
        if key_code == 27:  # escape
            self.__filter_window_active = False
        elif key_code == ord(' '):
            if self.__filter:
                self.__filter[self.__filter_selected_module] = not self.__filter[self.__filter_selected_module]
                self.__auto_scroll_position = True
                self.__delete_cached_logs()
        elif key_code == curses.KEY_UP:
            self.__filter_selected_index -= 1
            if self.__filter_selected_index == -1:
                self.__filter_selected_index = len(self.__filter) - 1
        elif key_code == curses.KEY_DOWN:
            self.__filter_selected_index += 1
            if self.__filter_selected_index == len(self.__filter):
                self.__filter_selected_index = 0

    def __main_window_process_event(self, key_code):
        if key_code == curses.KEY_F2:
            self.__logs_auto_scrolling = not self.__logs_auto_scrolling
        elif key_code == curses.KEY_F3:
            self.__filter_window_active = True
            self.__filter_selected_index = 0
        elif key_code == curses.KEY_F4:
            StatisticDataCollector().get_logger().clear_logs()
            self.__delete_cached_logs()
        elif key_code == curses.KEY_F9:
            curses.endwin()
            os.kill(os.getpid(), 15)
        elif key_code == curses.KEY_UP:
            self.__logs_auto_scrolling = False
            self.__scroll_position -= 1  # renderer will increase value if it is too small
        elif key_code == curses.KEY_DOWN:
            if self.__cached_processed_logs:
                if self.__scroll_position >= self.__log_index_to_last_line_no[-1]:
                    self.__logs_auto_scrolling = True
                else:
                    self.__scroll_position += 1

    def __render_frame(self):
        lines, cols = self.__screen.getmaxyx()
        min_lines, min_cols = self.__SCREEN_MINIMAL_SIZE
        if lines < min_lines or cols < min_cols:
            self.__screen.clear()
            self.__screen.addstr('Terminal size should be at least {}x{}!\n'.format(min_cols, min_lines))
            self.__screen.refresh()
            return
        statistics_window_width = 40

        logs_windows = self.__screen.subwin(lines - 1, cols - statistics_window_width, 0, 0)
        self.__render_logs_window(logs_windows)
        statistics_window = self.__screen.subwin(lines - 1, statistics_window_width, 0, cols - statistics_window_width)
        self.__render_statistics_window(statistics_window)
        info_bar_window = self.__screen.subwin(1, cols, lines - 1, 0)
        self.__render_info_bar(info_bar_window)
        if self.__filter_window_active:
            width, height = 60, 20
            filter_window = self.__screen.subwin(height, width, (lines - height) // 2, (cols - width) // 2)
            self.__render_filter_window(filter_window)
        self.__screen.refresh()

    def __render_logs_window(self, window):
        window.clear()
        self.__draw_entitled_box(window, 'Logs')

        sub_win = self.__get_sub_window(window)
        lines, cols = sub_win.getmaxyx()
        selected_modules = [module_name for module_name in self.__filter if self.__filter[module_name]]
        logs = StatisticDataCollector().get_logger().get_logs(selected_modules)
        new_logs = logs[len(self.__cached_processed_logs):]
        self.__cache_new_log_entries(new_logs, cols)
        self.__render_visible_log_entries(sub_win)

    def __cache_new_log_entries(self, new_entries, line_width):
        for timestamp, module_name, message in new_entries:
            timestamp_str = '{:02}:{:02}:{:02}.{:06}'.format(timestamp.hour, timestamp.minute,
                                                             timestamp.second, timestamp.microsecond)
            module_name = module_name.replace('\n', '<nl> ')
            self.__cached_processed_logs.append((timestamp_str, module_name, message))
            whole_message = timestamp_str + ' ' + module_name + ' ' + message
            lines_needed = 0
            for pseudo_line in whole_message.split('\n')[:-1]:
                pseudo_line += '\n'
                lines_needed += math.ceil(len(pseudo_line) / line_width)
            lines_needed += math.ceil(len(whole_message.split('\n')[-1]) / line_width)
            previous_entry_last_line = self.__log_index_to_last_line_no[-1] if self.__log_index_to_last_line_no else -1
            self.__log_index_to_last_line_no.append(previous_entry_last_line + lines_needed)

    def __render_visible_log_entries(self, window):
        lines, cols = window.getmaxyx()
        lines -= 1  # last line should be empty (cursor will be there)
                    # otherwise cursor will land below the window (what causes curses error)
        if not self.__cached_processed_logs:
            window.addstr(lines, 0, '(no logs)', curses.A_DIM)
            return
        if self.__logs_auto_scrolling:
            self.__scroll_position = self.__log_index_to_last_line_no[-1]
        self.__scroll_position = max(self.__scroll_position, lines - 1)
        first_line_no = max(0, self.__scroll_position - lines + 1)
        first_log_entry_index = bisect_left(self.__log_index_to_last_line_no, first_line_no)
        log_entry_index = first_log_entry_index
        while log_entry_index < len(self.__cached_processed_logs):
            last_entry_line_no = self.__log_index_to_last_line_no[log_entry_index]
            first_entry_line_no = self.__log_index_to_last_line_no[log_entry_index-1] + 1 if log_entry_index > 0 else 0
            self.__render_log_entry(window, log_entry_index,
                                    max(0, first_line_no - first_entry_line_no),
                                    max(0, last_entry_line_no - self.__scroll_position))

            log_entry_index += 1
            if last_entry_line_no >= self.__scroll_position:
                break

        if self.__scroll_position >= self.__log_index_to_last_line_no[-1]:
            info_message = '(end)'
        else:
            left_lines_count = self.__log_index_to_last_line_no[-1] - self.__scroll_position
            left_lines_count = str(left_lines_count) if left_lines_count < 10**9 else '>=10e9'
            info_message = '({} more lines)'.format(left_lines_count)
        window.addstr(lines, 0, info_message, curses.A_DIM)

    def __render_log_entry(self, window, log_entry_index, omitted_first_lines, omitted_last_lines):
        lines, cols = window.getmaxyx()
        timestamp_str, module_name, message = self.__cached_processed_logs[log_entry_index]
        colored_prefix = timestamp_str + ' ' + module_name + ' '
        whole_message = colored_prefix + message
        entry_lines = list()
        while len(whole_message) > 0:
            next_new_line = whole_message.find('\n') + 1
            if next_new_line == 0:
                next_new_line = cols
            split_point = min(cols, len(whole_message), next_new_line)
            entry_lines.append(whole_message[:split_point])
            whole_message = whole_message[split_point:]
            if len(entry_lines[-1]) < cols and not entry_lines[-1].endswith('\n'):
                    entry_lines[-1] += '\n'
        if omitted_first_lines + omitted_last_lines >= len(entry_lines):
            return

        if omitted_first_lines == 0:
            if len(colored_prefix) <= cols:
                rest_of_first_line = entry_lines[0][len(colored_prefix):]
                window.addstr(timestamp_str + ' ', curses.color_pair(self.__TIMESTAMP_COLOR))
                window.addstr(module_name + ' ', curses.color_pair(self.__MODULE_NAME_COLOR))
                window.addstr(rest_of_first_line)
            else:
                window.addstr(entry_lines[0])
        entry_lines = entry_lines[max(1, omitted_first_lines):len(entry_lines)-omitted_last_lines]
        for line in entry_lines:
            window.addstr(line)

    def __render_statistics_window(self, window):
        window.clear()
        self.__draw_entitled_box(window, 'Statistics')
        sub_win = self.__get_sub_window(window)
        stats_to_display = self.__prepare_stats()
        for name, value in stats_to_display:
            sub_win.addstr('{}: {}\n'.format(name, value))

    @staticmethod
    def __prepare_stats():
        def timedelta_to_str(timedelta_obj):
            seconds = timedelta_obj.total_seconds()
            if seconds < 1:
                return '<1 sec'
            seconds = int(seconds)
            result = '{} sec'.format(seconds % 60)
            if seconds >= 60:
                result = '{} min {}'.format(seconds // 60, result)
            return result

        def data_amount_to_str(data_amount):
            steps = (
                (1, 'B'),
                (10**3, 'kB'),
                (10**6, 'MB'),
                (10**9, 'GB')
            )
            last_good_result = '0 B'
            for bound, unit in steps:
                if data_amount < bound:
                    return last_good_result
                last_good_result = '{:.3f} {}'.format(data_amount / bound, unit)
            return last_good_result

        sdc = StatisticDataCollector()
        stats_scheme = (
            ('Time since last receiving', timedelta_to_str(sdc.get_time_since_last_data_receiving())),
            ('Time since start', timedelta_to_str(sdc.get_time_since_start())),
            ('Data receiving speed', data_amount_to_str(sdc.get_data_receiving_speed())+'/s'),
            ('Total data received', data_amount_to_str(sdc.get_total_data_received())),
            ('Queued requests', str(sdc.get_count_of_queued_requests())),
            ('Total sent requests', str(sdc.get_total_count_of_sent_requests()))
        )
        return stats_scheme

    @staticmethod
    def __draw_entitled_box(window, title):
        window.box()
        window.addstr(0, 1, title)

    @staticmethod
    def __get_sub_window(window, margin=1):
        lines, cols = window.getmaxyx()
        beg_y, beg_x = window.getbegyx()
        return window.subwin(lines - margin * 2, cols - margin * 2, beg_y + margin, beg_x + margin)

    def __render_info_bar(self, window):
        window.clear()
        info_bar_scheme = (
            ('F2', '{} auto scrolling'.format('Disable' if self.__logs_auto_scrolling else 'Enable')),
            ('F3', 'Filter'),
            ('F4', 'Clear'),
            ('F9', 'Exit'),
            ('↑↓', 'Scroll')
        )
        for key, description in info_bar_scheme:
            window.addstr(key)
            window.addstr(description, curses.color_pair(self.__INFO_BAR_DESC_COLOR))

    def __render_filter_window(self, window):
        window.clear()
        window.bkgd(' ', curses.color_pair(self.__FILTER_WINDOW_BACKGROUND))
        self.__draw_entitled_box(window, 'Filter')
        sub_win = self.__get_sub_window(window)
        sub_win.addstr('Please use arrows, space and escape to navigate.\n\n')
        self.__render_filter_list(sub_win)

    def __render_filter_list(self, window):
        lines, cols = window.getmaxyx()
        loop_cnt = 0
        for module_name, is_checked in self.__filter.items():
            color = self.__FILTER_WINDOW_BACKGROUND
            is_entry_selected = self.__filter_selected_index == loop_cnt
            if is_entry_selected:
                color = self.__FILTER_WINDOW_SELECTION
                self.__filter_selected_module = module_name

            prefix = '[x] ' if is_checked else '[ ] '
            module_name_needed_length = cols - len(prefix)
            if len(module_name) > module_name_needed_length:
                module_name = module_name[:module_name_needed_length-3] + '...'
            else:
                module_name += ' ' * (module_name_needed_length - len(module_name))
            try:  # TODO if there are many modules (False on May 18th, 2015), implement scrolling for this window
                window.addstr(prefix + module_name, curses.color_pair(color))
            except _curses.error:
                break
            loop_cnt += 1
<|endoftext|>


<commit_before>""" Parse tsrc config files """

from path import Path
import ruamel.yaml
import schema
from typing import Any, Dict, NewType, Optional
import xdg.BaseDirectory

import tsrc

Config = NewType("Config", Dict[str, Any])

def parse_config(
    file_path: Path,
    config_schema: Optional[schema.Schema] = None,
    roundtrip: bool = False,
) -> Config:
    try:
        contents = file_path.text()
    except OSError as os_error:
        raise tsrc.InvalidConfig(file_path, os_error)
    try:
        if roundtrip:
            yaml = ruamel.yaml.YAML(typ="rt")
        else:
            yaml = ruamel.yaml.YAML(typ="safe", pure=True)
        parsed = yaml.load(contents)
    except ruamel.yaml.error.YAMLError as yaml_error:
        raise tsrc.InvalidConfig(file_path, yaml_error)
    if config_schema:
        try:
            config_schema.validate(parsed)
        except schema.SchemaError as schema_error:
            raise tsrc.InvalidConfig(file_path, schema_error)
    return Config(parsed)

def dump_config(config: Config, path: Path) -> None:
    yaml = ruamel.yaml.YAML()
    with path.open("w") as fileobj:
        yaml.dump(config, fileobj)

def get_tsrc_config_path() -> Path:
    config_path = xdg.BaseDirectory.save_config_path("")
    config_path = Path(config_path) / "tsrc.yml"
    return config_path

def dump_tsrc_config(config: Config) -> None:
    file_path = get_tsrc_config_path()
    dump_config(config, file_path)

def parse_tsrc_config(config_path: Path = None, roundtrip: bool = False) -> Config:
    auth_schema = {
        schema.Optional("gitlab"): {"token": str},
        schema.Optional("github"): {"token": str},
    }
    tsrc_schema = schema.Schema({"auth": auth_schema})
    if not config_path:
        config_path = get_tsrc_config_path()
    return parse_config(config_path, tsrc_schema, roundtrip=roundtrip)
<commit_msg>Allow the tsrc.yml config file to be created if it doesn't already exist Github sources will generate the token file on the user's behalf.<commit_after>""" Parse tsrc config files """

from path import Path
import ruamel.yaml
import schema
from typing import Any, Dict, NewType, Optional
import xdg.BaseDirectory

import tsrc

Config = NewType("Config", Dict[str, Any])

def parse_config(
    file_path: Path,
    config_schema: Optional[schema.Schema] = None,
    roundtrip: bool = False,
) -> Config:
    if not file_path.exists():
        dump_config({"auth": {}}, file_path)

    try:
        contents = file_path.text()
    except OSError as os_error:
        raise tsrc.InvalidConfig(file_path, os_error)
    try:
        if roundtrip:
            yaml = ruamel.yaml.YAML(typ="rt")
        else:
            yaml = ruamel.yaml.YAML(typ="safe", pure=True)
        parsed = yaml.load(contents)
    except ruamel.yaml.error.YAMLError as yaml_error:
        raise tsrc.InvalidConfig(file_path, yaml_error)
    if config_schema:
        try:
            config_schema.validate(parsed)
        except schema.SchemaError as schema_error:
            raise tsrc.InvalidConfig(file_path, schema_error)
    return Config(parsed)

def dump_config(config: Config, path: Path) -> None:
    yaml = ruamel.yaml.YAML()
    with path.open("w") as fileobj:
        yaml.dump(config, fileobj)

def get_tsrc_config_path() -> Path:
    config_path = xdg.BaseDirectory.save_config_path("")
    config_path = Path(config_path) / "tsrc.yml"
    return config_path

def dump_tsrc_config(config: Config) -> None:
    file_path = get_tsrc_config_path()
    dump_config(config, file_path)

def parse_tsrc_config(config_path: Path = None, roundtrip: bool = False) -> Config:
    auth_schema = {
        schema.Optional("gitlab"): {"token": str},
        schema.Optional("github"): {"token": str},
        schema.Optional("github_enterprise"): {"token": str},
    }
    tsrc_schema = schema.Schema({"auth": auth_schema})
    if not config_path:
        config_path = get_tsrc_config_path()
    return parse_config(config_path, tsrc_schema, roundtrip=roundtrip)
<|endoftext|>


<commit_before>#! /usr/bin/env python3
import os
import hashlib
import argparse
import functools
import re
import logging
import textwrap
from collections import defaultdict, namedtuple

logger = logging.getLogger('')

Filename = namedtuple('Filename', ['name', 'base', 'ext', 'path'])

def create_filenames(filenames, root):
    '''
    Makes a generator that yields Filename objects

    Filename objects are a helper to allow multiple representations
    of the same file to be transfered cleanly.
    '''
    logger.info("Creating Filename objects")
    for filename in filenames:
        yield Filename(filename, *os.path.splitext(filename),
                       path=os.path.join(root, filename))

def generate_checksum(filename):
    '''
    A helper function that will generate the
    check sum of a file.
    '''
    logger.info("Generating checksum for {0}".format(filename))
    md5 = hashlib.md5()
    with open(filename, 'rb') as f:
        for chunk in iter(lambda: f.read(128 * md5.block_size), b''):
            md5.update(chunk)
    return md5.digest()

def is_substring(string1, string2):
    '''
    Returns a match if one string is a substring of the other

    For example::

        is_substring('this', 'this1') -> True
        is_substring('that1', 'that') -> True

    But::

        is_substring('that', 'this')  -> False

    '''
    return string1 in string2 or string2 in string1

def pick_shorter_name(file1, file2):
    '''
    This convenience function will help to find the shorter (often better)
    filename.  If the file names are the same length it returns the file
    that is less, hoping for numerically.

    It picks "file.txt" over "file (1).txt", but beware it also picks
    "f.txt" over "file.txt".

    It also picks "file (1).txt" over "file (2).txt"
    '''
    logger.debug("Finding the shortest of {0} and {1}".format(file1.name,
                                                              file2.name))
    if len(file1.name) > len(file2.name):
        return file2
    elif len(file1.name) < len(file2.name) or file1.name < file2.name:
        return file1
    else:
        return file2

def generate_checksum_dict(filenames):
    '''
    This function will create a dictionary of checksums mapped to
    a list of filenames.
    '''
    logger.info("Generating dictionary based on checksum")
    checksum_dict = defaultdict(set)

    for filename in filenames:
        checksum_dict[generate_checksum(filename.path)].add(filename)

    return checksum_dict

def generate_filename_dict(filenames, regex=r'(^.+?)(?: \(\d\))*(\..+)$'):
    '''
    This function will create a dictionary of filename parts mapped to a list
    of the real filenames.
    '''
    logger.info("Generating dictionary based on regular expression")
    filename_dict = defaultdict(set)

    regex = re.compile(regex)

    for filename in filenames:
        match = regex.match(filename.name)
        if match:
            logger.debug('Regex groups for {0}: {1}'.format(
                filename.name, str(match.groups())))
            filename_dict[match.groups()].add(filename)

    return filename_dict

def remove_by_checksum(list_of_names, no_action):
    hashes = generate_checksum_dict(list_of_names)
    for hash in hashes:
        if len(hashes[hash]) > 1:
            logger.info("Investigating duplicate checksum {0}".format(hash))
            logger.debug("Keys for {0} are {1}".format(hash, ', '.join([
                item.name for item in hashes[hash]
            ])))
            best = functools.reduce(pick_shorter_name, hashes[hash])
            for bad in hashes[hash] - {best}:
                if no_action:
                    print('{0} to be deleted'.format(bad.path))
                    logger.info('{0} would have been deleted'.format(bad.path))
                else:
                    logger.info('{0} was deleted'.format(bad.path))
                    os.remove(bad.path)
            logger.info('{0} was kept as only copy'.format(best.path))

        else:
            logger.debug(
                'Skipping non duplicate checksum {0} for key {1}'.format(
                    hash, ', '.join([item.name for item in hashes[hash]])))

def main(path, no_action, recursive, skip_regex, regex_pattern):
    '''
    This function handles all options and steps through the directory
    '''
    for root, dirs, filenames in os.walk(path):
        if not recursive and root != path:
            logger.debug("Skipping child directory {0}".format(root))
            continue

        if not skip_regex:
            names = generate_filename_dict(create_filenames(filenames, root),
                                           regex_pattern)

            for name in names:
                if len(names[name]) > 1:
                    logger.info("Investigating duplicate name {0}".format(name))
                    logger.debug("Keys for {0} are {1}".format(
                        name, ', '.join([item.name for item in names[name]])))
                    remove_by_checksum(names[name], no_action)
                else:
                    logger.debug(
                        'Skipping non duplicate name {0} for key {1}'.format(
                            name, ', '.join([item.name
                                             for item in names[name]])))
        else:
            remove_by_checksum(create_filenames(filenames, root), no_action)

if __name__ == '__main__':
    epilog = '''
    Examples:

        Find matches with default regex:

            $ ./twintrimmer.py -n ~/downloads

        Find matches ignoring the extension:

            $  ls examples/
            Google.html  Google.html~
            $ ./twintrimmer.py -n -p '(^.+?)(?: \(\d\))*\..+' examples/
            examples/Google.html~ to be deleted

        Find matches with "__1" added to basename:

            $ ls examples/underscore/
            file__1.txt  file.txt
            $ ./twintrimmer.py -n -p '(.+?)(?:__\d)*\..*' examples/underscore/
            examples/underscore/file__1.txt to be deleted
    '''

    parser = argparse.ArgumentParser(
        description='Tool for removing duplicate files',
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog=textwrap.dedent(epilog))
    parser.add_argument(
        'path',
        help='This is the path you want to run the checker against')
    parser.add_argument(
        '-n', '--no-action',
        default=False,
        action='store_true',
        help='This will print the output without changing anyfiles')
    parser.add_argument(
        '-r', '--recursive',
        default=False,
        action='store_true',
        help='This option toggles whether the program should search recursively')
    parser.add_argument('--verbosity',
                        type=int,
                        default=0,
                        help='Set debug level')
    parser.add_argument('--log-file',
                        help='This option sets a log file to write.')
    parser.add_argument('--log-level',
                        type=int,
                        default=3,
                        help='Set debug level in log file')
    parser.add_argument(
        '-p', '--pattern',
        type=str,
        default=r'(^.+?)(?: \(\d\))*(\..+)$',
        help='Filename regex (default: r\'(^.+?)(?: \(\d\))*(\..+)$\')')
    parser.add_argument(
        '-c', '--only-checksum',
        default=False,
        action='store_true',
        dest='skip_regex',
        help=
        'This option toggles whether the program searches only by checksum rather than name first')

    args = parser.parse_args()

    if args.log_level != 3 and not args.log_file:
        parser.error('Log level set without log file')

    if args.pattern != r'(^.+?)(?: \(\d\))*(\..+)$' and args.skip_regex:
        parser.error('Pattern set while skipping regex checking')

    stream = logging.StreamHandler()
    stream.setLevel((5 - args.verbosity) * 10)
    formatter_simple = logging.Formatter(
        '%(asctime)s - %(levelname)s - %(message)s')
    stream.setFormatter(formatter_simple)
    logger.addHandler(stream)
    logger.setLevel(logging.DEBUG)

    if args.log_file:
        log_file = logging.FileHandler(args.log_file)
        log_file.setFormatter(formatter_simple)
        log_file.setLevel((5 - args.log_level) * 10)
        logger.addHandler(log_file)

    logger.debug("Args: {0}".format(args))

    main(args.path, args.no_action, args.recursive, args.skip_regex,
         args.pattern)
<commit_msg>Tweaked the --help information<commit_after>#! /usr/bin/env python3
import os
import hashlib
import argparse
import functools
import re
import logging
import textwrap
from collections import defaultdict, namedtuple

logger = logging.getLogger('')

Filename = namedtuple('Filename', ['name', 'base', 'ext', 'path'])

def create_filenames(filenames, root):
    '''
    Makes a generator that yields Filename objects

    Filename objects are a helper to allow multiple representations
    of the same file to be transfered cleanly.
    '''
    logger.info("Creating Filename objects")
    for filename in filenames:
        yield Filename(filename, *os.path.splitext(filename),
                       path=os.path.join(root, filename))

def generate_checksum(filename):
    '''
    A helper function that will generate the
    check sum of a file.
    '''
    logger.info("Generating checksum for {0}".format(filename))
    md5 = hashlib.md5()
    with open(filename, 'rb') as f:
        for chunk in iter(lambda: f.read(128 * md5.block_size), b''):
            md5.update(chunk)
    return md5.digest()

def is_substring(string1, string2):
    '''
    Returns a match if one string is a substring of the other

    For example::

        is_substring('this', 'this1') -> True
        is_substring('that1', 'that') -> True

    But::

        is_substring('that', 'this')  -> False

    '''
    return string1 in string2 or string2 in string1

def pick_shorter_name(file1, file2):
    '''
    This convenience function will help to find the shorter (often better)
    filename.  If the file names are the same length it returns the file
    that is less, hoping for numerically.

    It picks "file.txt" over "file (1).txt", but beware it also picks
    "f.txt" over "file.txt".

    It also picks "file (1).txt" over "file (2).txt"
    '''
    logger.debug("Finding the shortest of {0} and {1}".format(file1.name,
                                                              file2.name))
    if len(file1.name) > len(file2.name):
        return file2
    elif len(file1.name) < len(file2.name) or file1.name < file2.name:
        return file1
    else:
        return file2

def generate_checksum_dict(filenames):
    '''
    This function will create a dictionary of checksums mapped to
    a list of filenames.
    '''
    logger.info("Generating dictionary based on checksum")
    checksum_dict = defaultdict(set)

    for filename in filenames:
        checksum_dict[generate_checksum(filename.path)].add(filename)

    return checksum_dict

def generate_filename_dict(filenames, regex=r'(^.+?)(?: \(\d\))*(\..+)$'):
    '''
    This function will create a dictionary of filename parts mapped to a list
    of the real filenames.
    '''
    logger.info("Generating dictionary based on regular expression")
    filename_dict = defaultdict(set)

    regex = re.compile(regex)

    for filename in filenames:
        match = regex.match(filename.name)
        if match:
            logger.debug('Regex groups for {0}: {1}'.format(
                filename.name, str(match.groups())))
            filename_dict[match.groups()].add(filename)

    return filename_dict

def remove_by_checksum(list_of_names, no_action):
    hashes = generate_checksum_dict(list_of_names)
    for hash in hashes:
        if len(hashes[hash]) > 1:
            logger.info("Investigating duplicate checksum {0}".format(hash))
            logger.debug("Keys for {0} are {1}".format(hash, ', '.join([
                item.name for item in hashes[hash]
            ])))
            best = functools.reduce(pick_shorter_name, hashes[hash])
            for bad in hashes[hash] - {best}:
                if no_action:
                    print('{0} to be deleted'.format(bad.path))
                    logger.info('{0} would have been deleted'.format(bad.path))
                else:
                    logger.info('{0} was deleted'.format(bad.path))
                    os.remove(bad.path)
            logger.info('{0} was kept as only copy'.format(best.path))

        else:
            logger.debug(
                'Skipping non duplicate checksum {0} for key {1}'.format(
                    hash, ', '.join([item.name for item in hashes[hash]])))

def main(path, no_action, recursive, skip_regex, regex_pattern):
    '''
    This function handles all options and steps through the directory
    '''
    for root, dirs, filenames in os.walk(path):
        if not recursive and root != path:
            logger.debug("Skipping child directory {0}".format(root))
            continue

        if not skip_regex:
            names = generate_filename_dict(create_filenames(filenames, root),
                                           regex_pattern)

            for name in names:
                if len(names[name]) > 1:
                    logger.info("Investigating duplicate name {0}".format(name))
                    logger.debug("Keys for {0} are {1}".format(
                        name, ', '.join([item.name for item in names[name]])))
                    remove_by_checksum(names[name], no_action)
                else:
                    logger.debug(
                        'Skipping non duplicate name {0} for key {1}'.format(
                            name, ', '.join([item.name
                                             for item in names[name]])))
        else:
            remove_by_checksum(create_filenames(filenames, root), no_action)

if __name__ == '__main__':
    epilog = '''
    Examples:

        Find matches with default regex:

            $ ./twintrimmer.py -n ~/downloads

        Find matches ignoring the extension:

            $  ls examples/
            Google.html  Google.html~
            $ ./twintrimmer.py -n -p '(^.+?)(?: \(\d\))*\..+' examples/
            examples/Google.html~ to be deleted

        Find matches with "__1" added to basename:

            $ ls examples/underscore/
            file__1.txt  file.txt
            $ ./twintrimmer.py -n -p '(.+?)(?:__\d)*\..*' examples/underscore/
            examples/underscore/file__1.txt to be deleted
    '''

    parser = argparse.ArgumentParser(
        description='Tool for removing duplicate files',
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog=textwrap.dedent(epilog))
    parser.add_argument('path', help='path to check')
    parser.add_argument('-n', '--no-action',
                        default=False,
                        action='store_true',
                        help='show what files would have been deleted')
    parser.add_argument('-r', '--recursive',
                        default=False,
                        action='store_true',
                        help='search directories recursively')
    parser.add_argument('--verbosity',
                        type=int,
                        default=0,
                        help='set print debug level')
    parser.add_argument('--log-file', help='write to log file.')
    parser.add_argument('--log-level',
                        type=int,
                        default=3,
                        help='set log file debug level')
    parser.add_argument('-p', '--pattern',
                        type=str,
                        default=r'(^.+?)(?: \(\d\))*(\..+)$',
                        help='set filename matching regex')
    parser.add_argument(
        '-c', '--only-checksum',
        default=False,
        action='store_true',
        dest='skip_regex',
        help='toggle searching by checksum rather than name first')

    args = parser.parse_args()

    if args.log_level != 3 and not args.log_file:
        parser.error('Log level set without log file')

    if args.pattern != r'(^.+?)(?: \(\d\))*(\..+)$' and args.skip_regex:
        parser.error('Pattern set while skipping regex checking')

    stream = logging.StreamHandler()
    stream.setLevel((5 - args.verbosity) * 10)
    formatter_simple = logging.Formatter(
        '%(asctime)s - %(levelname)s - %(message)s')
    stream.setFormatter(formatter_simple)
    logger.addHandler(stream)
    logger.setLevel(logging.DEBUG)

    if args.log_file:
        log_file = logging.FileHandler(args.log_file)
        log_file.setFormatter(formatter_simple)
        log_file.setLevel((5 - args.log_level) * 10)
        logger.addHandler(log_file)

    logger.debug("Args: {0}".format(args))

    main(args.path, args.no_action, args.recursive, args.skip_regex,
         args.pattern)
<|endoftext|>