<commit_before># The following comment should be removed at some point in the future.
# mypy: disallow-untyped-defs=False

from __future__ import absolute_import

import contextlib
import errno
import logging
import logging.handlers
import os
import sys
from logging import Filter, getLogger

from pip._vendor.six import PY2

from pip._internal.utils.compat import WINDOWS
from pip._internal.utils.deprecation import DEPRECATION_MSG_PREFIX
from pip._internal.utils.misc import ensure_dir

try:
    import threading
except ImportError:
    import dummy_threading as threading  # type: ignore

try:
    # Use "import as" and set colorama in the else clause to avoid mypy
    # errors and get the following correct revealed type for colorama:
    # `Union[_importlib_modulespec.ModuleType, None]`
    # Otherwise, we get an error like the following in the except block:
    #  > Incompatible types in assignment (expression has type "None",
    #   variable has type Module)
    # TODO: eliminate the need to use "import as" once mypy addresses some
    #  of its issues with conditional imports. Here is an umbrella issue:
    #  https://github.com/python/mypy/issues/1297
    from pip._vendor import colorama as _colorama
# Lots of different errors can come from this, including SystemError and
# ImportError.
except Exception:
    colorama = None
else:
    # Import Fore explicitly rather than accessing below as colorama.Fore
    # to avoid the following error running mypy:
    # > Module has no attribute "Fore"
    # TODO: eliminate the need to import Fore once mypy addresses some of its
    #  issues with conditional imports. This particular case could be an
    #  instance of the following issue (but also see the umbrella issue above):
    #  https://github.com/python/mypy/issues/3500
    from pip._vendor.colorama import Fore

    colorama = _colorama

_log_state = threading.local()
_log_state.indentation = 0
subprocess_logger = getLogger('pip.subprocessor')

class BrokenStdoutLoggingError(Exception):
    """
    Raised if BrokenPipeError occurs for the stdout stream while logging.
    """
    pass

# BrokenPipeError does not exist in Python 2 and, in addition, manifests
# differently in Windows and non-Windows.
if WINDOWS:
    # In Windows, a broken pipe can show up as EINVAL rather than EPIPE:
    # https://bugs.python.org/issue19612
    # https://bugs.python.org/issue30418
    if PY2:
        def _is_broken_pipe_error(exc_class, exc):
            """See the docstring for non-Windows Python 3 below."""
            return (exc_class is IOError and
                    exc.errno in (errno.EINVAL, errno.EPIPE))
    else:
        # In Windows, a broken pipe IOError became OSError in Python 3.
        def _is_broken_pipe_error(exc_class, exc):
            """See the docstring for non-Windows Python 3 below."""
            return ((exc_class is BrokenPipeError) or  # noqa: F821
                    (exc_class is OSError and
                     exc.errno in (errno.EINVAL, errno.EPIPE)))
elif PY2:
    def _is_broken_pipe_error(exc_class, exc):
        """See the docstring for non-Windows Python 3 below."""
        return (exc_class is IOError and exc.errno == errno.EPIPE)
else:
    # Then we are in the non-Windows Python 3 case.
    def _is_broken_pipe_error(exc_class, exc):
        """
        Return whether an exception is a broken pipe error.

        Args:
          exc_class: an exception class.
          exc: an exception instance.
        """
        return (exc_class is BrokenPipeError)  # noqa: F821

@contextlib.contextmanager
def indent_log(num=2):
    """
    A context manager which will cause the log output to be indented for any
    log messages emitted inside it.
    """
    # For thread-safety
    _log_state.indentation = get_indentation()
    _log_state.indentation += num
    try:
        yield
    finally:
        _log_state.indentation -= num

def get_indentation():
    return getattr(_log_state, 'indentation', 0)

class IndentingFormatter(logging.Formatter):

    def __init__(self, *args, **kwargs):
        """
        A logging.Formatter that obeys the indent_log() context manager.

        :param add_timestamp: A bool indicating output lines should be prefixed
            with their record's timestamp.
        """
        self.add_timestamp = kwargs.pop("add_timestamp", False)
        super(IndentingFormatter, self).__init__(*args, **kwargs)

    def get_message_start(self, formatted, levelno):
        """
        Return the start of the formatted log message (not counting the
        prefix to add to each line).
        """
        if levelno < logging.WARNING:
            return ''
        if formatted.startswith(DEPRECATION_MSG_PREFIX):
            # Then the message already has a prefix.  We don't want it to
            # look like "WARNING: DEPRECATION: ...."
            return ''
        if levelno < logging.ERROR:
            return 'WARNING: '

        return 'ERROR: '

    def format(self, record):
        """
        Calls the standard formatter, but will indent all of the log message
        lines by our current indentation level.
        """
        formatted = super(IndentingFormatter, self).format(record)
        message_start = self.get_message_start(formatted, record.levelno)
        formatted = message_start + formatted

        prefix = ''
        if self.add_timestamp:
            # TODO: Use Formatter.default_time_format after dropping PY2.
            t = self.formatTime(record, "%Y-%m-%dT%H:%M:%S")
            prefix = '{t},{record.msecs:03.0f} '.format(**locals())
        prefix += " " * get_indentation()
        formatted = "".join([
            prefix + line
            for line in formatted.splitlines(True)
        ])
        return formatted

def _color_wrap(*colors):
    def wrapped(inp):
        return "".join(list(colors) + [inp, colorama.Style.RESET_ALL])
    return wrapped

class ColorizedStreamHandler(logging.StreamHandler):

    # Don't build up a list of colors if we don't have colorama
    if colorama:
        COLORS = [
            # This needs to be in order from highest logging level to lowest.
            (logging.ERROR, _color_wrap(Fore.RED)),
            (logging.WARNING, _color_wrap(Fore.YELLOW)),
        ]
    else:
        COLORS = []

    def __init__(self, stream=None, no_color=None):
        logging.StreamHandler.__init__(self, stream)
        self._no_color = no_color

        if WINDOWS and colorama:
            self.stream = colorama.AnsiToWin32(self.stream)

    def _using_stdout(self):
        """
        Return whether the handler is using sys.stdout.
        """
        if WINDOWS and colorama:
            # Then self.stream is an AnsiToWin32 object.
            return self.stream.wrapped is sys.stdout

        return self.stream is sys.stdout

    def should_color(self):
        # Don't colorize things if we do not have colorama or if told not to
        if not colorama or self._no_color:
            return False

        real_stream = (
            self.stream if not isinstance(self.stream, colorama.AnsiToWin32)
            else self.stream.wrapped
        )

        # If the stream is a tty we should color it
        if hasattr(real_stream, "isatty") and real_stream.isatty():
            return True

        # If we have an ANSI term we should color it
        if os.environ.get("TERM") == "ANSI":
            return True

        # If anything else we should not color it
        return False

    def format(self, record):
        msg = logging.StreamHandler.format(self, record)

        if self.should_color():
            for level, color in self.COLORS:
                if record.levelno >= level:
                    msg = color(msg)
                    break

        return msg

    # The logging module says handleError() can be customized.
    def handleError(self, record):
        exc_class, exc = sys.exc_info()[:2]
        # If a broken pipe occurred while calling write() or flush() on the
        # stdout stream in logging's Handler.emit(), then raise our special
        # exception so we can handle it in main() instead of logging the
        # broken pipe error and continuing.
        if (exc_class and self._using_stdout() and
                _is_broken_pipe_error(exc_class, exc)):
            raise BrokenStdoutLoggingError()

        return super(ColorizedStreamHandler, self).handleError(record)

class BetterRotatingFileHandler(logging.handlers.RotatingFileHandler):

    def _open(self):
        ensure_dir(os.path.dirname(self.baseFilename))
        return logging.handlers.RotatingFileHandler._open(self)

class MaxLevelFilter(Filter):

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

    def filter(self, record):
        return record.levelno < self.level

class ExcludeLoggerFilter(Filter):

    """
    A logging Filter that excludes records from a logger (or its children).
    """

    def filter(self, record):
        # The base Filter class allows only records from a logger (or its
        # children).
        return not super(ExcludeLoggerFilter, self).filter(record)

def setup_logging(verbosity, no_color, user_log_file):
    """Configures and sets up all of the logging

    Returns the requested logging level, as its integer value.
    """

    # Determine the level to be logging at.
    if verbosity >= 1:
        level = "DEBUG"
    elif verbosity == -1:
        level = "WARNING"
    elif verbosity == -2:
        level = "ERROR"
    elif verbosity <= -3:
        level = "CRITICAL"
    else:
        level = "INFO"

    level_number = getattr(logging, level)

    # The "root" logger should match the "console" level *unless* we also need
    # to log to a user log file.
    include_user_log = user_log_file is not None
    if include_user_log:
        additional_log_file = user_log_file
        root_level = "DEBUG"
    else:
        additional_log_file = "/dev/null"
        root_level = level

    # Disable any logging besides WARNING unless we have DEBUG level logging
    # enabled for vendored libraries.
    vendored_log_level = "WARNING" if level in ["INFO", "ERROR"] else "DEBUG"

    # Shorthands for clarity
    log_streams = {
        "stdout": "ext://sys.stdout",
        "stderr": "ext://sys.stderr",
    }
    handler_classes = {
        "stream": "pip._internal.utils.logging.ColorizedStreamHandler",
        "file": "pip._internal.utils.logging.BetterRotatingFileHandler",
    }
    handlers = ["console", "console_errors", "console_subprocess"] + (
        ["user_log"] if include_user_log else []
    )

    logging.config.dictConfig({
        "version": 1,
        "disable_existing_loggers": False,
        "filters": {
            "exclude_warnings": {
                "()": "pip._internal.utils.logging.MaxLevelFilter",
                "level": logging.WARNING,
            },
            "restrict_to_subprocess": {
                "()": "logging.Filter",
                "name": subprocess_logger.name,
            },
            "exclude_subprocess": {
                "()": "pip._internal.utils.logging.ExcludeLoggerFilter",
                "name": subprocess_logger.name,
            },
        },
        "formatters": {
            "indent": {
                "()": IndentingFormatter,
                "format": "%(message)s",
            },
            "indent_with_timestamp": {
                "()": IndentingFormatter,
                "format": "%(message)s",
                "add_timestamp": True,
            },
        },
        "handlers": {
            "console": {
                "level": level,
                "class": handler_classes["stream"],
                "no_color": no_color,
                "stream": log_streams["stdout"],
                "filters": ["exclude_subprocess", "exclude_warnings"],
                "formatter": "indent",
            },
            "console_errors": {
                "level": "WARNING",
                "class": handler_classes["stream"],
                "no_color": no_color,
                "stream": log_streams["stderr"],
                "filters": ["exclude_subprocess"],
                "formatter": "indent",
            },
            # A handler responsible for logging to the console messages
            # from the "subprocessor" logger.
            "console_subprocess": {
                "level": level,
                "class": handler_classes["stream"],
                "no_color": no_color,
                "stream": log_streams["stderr"],
                "filters": ["restrict_to_subprocess"],
                "formatter": "indent",
            },
            "user_log": {
                "level": "DEBUG",
                "class": handler_classes["file"],
                "filename": additional_log_file,
                "delay": True,
                "formatter": "indent_with_timestamp",
            },
        },
        "root": {
            "level": root_level,
            "handlers": handlers,
        },
        "loggers": {
            "pip._vendor": {
                "level": vendored_log_level
            }
        },
    })

    return level_number
<commit_msg>Delete misleading line of code<commit_after># The following comment should be removed at some point in the future.
# mypy: disallow-untyped-defs=False

from __future__ import absolute_import

import contextlib
import errno
import logging
import logging.handlers
import os
import sys
from logging import Filter, getLogger

from pip._vendor.six import PY2

from pip._internal.utils.compat import WINDOWS
from pip._internal.utils.deprecation import DEPRECATION_MSG_PREFIX
from pip._internal.utils.misc import ensure_dir

try:
    import threading
except ImportError:
    import dummy_threading as threading  # type: ignore

try:
    # Use "import as" and set colorama in the else clause to avoid mypy
    # errors and get the following correct revealed type for colorama:
    # `Union[_importlib_modulespec.ModuleType, None]`
    # Otherwise, we get an error like the following in the except block:
    #  > Incompatible types in assignment (expression has type "None",
    #   variable has type Module)
    # TODO: eliminate the need to use "import as" once mypy addresses some
    #  of its issues with conditional imports. Here is an umbrella issue:
    #  https://github.com/python/mypy/issues/1297
    from pip._vendor import colorama as _colorama
# Lots of different errors can come from this, including SystemError and
# ImportError.
except Exception:
    colorama = None
else:
    # Import Fore explicitly rather than accessing below as colorama.Fore
    # to avoid the following error running mypy:
    # > Module has no attribute "Fore"
    # TODO: eliminate the need to import Fore once mypy addresses some of its
    #  issues with conditional imports. This particular case could be an
    #  instance of the following issue (but also see the umbrella issue above):
    #  https://github.com/python/mypy/issues/3500
    from pip._vendor.colorama import Fore

    colorama = _colorama

_log_state = threading.local()
subprocess_logger = getLogger('pip.subprocessor')

class BrokenStdoutLoggingError(Exception):
    """
    Raised if BrokenPipeError occurs for the stdout stream while logging.
    """
    pass

# BrokenPipeError does not exist in Python 2 and, in addition, manifests
# differently in Windows and non-Windows.
if WINDOWS:
    # In Windows, a broken pipe can show up as EINVAL rather than EPIPE:
    # https://bugs.python.org/issue19612
    # https://bugs.python.org/issue30418
    if PY2:
        def _is_broken_pipe_error(exc_class, exc):
            """See the docstring for non-Windows Python 3 below."""
            return (exc_class is IOError and
                    exc.errno in (errno.EINVAL, errno.EPIPE))
    else:
        # In Windows, a broken pipe IOError became OSError in Python 3.
        def _is_broken_pipe_error(exc_class, exc):
            """See the docstring for non-Windows Python 3 below."""
            return ((exc_class is BrokenPipeError) or  # noqa: F821
                    (exc_class is OSError and
                     exc.errno in (errno.EINVAL, errno.EPIPE)))
elif PY2:
    def _is_broken_pipe_error(exc_class, exc):
        """See the docstring for non-Windows Python 3 below."""
        return (exc_class is IOError and exc.errno == errno.EPIPE)
else:
    # Then we are in the non-Windows Python 3 case.
    def _is_broken_pipe_error(exc_class, exc):
        """
        Return whether an exception is a broken pipe error.

        Args:
          exc_class: an exception class.
          exc: an exception instance.
        """
        return (exc_class is BrokenPipeError)  # noqa: F821

@contextlib.contextmanager
def indent_log(num=2):
    """
    A context manager which will cause the log output to be indented for any
    log messages emitted inside it.
    """
    # For thread-safety
    _log_state.indentation = get_indentation()
    _log_state.indentation += num
    try:
        yield
    finally:
        _log_state.indentation -= num

def get_indentation():
    return getattr(_log_state, 'indentation', 0)

class IndentingFormatter(logging.Formatter):

    def __init__(self, *args, **kwargs):
        """
        A logging.Formatter that obeys the indent_log() context manager.

        :param add_timestamp: A bool indicating output lines should be prefixed
            with their record's timestamp.
        """
        self.add_timestamp = kwargs.pop("add_timestamp", False)
        super(IndentingFormatter, self).__init__(*args, **kwargs)

    def get_message_start(self, formatted, levelno):
        """
        Return the start of the formatted log message (not counting the
        prefix to add to each line).
        """
        if levelno < logging.WARNING:
            return ''
        if formatted.startswith(DEPRECATION_MSG_PREFIX):
            # Then the message already has a prefix.  We don't want it to
            # look like "WARNING: DEPRECATION: ...."
            return ''
        if levelno < logging.ERROR:
            return 'WARNING: '

        return 'ERROR: '

    def format(self, record):
        """
        Calls the standard formatter, but will indent all of the log message
        lines by our current indentation level.
        """
        formatted = super(IndentingFormatter, self).format(record)
        message_start = self.get_message_start(formatted, record.levelno)
        formatted = message_start + formatted

        prefix = ''
        if self.add_timestamp:
            # TODO: Use Formatter.default_time_format after dropping PY2.
            t = self.formatTime(record, "%Y-%m-%dT%H:%M:%S")
            prefix = '{t},{record.msecs:03.0f} '.format(**locals())
        prefix += " " * get_indentation()
        formatted = "".join([
            prefix + line
            for line in formatted.splitlines(True)
        ])
        return formatted

def _color_wrap(*colors):
    def wrapped(inp):
        return "".join(list(colors) + [inp, colorama.Style.RESET_ALL])
    return wrapped

class ColorizedStreamHandler(logging.StreamHandler):

    # Don't build up a list of colors if we don't have colorama
    if colorama:
        COLORS = [
            # This needs to be in order from highest logging level to lowest.
            (logging.ERROR, _color_wrap(Fore.RED)),
            (logging.WARNING, _color_wrap(Fore.YELLOW)),
        ]
    else:
        COLORS = []

    def __init__(self, stream=None, no_color=None):
        logging.StreamHandler.__init__(self, stream)
        self._no_color = no_color

        if WINDOWS and colorama:
            self.stream = colorama.AnsiToWin32(self.stream)

    def _using_stdout(self):
        """
        Return whether the handler is using sys.stdout.
        """
        if WINDOWS and colorama:
            # Then self.stream is an AnsiToWin32 object.
            return self.stream.wrapped is sys.stdout

        return self.stream is sys.stdout

    def should_color(self):
        # Don't colorize things if we do not have colorama or if told not to
        if not colorama or self._no_color:
            return False

        real_stream = (
            self.stream if not isinstance(self.stream, colorama.AnsiToWin32)
            else self.stream.wrapped
        )

        # If the stream is a tty we should color it
        if hasattr(real_stream, "isatty") and real_stream.isatty():
            return True

        # If we have an ANSI term we should color it
        if os.environ.get("TERM") == "ANSI":
            return True

        # If anything else we should not color it
        return False

    def format(self, record):
        msg = logging.StreamHandler.format(self, record)

        if self.should_color():
            for level, color in self.COLORS:
                if record.levelno >= level:
                    msg = color(msg)
                    break

        return msg

    # The logging module says handleError() can be customized.
    def handleError(self, record):
        exc_class, exc = sys.exc_info()[:2]
        # If a broken pipe occurred while calling write() or flush() on the
        # stdout stream in logging's Handler.emit(), then raise our special
        # exception so we can handle it in main() instead of logging the
        # broken pipe error and continuing.
        if (exc_class and self._using_stdout() and
                _is_broken_pipe_error(exc_class, exc)):
            raise BrokenStdoutLoggingError()

        return super(ColorizedStreamHandler, self).handleError(record)

class BetterRotatingFileHandler(logging.handlers.RotatingFileHandler):

    def _open(self):
        ensure_dir(os.path.dirname(self.baseFilename))
        return logging.handlers.RotatingFileHandler._open(self)

class MaxLevelFilter(Filter):

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

    def filter(self, record):
        return record.levelno < self.level

class ExcludeLoggerFilter(Filter):

    """
    A logging Filter that excludes records from a logger (or its children).
    """

    def filter(self, record):
        # The base Filter class allows only records from a logger (or its
        # children).
        return not super(ExcludeLoggerFilter, self).filter(record)

def setup_logging(verbosity, no_color, user_log_file):
    """Configures and sets up all of the logging

    Returns the requested logging level, as its integer value.
    """

    # Determine the level to be logging at.
    if verbosity >= 1:
        level = "DEBUG"
    elif verbosity == -1:
        level = "WARNING"
    elif verbosity == -2:
        level = "ERROR"
    elif verbosity <= -3:
        level = "CRITICAL"
    else:
        level = "INFO"

    level_number = getattr(logging, level)

    # The "root" logger should match the "console" level *unless* we also need
    # to log to a user log file.
    include_user_log = user_log_file is not None
    if include_user_log:
        additional_log_file = user_log_file
        root_level = "DEBUG"
    else:
        additional_log_file = "/dev/null"
        root_level = level

    # Disable any logging besides WARNING unless we have DEBUG level logging
    # enabled for vendored libraries.
    vendored_log_level = "WARNING" if level in ["INFO", "ERROR"] else "DEBUG"

    # Shorthands for clarity
    log_streams = {
        "stdout": "ext://sys.stdout",
        "stderr": "ext://sys.stderr",
    }
    handler_classes = {
        "stream": "pip._internal.utils.logging.ColorizedStreamHandler",
        "file": "pip._internal.utils.logging.BetterRotatingFileHandler",
    }
    handlers = ["console", "console_errors", "console_subprocess"] + (
        ["user_log"] if include_user_log else []
    )

    logging.config.dictConfig({
        "version": 1,
        "disable_existing_loggers": False,
        "filters": {
            "exclude_warnings": {
                "()": "pip._internal.utils.logging.MaxLevelFilter",
                "level": logging.WARNING,
            },
            "restrict_to_subprocess": {
                "()": "logging.Filter",
                "name": subprocess_logger.name,
            },
            "exclude_subprocess": {
                "()": "pip._internal.utils.logging.ExcludeLoggerFilter",
                "name": subprocess_logger.name,
            },
        },
        "formatters": {
            "indent": {
                "()": IndentingFormatter,
                "format": "%(message)s",
            },
            "indent_with_timestamp": {
                "()": IndentingFormatter,
                "format": "%(message)s",
                "add_timestamp": True,
            },
        },
        "handlers": {
            "console": {
                "level": level,
                "class": handler_classes["stream"],
                "no_color": no_color,
                "stream": log_streams["stdout"],
                "filters": ["exclude_subprocess", "exclude_warnings"],
                "formatter": "indent",
            },
            "console_errors": {
                "level": "WARNING",
                "class": handler_classes["stream"],
                "no_color": no_color,
                "stream": log_streams["stderr"],
                "filters": ["exclude_subprocess"],
                "formatter": "indent",
            },
            # A handler responsible for logging to the console messages
            # from the "subprocessor" logger.
            "console_subprocess": {
                "level": level,
                "class": handler_classes["stream"],
                "no_color": no_color,
                "stream": log_streams["stderr"],
                "filters": ["restrict_to_subprocess"],
                "formatter": "indent",
            },
            "user_log": {
                "level": "DEBUG",
                "class": handler_classes["file"],
                "filename": additional_log_file,
                "delay": True,
                "formatter": "indent_with_timestamp",
            },
        },
        "root": {
            "level": root_level,
            "handlers": handlers,
        },
        "loggers": {
            "pip._vendor": {
                "level": vendored_log_level
            }
        },
    })

    return level_number
<|endoftext|>


<commit_before>import pymongo
import pandas as pd
from math import log
from time import time
import sys
from scipy.stats import linregress
import pickle

client = pymongo.MongoClient()
db = client['bitmicro']

def get_book_df(symbol, limit, convert_timestamps=False):
    '''
    Returns a DataFrame of book data
    '''
    books_db = db[symbol+'_books']
    cursor = books_db.find().sort('_id', -1).limit(limit)
    books = pd.DataFrame(list(cursor))
    books = books.set_index('_id')
    if convert_timestamps:
        books.index = pd.to_datetime(books.index, unit='s')

    def to_df(x):
        return pd.DataFrame(x[:10])
    return books.applymap(to_df).sort_index()

def get_width_and_mid(books):
    '''
    Returns width of best market and midpoint for each data point in DataFrame
    of book data
    '''
    best_bid = books.bids.apply(lambda x: x.price[0])
    best_ask = books.asks.apply(lambda x: x.price[0])
    return best_ask-best_bid, (best_bid + best_ask)/2

def get_future_mid(books, offset, sensitivity):
    '''
    Returns percent change of future midpoints for each data point in DataFrame
    of book data
    '''

    def future(timestamp):
        i = books.index.get_loc(timestamp+offset, method='nearest')
        if abs(books.index[i] - (timestamp+offset)) < sensitivity:
            return books.mid.iloc[i]
    return (books.index.map(future)/books.mid).apply(log)

def get_power_imbalance(books, n=10, power=2):
    '''
    Returns a measure of the imbalance between bids and offers for each data
    point in DataFrame of book data
    '''

    def calc_imbalance(book):
        def calc(x):
            return x.amount*(.5*book.width/(x.price-book.mid))**power
        bid_imbalance = book.bids.iloc[:n].apply(calc, axis=1)
        ask_imbalance = book.asks.iloc[:n].apply(calc, axis=1)
        return (bid_imbalance-ask_imbalance).sum()
    imbalance = books.apply(calc_imbalance, axis=1)
    return imbalance

def get_power_adjusted_price(books, n=10, power=2):
    '''
    Returns the percent change of an average of order prices weighted by inverse
    distance-wieghted volume for each data point in DataFrame of book data
    '''

    def calc_adjusted_price(book):
        def calc(x):
            return x.amount*(.5*book.width/(x.price-book.mid))**power
        bid_inv = 1/book.bids.iloc[:n].apply(calc, axis=1)
        ask_inv = 1/book.asks.iloc[:n].apply(calc, axis=1)
        bid_price = book.bids.price.iloc[:n]
        ask_price = book.asks.price.iloc[:n]
        return (bid_price*bid_inv + ask_price*ask_inv).sum() /\
            (bid_inv + ask_inv).sum()
    adjusted = books.apply(calc_adjusted_price, axis=1)
    return (adjusted/books.mid).apply(log).fillna(0)

def get_trade_df(symbol, min_ts, max_ts, convert_timestamps=False):
    '''
    Returns a DataFrame of trades for symbol in time range
    '''
    trades_db = db[symbol+'_trades']
    query = {'timestamp': {'$gt': min_ts, '$lt': max_ts}}
    cursor = trades_db.find(query).sort('_id', pymongo.ASCENDING)
    trades = pd.DataFrame(list(cursor))
    if not trades.empty:
        trades = trades.set_index('_id')
        if convert_timestamps:
            trades.index = pd.to_datetime(trades.index, unit='s')
    return trades

def get_trades_indexes(books, trades, offset, live=False):
    '''
    Returns indexes of trades in offset range for each data point in DataFrame
    of book data
    '''
    def indexes(ts):
        ts = int(ts)
        i_0 = trades.timestamp.searchsorted([ts-offset], side='left')[0]
        if live:
            i_n = -1
        else:
            i_n = trades.timestamp.searchsorted([ts-1], side='right')[0]
        return (i_0, i_n)
    return books.index.map(indexes)

def get_trades_count(books, trades):
    '''
    Returns a count of trades for each data point in DataFrame of book data
    '''
    def count(x):
        return len(trades.iloc[x.indexes[0]:x.indexes[1]])
    return books.apply(count, axis=1)

def get_trades_average(books, trades):
    '''
    Returns the percent change of a volume-weighted average of trades for each
    data point in DataFrame of book data
    '''

    def mean_trades(x):
        trades_n = trades.iloc[x.indexes[0]:x.indexes[1]]
        if not trades_n.empty:
            return (trades_n.price*trades_n.amount).sum()/trades_n.amount.sum()
    return (books.mid/books.apply(mean_trades, axis=1)).apply(log).fillna(0)

def get_aggressor(books, trades):
    '''
    Returns a measure of whether trade aggressors were buyers or sellers for
    each data point in DataFrame of book data
    '''

    def aggressor(x):
        trades_n = trades.iloc[x.indexes[0]:x.indexes[1]]
        if trades_n.empty:
            return 0
        buys = trades_n['type'] == 'buy'
        buy_vol = trades_n[buys].amount.sum()
        sell_vol = trades_n[~buys].amount.sum()
        return buy_vol - sell_vol
    return books.apply(aggressor, axis=1)

def get_trend(books, trades):
    '''
    Returns the linear trend in previous trades for each data point in DataFrame
    of book data
    '''

    def trend(x):
        trades_n = trades.iloc[x.indexes[0]:x.indexes[1]]
        if len(trades_n) < 3:
            return 0
        else:
            return linregress(trades_n.index.values, trades_n.price.values)[0]
    return books.apply(trend, axis=1)

def check_times(books):
    '''
    Returns list of differences between collection time and max book timestamps
    for verification purposes
    '''
    time_diff = []
    for i in range(len(books)):
        book = books.iloc[i]
        ask_ts = max(book.asks.timestamp)
        bid_ts = max(book.bids.timestamp)
        ts = max(ask_ts, bid_ts)
        time_diff.append(book.name-ts)
    return time_diff

def make_features(symbol, sample, mid_offsets,
                  trades_offsets, powers, live=False):
    '''
    Returns a DataFrame with targets and features
    '''
    start = time()
    stage = time()

    # Book related features:
    books = get_book_df(symbol, sample)
    if not live:
        print 'get book data run time:', (time()-stage)/60, 'minutes'
        stage = time()
    books['width'], books['mid'] = get_width_and_mid(books)
    if not live:
        print 'width and mid run time:', (time()-stage)/60, 'minutes'
        stage = time()
    for n in mid_offsets:
        books['mid{}'.format(n)] = get_future_mid(books, n, sensitivity=1)
    if not live:
        books = books.dropna()
        print 'offset mids run time:', (time()-stage)/60, 'minutes'
        stage = time()
    for p in powers:
        books['imbalance{}'.format(p)] = get_power_imbalance(books, 10, p)
        books['adj_price{}'.format(p)] = get_power_adjusted_price(books, 10, p)
    if not live:
        print 'power calcs run time:', (time()-stage)/60, 'minutes'
        stage = time()
    books = books.drop(['bids', 'asks'], axis=1)

    # Trade related features:
    min_ts = books.index.min() - trades_offsets[-1]
    max_ts = books.index.max()
    if live:
        max_ts += 10
    trades = get_trade_df(symbol, min_ts, max_ts)
    for n in trades_offsets:
        if trades.empty:
            books['indexes'] = 0
            books['t{}_count'.format(n)] = 0
            books['t{}_av'.format(n)] = 0
            books['agg{}'.format(n)] = 0
            books['trend{}'.format(n)] = 0
        else:
            books['indexes'] = get_trades_indexes(books, trades, n, live)
            books['t{}_count'.format(n)] = get_trades_count(books, trades)
            books['t{}_av'.format(n)] = get_trades_average(books, trades)
            books['agg{}'.format(n)] = get_aggressor(books, trades)
            books['trend{}'.format(n)] = get_trend(books, trades)
    if not live:
        print 'trade features run time:', (time()-stage)/60, 'minutes'
        stage = time()
        print 'make_features run time:', (time()-start)/60, 'minutes'

    return books.drop('indexes', axis=1)

def make_data(symbol, sample):
    '''
    Convenience function for calling make_features
    '''
    data = make_features(symbol,
                         sample=sample,
                         mid_offsets=[30],
                         trades_offsets=[10, 30, 60, 120, 180],
                         powers=[2, 4, 8])
    return data

if __name__ == '__main__' and len(sys.argv) == 4:
    data = make_data(sys.argv[1], int(sys.argv[2]))
    with open(sys.argv[3], 'w+') as f:
        pickle.dump(data, f)
<commit_msg>remove 10-second trade offset<commit_after>import pymongo
import pandas as pd
from math import log
from time import time
import sys
from scipy.stats import linregress
import pickle

client = pymongo.MongoClient()
db = client['bitmicro']

def get_book_df(symbol, limit, convert_timestamps=False):
    '''
    Returns a DataFrame of book data
    '''
    books_db = db[symbol+'_books']
    cursor = books_db.find().sort('_id', -1).limit(limit)
    books = pd.DataFrame(list(cursor))
    books = books.set_index('_id')
    if convert_timestamps:
        books.index = pd.to_datetime(books.index, unit='s')

    def to_df(x):
        return pd.DataFrame(x[:10])
    return books.applymap(to_df).sort_index()

def get_width_and_mid(books):
    '''
    Returns width of best market and midpoint for each data point in DataFrame
    of book data
    '''
    best_bid = books.bids.apply(lambda x: x.price[0])
    best_ask = books.asks.apply(lambda x: x.price[0])
    return best_ask-best_bid, (best_bid + best_ask)/2

def get_future_mid(books, offset, sensitivity):
    '''
    Returns percent change of future midpoints for each data point in DataFrame
    of book data
    '''

    def future(timestamp):
        i = books.index.get_loc(timestamp+offset, method='nearest')
        if abs(books.index[i] - (timestamp+offset)) < sensitivity:
            return books.mid.iloc[i]
    return (books.index.map(future)/books.mid).apply(log)

def get_power_imbalance(books, n=10, power=2):
    '''
    Returns a measure of the imbalance between bids and offers for each data
    point in DataFrame of book data
    '''

    def calc_imbalance(book):
        def calc(x):
            return x.amount*(.5*book.width/(x.price-book.mid))**power
        bid_imbalance = book.bids.iloc[:n].apply(calc, axis=1)
        ask_imbalance = book.asks.iloc[:n].apply(calc, axis=1)
        return (bid_imbalance-ask_imbalance).sum()
    imbalance = books.apply(calc_imbalance, axis=1)
    return imbalance

def get_power_adjusted_price(books, n=10, power=2):
    '''
    Returns the percent change of an average of order prices weighted by inverse
    distance-wieghted volume for each data point in DataFrame of book data
    '''

    def calc_adjusted_price(book):
        def calc(x):
            return x.amount*(.5*book.width/(x.price-book.mid))**power
        bid_inv = 1/book.bids.iloc[:n].apply(calc, axis=1)
        ask_inv = 1/book.asks.iloc[:n].apply(calc, axis=1)
        bid_price = book.bids.price.iloc[:n]
        ask_price = book.asks.price.iloc[:n]
        return (bid_price*bid_inv + ask_price*ask_inv).sum() /\
            (bid_inv + ask_inv).sum()
    adjusted = books.apply(calc_adjusted_price, axis=1)
    return (adjusted/books.mid).apply(log).fillna(0)

def get_trade_df(symbol, min_ts, max_ts, convert_timestamps=False):
    '''
    Returns a DataFrame of trades for symbol in time range
    '''
    trades_db = db[symbol+'_trades']
    query = {'timestamp': {'$gt': min_ts, '$lt': max_ts}}
    cursor = trades_db.find(query).sort('_id', pymongo.ASCENDING)
    trades = pd.DataFrame(list(cursor))
    if not trades.empty:
        trades = trades.set_index('_id')
        if convert_timestamps:
            trades.index = pd.to_datetime(trades.index, unit='s')
    return trades

def get_trades_indexes(books, trades, offset, live=False):
    '''
    Returns indexes of trades in offset range for each data point in DataFrame
    of book data
    '''
    def indexes(ts):
        ts = int(ts)
        i_0 = trades.timestamp.searchsorted([ts-offset], side='left')[0]
        if live:
            i_n = -1
        else:
            i_n = trades.timestamp.searchsorted([ts-1], side='right')[0]
        return (i_0, i_n)
    return books.index.map(indexes)

def get_trades_count(books, trades):
    '''
    Returns a count of trades for each data point in DataFrame of book data
    '''
    def count(x):
        return len(trades.iloc[x.indexes[0]:x.indexes[1]])
    return books.apply(count, axis=1)

def get_trades_average(books, trades):
    '''
    Returns the percent change of a volume-weighted average of trades for each
    data point in DataFrame of book data
    '''

    def mean_trades(x):
        trades_n = trades.iloc[x.indexes[0]:x.indexes[1]]
        if not trades_n.empty:
            return (trades_n.price*trades_n.amount).sum()/trades_n.amount.sum()
    return (books.mid/books.apply(mean_trades, axis=1)).apply(log).fillna(0)

def get_aggressor(books, trades):
    '''
    Returns a measure of whether trade aggressors were buyers or sellers for
    each data point in DataFrame of book data
    '''

    def aggressor(x):
        trades_n = trades.iloc[x.indexes[0]:x.indexes[1]]
        if trades_n.empty:
            return 0
        buys = trades_n['type'] == 'buy'
        buy_vol = trades_n[buys].amount.sum()
        sell_vol = trades_n[~buys].amount.sum()
        return buy_vol - sell_vol
    return books.apply(aggressor, axis=1)

def get_trend(books, trades):
    '''
    Returns the linear trend in previous trades for each data point in DataFrame
    of book data
    '''

    def trend(x):
        trades_n = trades.iloc[x.indexes[0]:x.indexes[1]]
        if len(trades_n) < 3:
            return 0
        else:
            return linregress(trades_n.index.values, trades_n.price.values)[0]
    return books.apply(trend, axis=1)

def check_times(books):
    '''
    Returns list of differences between collection time and max book timestamps
    for verification purposes
    '''
    time_diff = []
    for i in range(len(books)):
        book = books.iloc[i]
        ask_ts = max(book.asks.timestamp)
        bid_ts = max(book.bids.timestamp)
        ts = max(ask_ts, bid_ts)
        time_diff.append(book.name-ts)
    return time_diff

def make_features(symbol, sample, mid_offsets,
                  trades_offsets, powers, live=False):
    '''
    Returns a DataFrame with targets and features
    '''
    start = time()
    stage = time()

    # Book related features:
    books = get_book_df(symbol, sample)
    if not live:
        print 'get book data run time:', (time()-stage)/60, 'minutes'
        stage = time()
    books['width'], books['mid'] = get_width_and_mid(books)
    if not live:
        print 'width and mid run time:', (time()-stage)/60, 'minutes'
        stage = time()
    for n in mid_offsets:
        books['mid{}'.format(n)] = get_future_mid(books, n, sensitivity=1)
    if not live:
        books = books.dropna()
        print 'offset mids run time:', (time()-stage)/60, 'minutes'
        stage = time()
    for p in powers:
        books['imbalance{}'.format(p)] = get_power_imbalance(books, 10, p)
        books['adj_price{}'.format(p)] = get_power_adjusted_price(books, 10, p)
    if not live:
        print 'power calcs run time:', (time()-stage)/60, 'minutes'
        stage = time()
    books = books.drop(['bids', 'asks'], axis=1)

    # Trade related features:
    min_ts = books.index.min() - trades_offsets[-1]
    max_ts = books.index.max()
    if live:
        max_ts += 10
    trades = get_trade_df(symbol, min_ts, max_ts)
    for n in trades_offsets:
        if trades.empty:
            books['indexes'] = 0
            books['t{}_count'.format(n)] = 0
            books['t{}_av'.format(n)] = 0
            books['agg{}'.format(n)] = 0
            books['trend{}'.format(n)] = 0
        else:
            books['indexes'] = get_trades_indexes(books, trades, n, live)
            books['t{}_count'.format(n)] = get_trades_count(books, trades)
            books['t{}_av'.format(n)] = get_trades_average(books, trades)
            books['agg{}'.format(n)] = get_aggressor(books, trades)
            books['trend{}'.format(n)] = get_trend(books, trades)
    if not live:
        print 'trade features run time:', (time()-stage)/60, 'minutes'
        stage = time()
        print 'make_features run time:', (time()-start)/60, 'minutes'

    return books.drop('indexes', axis=1)

def make_data(symbol, sample):
    '''
    Convenience function for calling make_features
    '''
    data = make_features(symbol,
                         sample=sample,
                         mid_offsets=[30],
                         trades_offsets=[30, 60, 120, 180],
                         powers=[2, 4, 8])
    return data

if __name__ == '__main__' and len(sys.argv) == 4:
    data = make_data(sys.argv[1], int(sys.argv[2]))
    with open(sys.argv[3], 'w+') as f:
        pickle.dump(data, f)
<|endoftext|>


<commit_before># -*- coding: utf-8 -*-
"""Class to handle the contents of a file containing unigrams."""
from __future__ import absolute_import, division, print_function, unicode_literals
from collections import Counter, OrderedDict
from model.datasets import load_articles

class Unigrams():
    """Class to handle the contents of a file containing unigrams."""
    def __init__(self, filepath=None, skip_first_n=0, max_count_words=None):
        """Initialize the unigrams list, optionally from a file.
        Args:
            filepath: Optional filepath to a file containing the unigrams in the
                form "word<tab>count<linebreak>word<tab>count...".
            skip_first_n: Number of words to skip at the start of the unigrams file.
            max_count_words: Maxmimum number of words to read from the unigrams file.
        """
        self.word_to_rank = OrderedDict()
        self.word_to_count = OrderedDict()
        self.sum_of_counts = 0
        if filepath is not None:
            self.fill_from_file(filepath, skip_first_n=skip_first_n, max_count_words=max_count_words)

    def clear(self):
        """Resets this class, empties all ranking and count dictionaries."""
        self.word_to_rank = OrderedDict()
        self.word_to_count = OrderedDict()

    def fill_from_file(self, filepath, skip_first_n=0, max_count_words=None):
        """Fills the dictionaries of this class from a file containing unigrams.
        Expected structure of the file:
            foo    120
            bar    105
            ...
        The words have to separated from the counts by a tab.
        The words have to ordered descencnding, from the most common to the least common word.

        Args:
            filepath: Optional filepath to a file containing the unigrams in the
                form "word<tab>count<linebreak>word<tab>count...".
            skip_first_n: Number of words to skip at the start of the unigrams file.
            max_count_words: Maxmimum number of words to read from the unigrams file.
        """
        skipped = 0
        added = 0
        with open(filepath, "r") as f:
            for line_idx, line in enumerate(f):
                columns = line.decode("utf-8").strip().split("\t")
                if len(columns) == 2:
                    if skipped < skip_first_n:
                        skipped += 1
                    else:
                        if max_count_words is not None and added >= max_count_words:
                            break
                        else:
                            word, count = columns
                            count = int(count)
                            self.word_to_rank[word] = added + 1
                            self.word_to_count[word] = count
                            self.sum_of_counts += count
                            added += 1
                else:
                    print("[Warning] Expected 2 columns in unigrams file at line %d, got %d" % (line_idx, len(columns)))

    def fill_from_articles(self, filepath, verbose=False):
        """Fills the dictionaries of this class from a corpus file.

        The corpus file is expected to contain one article/document per line.
        The corpus file may contain labels at each word, e.g. "John/PER Doe/PER did yesterday...".
        These labels will be ignored by this function (e.g. only "John" and "Doe" would be used).
        All count values and ranks will be automatically estimated.

        Note: This function is rather slow.

        Args:
            filepath: Filepath to the corpus file.
            verbose: Whether to output messages during parsing.
        """
        self.fill_from_articles_labels(filepath, labels=None, verbose=verbose)

    def fill_from_articles_labels(self, filepath, labels=None, verbose=False):
        """Fills the dictionaries of this class from a corpus file, optionally only with words of
        specific labels (e.g. only words labeled with "PER").

        The corpus file is expected to contain one article/document per line.
        The corpus file may contain labels at each word, e.g. "John/PER Doe/PER did yesterday...".
        All count values and ranks will be automatically estimated.

        Note: This function is rather slow.

        Args:
            filepath: Filepath to the corpus file.
            labels: Optionally one or more labels. If provided, only words that are annotated
                with any of these labels will be counted.
            verbose: Whether to output messages during parsing.
        """
        assert labels is None or type(labels) == type(list())
        assert labels is None or len(labels) > 0

        counts = Counter()
        self.sum_of_counts = 0

        articles = load_articles(filepath, start_at=0)
        for i, article in enumerate(articles):
            words = [token.word for token in article.tokens if labels is None or token.label in labels]
            counts.update(words)
            if verbose and i % 1000 == 0:
                print("Article %d" % (i))

        most_common = counts.most_common()
        for i, (word, count) in enumerate(most_common):
            self.word_to_count[word] = count
            self.word_to_rank[word] = i + 1
            self.sum_of_counts += count

    def write_to_file(self, filepath):
        """Writes the contents of this unigrams object to a file.
        The file can later on be loaded with fill_from_file().
        Args:
            filepath: Filepath to the file to which to write.
        """
        with open(filepath, "w") as f:
            for i, (word, count) in enumerate(self.word_to_count.iteritems()):
                print(word, count)
                if i > 0:
                    f.write("\n")
                f.write(word.encode("utf-8"))
                f.write("\t")
                f.write(str(count))

    def get_rank_of(self, word, default=-1):
        """Returns the rank of a word among all unigrams.
        The most common word has rank 1.
        Args:
            word: The word which to rank.
            default: A default value to return if the word was not found among the unigrams.
        Returns:
            integer or default value (-1).
        """
        if word in self.word_to_rank:
            return self.word_to_rank[word]
        else:
            return default

    def get_count_of(self, word, default=-1):
        """Returns the count of a word among all unigrams.
        Args:
            word: The word which to rank.
            default: A default value to return if the word was not found among the unigrams.
        Returns:
            integer or default value (-1).
        """
        if word in self.word_to_count:
            return self.word_to_count[word]
        else:
            return default

    def get_frequency_of(self, word, default=None):
        """Returns the frequency of a word among all unigrams.
        The frequency is calculated by count(word)/count(all words)
        Args:
            word: The word which to rank.
            default: A default value to return if the word was not found among the unigrams.
        Returns:
            double or default value (None).
        """
        count = self.get_count_of(word, None)
        if count is not None:
            return count / self.sum_of_counts
        else:
            return default
<commit_msg>Optimize styling according to pylint<commit_after># -*- coding: utf-8 -*-
"""Class to handle the contents of a file containing unigrams."""
from __future__ import absolute_import, division, print_function, unicode_literals
from collections import Counter, OrderedDict
from model.datasets import load_articles

class Unigrams(object):
    """Class to handle the contents of a file containing unigrams."""
    def __init__(self, filepath=None, skip_first_n=0, max_count_words=None):
        """Initialize the unigrams list, optionally from a file.
        Args:
            filepath: Optional filepath to a file containing the unigrams in the
                form "word<tab>count<linebreak>word<tab>count...".
            skip_first_n: Number of words to skip at the start of the unigrams file.
            max_count_words: Maxmimum number of words to read from the unigrams file.
        """
        self.word_to_rank = OrderedDict()
        self.word_to_count = OrderedDict()
        self.sum_of_counts = 0
        if filepath is not None:
            self.fill_from_file(filepath, skip_first_n=skip_first_n,
                                max_count_words=max_count_words)

    def clear(self):
        """Resets this class, empties all ranking and count dictionaries."""
        self.word_to_rank = OrderedDict()
        self.word_to_count = OrderedDict()

    def fill_from_file(self, filepath, skip_first_n=0, max_count_words=None):
        """Fills the dictionaries of this class from a file containing unigrams.
        Expected structure of the file:
            foo    120
            bar    105
            ...
        The words have to separated from the counts by a tab.
        The words have to ordered descencnding, from the most common to the least common word.

        Args:
            filepath: Optional filepath to a file containing the unigrams in the
                form "word<tab>count<linebreak>word<tab>count...".
            skip_first_n: Number of words to skip at the start of the unigrams file.
            max_count_words: Maxmimum number of words to read from the unigrams file.
        """
        skipped = 0
        added = 0
        with open(filepath, "r") as handle:
            for line_idx, line in enumerate(handle):
                columns = line.decode("utf-8").strip().split("\t")
                if len(columns) == 2:
                    if skipped < skip_first_n:
                        skipped += 1
                    else:
                        if max_count_words is not None and added >= max_count_words:
                            break
                        else:
                            word, count = columns
                            count = int(count)
                            self.word_to_rank[word] = added + 1
                            self.word_to_count[word] = count
                            self.sum_of_counts += count
                            added += 1
                else:
                    print("[Warning] Expected 2 columns in unigrams file at line %d, " \
                          "got %d" % (line_idx, len(columns)))

    def fill_from_articles(self, filepath, verbose=False):
        """Fills the dictionaries of this class from a corpus file.

        The corpus file is expected to contain one article/document per line.
        The corpus file may contain labels at each word, e.g. "John/PER Doe/PER did yesterday...".
        These labels will be ignored by this function (e.g. only "John" and "Doe" would be used).
        All count values and ranks will be automatically estimated.

        Note: This function is rather slow.

        Args:
            filepath: Filepath to the corpus file.
            verbose: Whether to output messages during parsing.
        """
        self.fill_from_articles_labels(filepath, labels=None, verbose=verbose)

    def fill_from_articles_labels(self, filepath, labels=None, verbose=False):
        """Fills the dictionaries of this class from a corpus file, optionally only with words of
        specific labels (e.g. only words labeled with "PER").

        The corpus file is expected to contain one article/document per line.
        The corpus file may contain labels at each word, e.g. "John/PER Doe/PER did yesterday...".
        All count values and ranks will be automatically estimated.

        Note: This function is rather slow.

        Args:
            filepath: Filepath to the corpus file.
            labels: Optionally one or more labels. If provided, only words that are annotated
                with any of these labels will be counted.
            verbose: Whether to output messages during parsing.
        """
        assert labels is None or isinstance(labels, list)
        assert labels is None or len(labels) > 0

        counts = Counter()
        self.sum_of_counts = 0

        articles = load_articles(filepath, start_at=0)
        for i, article in enumerate(articles):
            words = [token.word for token in article.tokens \
                                if labels is None or token.label in labels]
            counts.update(words)
            if verbose and i % 1000 == 0:
                print("Article %d" % (i))

        most_common = counts.most_common()
        for i, (word, count) in enumerate(most_common):
            self.word_to_count[word] = count
            self.word_to_rank[word] = i + 1
            self.sum_of_counts += count

    def write_to_file(self, filepath):
        """Writes the contents of this unigrams object to a file.
        The file can later on be loaded with fill_from_file().
        Args:
            filepath: Filepath to the file to which to write.
        """
        with open(filepath, "w") as handle:
            for i, (word, count) in enumerate(self.word_to_count.iteritems()):
                print(word, count)
                if i > 0:
                    handle.write("\n")
                handle.write(word.encode("utf-8"))
                handle.write("\t")
                handle.write(str(count))

    def get_rank_of(self, word, default=-1):
        """Returns the rank of a word among all unigrams.
        The most common word has rank 1.
        Args:
            word: The word which to rank.
            default: A default value to return if the word was not found among the unigrams.
        Returns:
            integer or default value (-1).
        """
        if word in self.word_to_rank:
            return self.word_to_rank[word]
        else:
            return default

    def get_count_of(self, word, default=-1):
        """Returns the count of a word among all unigrams.
        Args:
            word: The word which to rank.
            default: A default value to return if the word was not found among the unigrams.
        Returns:
            integer or default value (-1).
        """
        if word in self.word_to_count:
            return self.word_to_count[word]
        else:
            return default

    def get_frequency_of(self, word, default=None):
        """Returns the frequency of a word among all unigrams.
        The frequency is calculated by count(word)/count(all words)
        Args:
            word: The word which to rank.
            default: A default value to return if the word was not found among the unigrams.
        Returns:
            double or default value (None).
        """
        count = self.get_count_of(word, None)
        if count is not None:
            return count / self.sum_of_counts
        else:
            return default
<|endoftext|>


<commit_before>from odict import odict
from plumber import (
    plumber,
    Part,
)
from zope.interface import implements
from zope.interface.common.mapping import (
    IEnumerableMapping,
    IFullMapping,
)
from node.interfaces import IAliaser
from node.utils import ReverseMapping
from node.parts import (
    Adopt,
    Nodify,
    NodeChildValidate,
)

class DictAliaser(odict):
    """Uses its own dictionary for aliasing.

    ``__getitem__`` -> unalias
    """
    implements(IAliaser, IFullMapping)

    def __init__(self, data=(), strict=True):
        super(DictAliaser, self).__init__(data)
        self.strict = strict

    def alias(self, key):
        try:
            return ReverseMapping(self)[key]
        except KeyError, e:
            if not self.strict:
                return key
            raise e

    def unalias(self, aliased_key):
        try:
            return self[aliased_key]
        except KeyError, e:
            if not self.strict:
                return aliased_key
            raise e

class PrefixAliaser(object):
    """An aliaser that prefix all keys.

    As it never raise KeyError it is not whitelisting.
    """
    implements(IAliaser)

    def __init__(self, prefix=None):
        self.prefix = prefix

    def alias(self, key):
        return (self.prefix or '') + key

    def unalias(self, prefixed_key):
        """Returns the real key for a prefixed_key.
        """
        prefix = self.prefix or ''
        if not prefixed_key.startswith(prefix):
            raise KeyError(u"key '%s' does not match prefix '%s'" % \
                    (prefixed_key, prefix))
        return prefixed_key[len(prefix):]

class SuffixAliaser(object):
    """An aliaser that suffixes all keys.

    As it never raise KeyError it is not whitelisting.
    """
    implements(IAliaser)

    def __init__(self, suffix=None):
        self.suffix = suffix

    def alias(self, key):
        return key + (self.suffix or '')

    def unalias(self, suffixed_key):
        """returns the real key for a suffixed_key
        """
        suffix = self.suffix or ''
        if not suffixed_key.endswith(suffix):
            raise KeyError(
                    u"key '%s' does not match suffix '%s'" % \
                            (suffixed_key, suffix)
                    )
        return suffixed_key[:-len(suffix)]

class AliaserChain(object):
    """A chain of aliasers.

    chain = [aliaser1, aliaser2]
    chain.alias(key) == aliaser2.alias(aliaser1.alias(key))
    chain.unalias(alias_key) == aliaser2.unalias(aliaser1.unalias(aliased_key))
    """
    implements(IAliaser)

    # XXX: we are IEnumerableMapping if one of our childs is, which is
    #      important as we become a whitelist, eg. for Node.__iter__

    def __init__(self, chain=None):
        self.chain = chain

    def alias(self, key):
        for aliaser in self.chain:
            key = aliaser.alias(key)
        return key

    def unalias(self, key):
        for aliaser in reversed(self.chain):
            key = aliaser.unalias(key)
        return key

class PrefixSuffixAliaser(AliaserChain):
    """Prefixes and suffixes.
    """

    def __init__(self, prefix=None, suffix=None):
        self.chain = (
            PrefixAliaser(prefix),
            SuffixAliaser(suffix))

class NamedAliasers(dict):
    """A dictionary storing aliasers by name.

    XXX
    """

class AliasedNodespace(object):
    """Performs aliasing/unaliasing for node children.

    Is not the parent of its children, the children don't know about their name
    here.

    Future additional mode: children are wrapped, wrapper knows name and we are
    parent of wrapper.
    """
    __metaclass__ = plumber
    __plumbing__ = (
        NodeChildValidate,
        Adopt,
        Nodify,
    )

    def __init__(self, context, aliaser=None):
        """
        context
            the node whose children to alias

        aliaser
            the aliaser to be used
        """
        # XXX: is just taking over the name ok for all use cases?
        # XXX: case where not? -rn
        #super(AliasedNodespace, self).__init__(context.__name__)

        self.__name__ = context.__name__
        self.__parent__ = None

        self.context = context
        self.aliaser = aliaser

    @property
    def changed(self):
        # XXX: move this out of here. Base AliasedNodespace should not take
        #      care about special contracts. this one is currently used in
        #      node.ext.ldap, subclass AliasedNodespace there.
        return self.context.changed                         #pragma NO COVERAGE

    @property
    def __call__(self):
        # XXX: propably move this out here as well.
        return self.context()                               #pragma NO COVERAGE

    def __delitem__(self, key):
        unaliased_key = self.aliaser and self.aliaser.unalias(key) or key
        try:
            del self.context[unaliased_key]
        except KeyError:
            raise KeyError(key)

    def __getitem__(self, key):
        unaliased_key = self.aliaser and self.aliaser.unalias(key) or key
        try:
            return self.context[unaliased_key]
        except KeyError:
            raise KeyError(key)

    def __setitem__(self, key, val):
        unaliased_key = self.aliaser and self.aliaser.unalias(key) or key
        self.context[unaliased_key] = val

    def __iter__(self):
        for key in self.context:
            try:
                yield self.aliaser and self.aliaser.alias(key) or key
            except KeyError:
                if IEnumerableMapping.providedBy(self.aliaser):
                    # an enumerable aliaser whitelists, we skip non-listed keys
                    continue
                # no whitelisting and a KeyError on our internal data: that's
                # bad! Most probably not triggered on _Node but a subclass
                # XXX: test case showing this
                raise RuntimeError(                         #pragma NO COVERAGE
                    u"Inconsist internal node state")       #pragma NO COVERAGE

    def __repr__(self):
        return "Aliased " + self.context.__repr__()
<commit_msg>remove non storage contract specific properties from AliasedNodespace<commit_after>from odict import odict
from plumber import (
    plumber,
    Part,
)
from zope.interface import implements
from zope.interface.common.mapping import (
    IEnumerableMapping,
    IFullMapping,
)
from node.interfaces import IAliaser
from node.utils import ReverseMapping
from node.parts import (
    Adopt,
    Nodify,
    NodeChildValidate,
)

class DictAliaser(odict):
    """Uses its own dictionary for aliasing.

    ``__getitem__`` -> unalias
    """
    implements(IAliaser, IFullMapping)

    def __init__(self, data=(), strict=True):
        super(DictAliaser, self).__init__(data)
        self.strict = strict

    def alias(self, key):
        try:
            return ReverseMapping(self)[key]
        except KeyError, e:
            if not self.strict:
                return key
            raise e

    def unalias(self, aliased_key):
        try:
            return self[aliased_key]
        except KeyError, e:
            if not self.strict:
                return aliased_key
            raise e

class PrefixAliaser(object):
    """An aliaser that prefix all keys.

    As it never raise KeyError it is not whitelisting.
    """
    implements(IAliaser)

    def __init__(self, prefix=None):
        self.prefix = prefix

    def alias(self, key):
        return (self.prefix or '') + key

    def unalias(self, prefixed_key):
        """Returns the real key for a prefixed_key.
        """
        prefix = self.prefix or ''
        if not prefixed_key.startswith(prefix):
            raise KeyError(u"key '%s' does not match prefix '%s'" % \
                    (prefixed_key, prefix))
        return prefixed_key[len(prefix):]

class SuffixAliaser(object):
    """An aliaser that suffixes all keys.

    As it never raise KeyError it is not whitelisting.
    """
    implements(IAliaser)

    def __init__(self, suffix=None):
        self.suffix = suffix

    def alias(self, key):
        return key + (self.suffix or '')

    def unalias(self, suffixed_key):
        """returns the real key for a suffixed_key
        """
        suffix = self.suffix or ''
        if not suffixed_key.endswith(suffix):
            raise KeyError(
                    u"key '%s' does not match suffix '%s'" % \
                            (suffixed_key, suffix)
                    )
        return suffixed_key[:-len(suffix)]

class AliaserChain(object):
    """A chain of aliasers.

    chain = [aliaser1, aliaser2]
    chain.alias(key) == aliaser2.alias(aliaser1.alias(key))
    chain.unalias(alias_key) == aliaser2.unalias(aliaser1.unalias(aliased_key))
    """
    implements(IAliaser)

    # XXX: we are IEnumerableMapping if one of our childs is, which is
    #      important as we become a whitelist, eg. for Node.__iter__

    def __init__(self, chain=None):
        self.chain = chain

    def alias(self, key):
        for aliaser in self.chain:
            key = aliaser.alias(key)
        return key

    def unalias(self, key):
        for aliaser in reversed(self.chain):
            key = aliaser.unalias(key)
        return key

class PrefixSuffixAliaser(AliaserChain):
    """Prefixes and suffixes.
    """

    def __init__(self, prefix=None, suffix=None):
        self.chain = (
            PrefixAliaser(prefix),
            SuffixAliaser(suffix))

#class NamedAliasers(dict):
#    """A dictionary storing aliasers by name.
#
#    XXX
#    """

class AliasedNodespace(object):
    """Performs aliasing/unaliasing for node children.

    Is not the parent of its children, the children don't know about their name
    here.

    Future additional mode: children are wrapped, wrapper knows name and we are
    parent of wrapper.
    """
    __metaclass__ = plumber
    __plumbing__ = (
        NodeChildValidate,
        Adopt,
        Nodify,
    )

    def __init__(self, context, aliaser=None):
        """
        context
            the node whose children to alias

        aliaser
            the aliaser to be used
        """
        # XXX: is just taking over the name ok for all use cases?
        # XXX: case where not? -rn
        #super(AliasedNodespace, self).__init__(context.__name__)

        self.__name__ = context.__name__
        self.__parent__ = None

        self.context = context
        self.aliaser = aliaser

    def __delitem__(self, key):
        unaliased_key = self.aliaser and self.aliaser.unalias(key) or key
        try:
            del self.context[unaliased_key]
        except KeyError:
            raise KeyError(key)

    def __getitem__(self, key):
        unaliased_key = self.aliaser and self.aliaser.unalias(key) or key
        try:
            return self.context[unaliased_key]
        except KeyError:
            raise KeyError(key)

    def __setitem__(self, key, val):
        unaliased_key = self.aliaser and self.aliaser.unalias(key) or key
        self.context[unaliased_key] = val

    def __iter__(self):
        for key in self.context:
            try:
                yield self.aliaser and self.aliaser.alias(key) or key
            except KeyError:
                if IEnumerableMapping.providedBy(self.aliaser):
                    # an enumerable aliaser whitelists, we skip non-listed keys
                    continue
                # no whitelisting and a KeyError on our internal data: that's
                # bad! Most probably not triggered on _Node but a subclass
                # XXX: test case showing this.
                raise RuntimeError(                         #pragma NO COVERAGE
                    u"Inconsist internal node state")       #pragma NO COVERAGE

    def __repr__(self):
        return "Aliased " + self.context.__repr__()
<|endoftext|>


<commit_before># Copyright (c) 2013-2016 Jeffrey Pfau
#
# This Source Code Form is subject to the terms of the Mozilla Public
# License, v. 2.0. If a copy of the MPL was not distributed with this
# file, You can obtain one at http://mozilla.org/MPL/2.0/.
from ._pylib import ffi, lib

class MemoryView(object):
    def __init__(self, core, width, size, base=0, sign="u"):
        self._core = core
        self._width = width
        self._size = size
        self._base = base
        self._busRead = getattr(self._core, "busRead" + str(width * 8))
        self._busWrite = getattr(self._core, "busWrite" + str(width * 8))
        self._rawRead = getattr(self._core, "rawRead" + str(width * 8))
        self._rawWrite = getattr(self._core, "rawWrite" + str(width * 8))
        self._mask = (1 << (width * 8)) - 1 # Used to force values to fit within range so that negative values work
        if sign == "u" or sign == "unsigned":
            self._type = "uint{}_t".format(width * 8)
        elif sign == "i" or sign == "s" or sign == "signed":
            self._type = "int{}_t".format(width * 8)
        else:
            raise ValueError("Invalid sign type: '{}'".format(sign))

    def _addrCheck(self, address):
        if isinstance(address, slice):
            start = address.start or 0
            stop = self._size - self._width if address.stop is None else address.stop
        else:
            start = address
            stop = address + self._width
        if start >= self._size or stop > self._size:
            raise IndexError()
        if start < 0 or stop < 0:
            raise IndexError()

    def __len__(self):
        return self._size

    def __getitem__(self, address):
        self._addrCheck(address)
        if isinstance(address, slice):
            start = address.start or 0
            stop = self._size - self._width if address.stop is None else address.stop
            step = address.step or self._width
            return [int(ffi.cast(self._type, self._busRead(self._core, self._base + a))) for a in range(start, stop, step)]
        else:
            return int(ffi.cast(self._type, self._busRead(self._core, self._base + address)))

    def __setitem__(self, address, value):
        self._addrCheck(address)
        if isinstance(address, slice):
            start = address.start or 0
            stop = self._size - self._width if address.stop is None else address.stop
            step = address.step or self._width
            for a in range(start, stop, step):
                self._busWrite(self._core, self._base + a, value[a] & self._mask)
        else:
            self._busWrite(self._core, self._base + address, value & self._mask)

    def rawRead(self, address, segment=-1):
        self._addrCheck(address)
        return int(ffi.cast(self._type, self._rawRead(self._core, self._base + address, segment)))

    def rawWrite(self, address, value, segment=-1):
        self._addrCheck(address)
        self._rawWrite(self._core, self._base + address, segment, value & self._mask)

class MemorySearchResult(object):
    def __init__(self, memory, result):
        self.address = result.address
        self.segment = result.segment
        self.guessDivisor = result.guessDivisor
        self.type = result.type

        if result.type == Memory.SEARCH_8:
            self._memory = memory.u8
        elif result.type == Memory.SEARCH_16:
            self._memory = memory.u16
        elif result.type == Memory.SEARCH_32:
            self._memory = memory.u32
        elif result.type == Memory.SEARCH_STRING:
            self._memory = memory.u8
        else:
            raise ValueError("Unknown type: %X" % result.type)

    @property
    def value(self):
        if self.type == Memory.SEARCH_STRING:
            raise ValueError
        return self._memory[self.address] * self.guessDivisor

    @value.setter
    def value(self, v):
        if self.type == Memory.SEARCH_STRING:
            raise IndexError
        self._memory[self.address] = v // self.guessDivisor

class Memory(object):
    SEARCH_32 = lib.mCORE_MEMORY_SEARCH_32
    SEARCH_16 = lib.mCORE_MEMORY_SEARCH_16
    SEARCH_8 = lib.mCORE_MEMORY_SEARCH_8
    SEARCH_STRING = lib.mCORE_MEMORY_SEARCH_STRING
    SEARCH_GUESS = lib.mCORE_MEMORY_SEARCH_GUESS

    READ = lib.mCORE_MEMORY_READ
    WRITE = lib.mCORE_MEMORY_READ
    RW = lib.mCORE_MEMORY_RW

    def __init__(self, core, size, base=0):
        self.size = size
        self.base = base
        self._core = core

        self.u8 = MemoryView(core, 1, size, base, "u")
        self.u16 = MemoryView(core, 2, size, base, "u")
        self.u32 = MemoryView(core, 4, size, base, "u")
        self.s8 = MemoryView(core, 1, size, base, "s")
        self.s16 = MemoryView(core, 2, size, base, "s")
        self.s32 = MemoryView(core, 4, size, base, "s")

    def __len__(self):
        return self._size

    def search(self, value, type=SEARCH_GUESS, flags=RW, limit=10000, old_results=[]):
        results = ffi.new("struct mCoreMemorySearchResults*")
        lib.mCoreMemorySearchResultsInit(results, len(old_results))
        params = ffi.new("struct mCoreMemorySearchParams*")
        params.memoryFlags = flags
        params.type = type
        if type == self.SEARCH_8:
            params.value8 = int(value)
        elif type == self.SEARCH_16:
            params.value16 = int(value)
        elif type == self.SEARCH_32:
            params.value32 = int(value)
        else:
            params.valueStr = ffi.new("char[]", str(value).encode("ascii"))

        for result in old_results:
            r = lib.mCoreMemorySearchResultsAppend(results)
            r.address = result.address
            r.segment = result.segment
            r.guessDivisor = result.guessDivisor
            r.type = result.type
        if old_results:
            lib.mCoreMemorySearchRepeat(self._core, params, results)
        else:
            lib.mCoreMemorySearch(self._core, params, results, limit)
        new_results = [MemorySearchResult(self, lib.mCoreMemorySearchResultsGetPointer(results, i)) for i in range(lib.mCoreMemorySearchResultsSize(results))]
        lib.mCoreMemorySearchResultsDeinit(results)
        return new_results
<commit_msg>Python: Add subscripting to root memory object<commit_after># Copyright (c) 2013-2016 Jeffrey Pfau
#
# This Source Code Form is subject to the terms of the Mozilla Public
# License, v. 2.0. If a copy of the MPL was not distributed with this
# file, You can obtain one at http://mozilla.org/MPL/2.0/.
from ._pylib import ffi, lib

class MemoryView(object):
    def __init__(self, core, width, size, base=0, sign="u"):
        self._core = core
        self._width = width
        self._size = size
        self._base = base
        self._busRead = getattr(self._core, "busRead" + str(width * 8))
        self._busWrite = getattr(self._core, "busWrite" + str(width * 8))
        self._rawRead = getattr(self._core, "rawRead" + str(width * 8))
        self._rawWrite = getattr(self._core, "rawWrite" + str(width * 8))
        self._mask = (1 << (width * 8)) - 1 # Used to force values to fit within range so that negative values work
        if sign == "u" or sign == "unsigned":
            self._type = "uint{}_t".format(width * 8)
        elif sign == "i" or sign == "s" or sign == "signed":
            self._type = "int{}_t".format(width * 8)
        else:
            raise ValueError("Invalid sign type: '{}'".format(sign))

    def _addrCheck(self, address):
        if isinstance(address, slice):
            start = address.start or 0
            stop = self._size - self._width if address.stop is None else address.stop
        else:
            start = address
            stop = address + self._width
        if start >= self._size or stop > self._size:
            raise IndexError()
        if start < 0 or stop < 0:
            raise IndexError()

    def __len__(self):
        return self._size

    def __getitem__(self, address):
        self._addrCheck(address)
        if isinstance(address, slice):
            start = address.start or 0
            stop = self._size - self._width if address.stop is None else address.stop
            step = address.step or self._width
            return [int(ffi.cast(self._type, self._busRead(self._core, self._base + a))) for a in range(start, stop, step)]
        else:
            return int(ffi.cast(self._type, self._busRead(self._core, self._base + address)))

    def __setitem__(self, address, value):
        self._addrCheck(address)
        if isinstance(address, slice):
            start = address.start or 0
            stop = self._size - self._width if address.stop is None else address.stop
            step = address.step or self._width
            for a in range(start, stop, step):
                self._busWrite(self._core, self._base + a, value[a] & self._mask)
        else:
            self._busWrite(self._core, self._base + address, value & self._mask)

    def rawRead(self, address, segment=-1):
        self._addrCheck(address)
        return int(ffi.cast(self._type, self._rawRead(self._core, self._base + address, segment)))

    def rawWrite(self, address, value, segment=-1):
        self._addrCheck(address)
        self._rawWrite(self._core, self._base + address, segment, value & self._mask)

class MemorySearchResult(object):
    def __init__(self, memory, result):
        self.address = result.address
        self.segment = result.segment
        self.guessDivisor = result.guessDivisor
        self.type = result.type

        if result.type == Memory.SEARCH_8:
            self._memory = memory.u8
        elif result.type == Memory.SEARCH_16:
            self._memory = memory.u16
        elif result.type == Memory.SEARCH_32:
            self._memory = memory.u32
        elif result.type == Memory.SEARCH_STRING:
            self._memory = memory.u8
        else:
            raise ValueError("Unknown type: %X" % result.type)

    @property
    def value(self):
        if self.type == Memory.SEARCH_STRING:
            raise ValueError
        return self._memory[self.address] * self.guessDivisor

    @value.setter
    def value(self, v):
        if self.type == Memory.SEARCH_STRING:
            raise IndexError
        self._memory[self.address] = v // self.guessDivisor

class Memory(object):
    SEARCH_32 = lib.mCORE_MEMORY_SEARCH_32
    SEARCH_16 = lib.mCORE_MEMORY_SEARCH_16
    SEARCH_8 = lib.mCORE_MEMORY_SEARCH_8
    SEARCH_STRING = lib.mCORE_MEMORY_SEARCH_STRING
    SEARCH_GUESS = lib.mCORE_MEMORY_SEARCH_GUESS

    READ = lib.mCORE_MEMORY_READ
    WRITE = lib.mCORE_MEMORY_READ
    RW = lib.mCORE_MEMORY_RW

    def __init__(self, core, size, base=0):
        self.size = size
        self.base = base
        self._core = core

        self.u8 = MemoryView(core, 1, size, base, "u")
        self.u16 = MemoryView(core, 2, size, base, "u")
        self.u32 = MemoryView(core, 4, size, base, "u")
        self.s8 = MemoryView(core, 1, size, base, "s")
        self.s16 = MemoryView(core, 2, size, base, "s")
        self.s32 = MemoryView(core, 4, size, base, "s")

    def __len__(self):
        return self._size

    def search(self, value, type=SEARCH_GUESS, flags=RW, limit=10000, old_results=[]):
        results = ffi.new("struct mCoreMemorySearchResults*")
        lib.mCoreMemorySearchResultsInit(results, len(old_results))
        params = ffi.new("struct mCoreMemorySearchParams*")
        params.memoryFlags = flags
        params.type = type
        if type == self.SEARCH_8:
            params.value8 = int(value)
        elif type == self.SEARCH_16:
            params.value16 = int(value)
        elif type == self.SEARCH_32:
            params.value32 = int(value)
        else:
            params.valueStr = ffi.new("char[]", str(value).encode("ascii"))

        for result in old_results:
            r = lib.mCoreMemorySearchResultsAppend(results)
            r.address = result.address
            r.segment = result.segment
            r.guessDivisor = result.guessDivisor
            r.type = result.type
        if old_results:
            lib.mCoreMemorySearchRepeat(self._core, params, results)
        else:
            lib.mCoreMemorySearch(self._core, params, results, limit)
        new_results = [MemorySearchResult(self, lib.mCoreMemorySearchResultsGetPointer(results, i)) for i in range(lib.mCoreMemorySearchResultsSize(results))]
        lib.mCoreMemorySearchResultsDeinit(results)
        return new_results

    def __getitem__(self, address):
        if isinstance(address, slice):
            return bytearray(self.u8[address])
        else:
            return self.u8[address]
<|endoftext|>


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

"""
Sith attacks (and helps debugging) Jedi.

Randomly search Python files and run Jedi on it.  Exception and used
arguments are recorded to ``./record.json`` (specified by --record)::

    %(prog)s random /path/to/sourcecode

Redo recorded exception::

    %(prog)s redo

Fallback to pdb when error is raised::

    %(prog)s --pdb random
    %(prog)s --pdb redo

"""

from __future__ import print_function
import json
import os
import random
import sys
import traceback
import argparse

try:
    from itertools import izip as zip
except ImportError:
    pass

import jedi

class SourceCode(object):

    def __init__(self, path):
        self.path = path
        with open(path) as f:
            self.source = f.read()
        self.lines = self.source.splitlines()
        self.maxline = len(self.lines)

    def choose_script_args(self):
        line = random.randint(1, self.maxline)
        column = random.randint(0, len(self.lines[line - 1]))
        return (self.source, line, column, self.path)

class SourceFinder(object):

    def __init__(self, rootpath):
        self.rootpath = rootpath
        self.files = list(self.search_files())

    def search_files(self):
        for root, dirnames, filenames in os.walk(self.rootpath):
            for name in filenames:
                if name.endswith('.py'):
                    yield os.path.join(root, name)

    def choose_source(self):
        # FIXME: try same file for several times
        return SourceCode(random.choice(self.files))

class BaseAttacker(object):

    def __init__(self):
        self.record = {'data': []}

    def attack(self, operation, *args):
        script = jedi.Script(*args)
        op = getattr(script, operation)
        op()

    def add_record(self, exc_info, operation, args):
        (_type, value, tb) = exc_info
        self.record['data'].append({
            'traceback': traceback.format_tb(tb),
            'error': repr(value),
            'operation': operation,
            'args': args,
        })

    def get_record(self, recid):
        return self.record['data'][recid]

    def save_record(self, path):
        with open(path, 'w') as f:
            json.dump(self.record, f)

    def load_record(self, path):
        with open(path) as f:
            self.record = json.load(f)
            return self.record

    def add_arguments(self, parser):
        parser.set_defaults(func=self.do_run)

    def get_help(self):
        for line in self.__doc__.splitlines():
            line = line.strip()
            if line:
                return line

class MixinPrinter(object):

    def print_record(self, recid=-1):
        data = self.get_record(recid)
        print(*data['traceback'], end='')
        print("""
{error} is raised by running Script(...).{operation}() with
line  : {args[1]}
column: {args[2]}
path  : {args[3]}
""".format(**data))

class MixinLoader(object):

    def add_arguments(self, parser):
        super(MixinLoader, self).add_arguments(parser)
        parser = parser.add_argument(
            'recid', default=0, nargs='?', type=int, help="""
            This option currently has no effect as random attack record
            only one error.
            """)

    def do_run(self, record, recid):
        self.load_record(record)

class AttackReporter(object):

    def __init__(self):
        self.tries = 0
        self.errors = 0

    def __iter__(self):
        return self

    def __next__(self):
        self.tries += 1
        sys.stderr.write('.')
        sys.stderr.flush()
        return self.tries

    next = __next__

    def error(self):
        self.errors += 1
        sys.stderr.write('\n')
        sys.stderr.flush()
        print('{0}th error is encountered after {1} tries.'
              .format(self.errors, self.tries))

class RandomAtaccker(MixinPrinter, BaseAttacker):

    """
    Randomly run Script().<method>() against files under <rootpath>.
    """

    operations = [
        'completions', 'goto_assignments', 'goto_definitions', 'usages',
        'call_signatures']

    def choose_operation(self):
        return random.choice(self.operations)

    def generate_attacks(self, maxtries, finder):
        for _ in range(maxtries):
            src = finder.choose_source()
            operation = self.choose_operation()
            yield (operation, src.choose_script_args())

    def do_run(self, record, rootpath, maxtries):
        finder = SourceFinder(rootpath)
        reporter = AttackReporter()
        for (operation, args) in self.generate_attacks(maxtries, finder):
            reporter.next()
            try:
                self.attack(operation, *args)
            except jedi.NotFoundError:
                pass
            except Exception:
                self.add_record(sys.exc_info(), operation, args)
                reporter.error()
                self.print_record()
                raise
            finally:
                self.save_record(record)

    def add_arguments(self, parser):
        super(RandomAtaccker, self).add_arguments(parser)
        parser.add_argument(
            '--maxtries', '-l', default=10000, type=int)
        parser.add_argument(
            'rootpath', default='.', nargs='?',
            help='root directory to look for Python files.')

class RedoAttacker(MixinLoader, BaseAttacker):

    """
    Redo recorded attack.
    """

    def do_run(self, record, recid):
        super(RedoAttacker, self).do_run(record, recid)
        data = self.get_record(recid)
        self.attack(data['operation'], *data['args'])

class ShowRecord(MixinLoader, MixinPrinter, BaseAttacker):

    """
    Show recorded errors.
    """

    def do_run(self, record, recid):
        super(ShowRecord, self).do_run(record, recid)
        self.print_record()

class AttackApp(object):

    def __init__(self):
        self.parsers = []
        self.attackers = []

    def run(self, args=None):
        parser = self.get_parser()
        self.do_run(**vars(parser.parse_args(args)))

    def do_run(self, func, debugger, **kwds):
        try:
            func(**kwds)
        except:
            exc_info = sys.exc_info()
            if debugger == 'pdb':
                import pdb
                pdb.post_mortem(exc_info[2])
            elif debugger == 'ipdb':
                import ipdb
                ipdb.post_mortem(exc_info[2])

    def add_parser(self, attacker_class, *args, **kwds):
        attacker = attacker_class()
        parser = self.subparsers.add_parser(
            *args,
            help=attacker.get_help(),
            description=attacker.__doc__,
            **kwds)
        attacker.add_arguments(parser)

        # Not required, just fore debugging:
        self.parsers.append(parser)
        self.attackers.append(attacker)

    def get_parser(self):
        parser = argparse.ArgumentParser(
            formatter_class=argparse.RawDescriptionHelpFormatter,
            description=__doc__)
        parser.add_argument(
            '--record', '-R', default='record.json',
            help='Exceptions are recorded in here (default: %(default)s).')
        parser.add_argument(
            '--pdb', dest='debugger', const='pdb', action='store_const',
            help="Launch pdb when error is raised.")
        parser.add_argument(
            '--ipdb', dest='debugger', const='ipdb', action='store_const',
            help="Launch ipdb when error is raised.")

        self.subparsers = parser.add_subparsers()
        self.add_parser(RandomAtaccker, 'random')
        self.add_parser(RedoAttacker, 'redo')
        self.add_parser(ShowRecord, 'show')

        return parser

if __name__ == '__main__':
    app = AttackApp()
    app.run()
<commit_msg>Removed unused import from sith.py<commit_after>#!/usr/bin/env python

"""
Sith attacks (and helps debugging) Jedi.

Randomly search Python files and run Jedi on it.  Exception and used
arguments are recorded to ``./record.json`` (specified by --record)::

    %(prog)s random /path/to/sourcecode

Redo recorded exception::

    %(prog)s redo

Fallback to pdb when error is raised::

    %(prog)s --pdb random
    %(prog)s --pdb redo

"""

from __future__ import print_function
import json
import os
import random
import sys
import traceback
import argparse

import jedi

class SourceCode(object):

    def __init__(self, path):
        self.path = path
        with open(path) as f:
            self.source = f.read()
        self.lines = self.source.splitlines()
        self.maxline = len(self.lines)

    def choose_script_args(self):
        line = random.randint(1, self.maxline)
        column = random.randint(0, len(self.lines[line - 1]))
        return (self.source, line, column, self.path)

class SourceFinder(object):

    def __init__(self, rootpath):
        self.rootpath = rootpath
        self.files = list(self.search_files())

    def search_files(self):
        for root, dirnames, filenames in os.walk(self.rootpath):
            for name in filenames:
                if name.endswith('.py'):
                    yield os.path.join(root, name)

    def choose_source(self):
        # FIXME: try same file for several times
        return SourceCode(random.choice(self.files))

class BaseAttacker(object):

    def __init__(self):
        self.record = {'data': []}

    def attack(self, operation, *args):
        script = jedi.Script(*args)
        op = getattr(script, operation)
        op()

    def add_record(self, exc_info, operation, args):
        (_type, value, tb) = exc_info
        self.record['data'].append({
            'traceback': traceback.format_tb(tb),
            'error': repr(value),
            'operation': operation,
            'args': args,
        })

    def get_record(self, recid):
        return self.record['data'][recid]

    def save_record(self, path):
        with open(path, 'w') as f:
            json.dump(self.record, f)

    def load_record(self, path):
        with open(path) as f:
            self.record = json.load(f)
            return self.record

    def add_arguments(self, parser):
        parser.set_defaults(func=self.do_run)

    def get_help(self):
        for line in self.__doc__.splitlines():
            line = line.strip()
            if line:
                return line

class MixinPrinter(object):

    def print_record(self, recid=-1):
        data = self.get_record(recid)
        print(*data['traceback'], end='')
        print("""
{error} is raised by running Script(...).{operation}() with
line  : {args[1]}
column: {args[2]}
path  : {args[3]}
""".format(**data))

class MixinLoader(object):

    def add_arguments(self, parser):
        super(MixinLoader, self).add_arguments(parser)
        parser = parser.add_argument(
            'recid', default=0, nargs='?', type=int, help="""
            This option currently has no effect as random attack record
            only one error.
            """)

    def do_run(self, record, recid):
        self.load_record(record)

class AttackReporter(object):

    def __init__(self):
        self.tries = 0
        self.errors = 0

    def __iter__(self):
        return self

    def __next__(self):
        self.tries += 1
        sys.stderr.write('.')
        sys.stderr.flush()
        return self.tries

    next = __next__

    def error(self):
        self.errors += 1
        sys.stderr.write('\n')
        sys.stderr.flush()
        print('{0}th error is encountered after {1} tries.'
              .format(self.errors, self.tries))

class RandomAtaccker(MixinPrinter, BaseAttacker):

    """
    Randomly run Script().<method>() against files under <rootpath>.
    """

    operations = [
        'completions', 'goto_assignments', 'goto_definitions', 'usages',
        'call_signatures']

    def choose_operation(self):
        return random.choice(self.operations)

    def generate_attacks(self, maxtries, finder):
        for _ in range(maxtries):
            src = finder.choose_source()
            operation = self.choose_operation()
            yield (operation, src.choose_script_args())

    def do_run(self, record, rootpath, maxtries):
        finder = SourceFinder(rootpath)
        reporter = AttackReporter()
        for (operation, args) in self.generate_attacks(maxtries, finder):
            reporter.next()
            try:
                self.attack(operation, *args)
            except jedi.NotFoundError:
                pass
            except Exception:
                self.add_record(sys.exc_info(), operation, args)
                reporter.error()
                self.print_record()
                raise
            finally:
                self.save_record(record)

    def add_arguments(self, parser):
        super(RandomAtaccker, self).add_arguments(parser)
        parser.add_argument(
            '--maxtries', '-l', default=10000, type=int)
        parser.add_argument(
            'rootpath', default='.', nargs='?',
            help='root directory to look for Python files.')

class RedoAttacker(MixinLoader, BaseAttacker):

    """
    Redo recorded attack.
    """

    def do_run(self, record, recid):
        super(RedoAttacker, self).do_run(record, recid)
        data = self.get_record(recid)
        self.attack(data['operation'], *data['args'])

class ShowRecord(MixinLoader, MixinPrinter, BaseAttacker):

    """
    Show recorded errors.
    """

    def do_run(self, record, recid):
        super(ShowRecord, self).do_run(record, recid)
        self.print_record()

class AttackApp(object):

    def __init__(self):
        self.parsers = []
        self.attackers = []

    def run(self, args=None):
        parser = self.get_parser()
        self.do_run(**vars(parser.parse_args(args)))

    def do_run(self, func, debugger, **kwds):
        try:
            func(**kwds)
        except:
            exc_info = sys.exc_info()
            if debugger == 'pdb':
                import pdb
                pdb.post_mortem(exc_info[2])
            elif debugger == 'ipdb':
                import ipdb
                ipdb.post_mortem(exc_info[2])

    def add_parser(self, attacker_class, *args, **kwds):
        attacker = attacker_class()
        parser = self.subparsers.add_parser(
            *args,
            help=attacker.get_help(),
            description=attacker.__doc__,
            **kwds)
        attacker.add_arguments(parser)

        # Not required, just fore debugging:
        self.parsers.append(parser)
        self.attackers.append(attacker)

    def get_parser(self):
        parser = argparse.ArgumentParser(
            formatter_class=argparse.RawDescriptionHelpFormatter,
            description=__doc__)
        parser.add_argument(
            '--record', '-R', default='record.json',
            help='Exceptions are recorded in here (default: %(default)s).')
        parser.add_argument(
            '--pdb', dest='debugger', const='pdb', action='store_const',
            help="Launch pdb when error is raised.")
        parser.add_argument(
            '--ipdb', dest='debugger', const='ipdb', action='store_const',
            help="Launch ipdb when error is raised.")

        self.subparsers = parser.add_subparsers()
        self.add_parser(RandomAtaccker, 'random')
        self.add_parser(RedoAttacker, 'redo')
        self.add_parser(ShowRecord, 'show')

        return parser

if __name__ == '__main__':
    app = AttackApp()
    app.run()
<|endoftext|>