<commit_before>import pytest
from pytorch_lightning import Trainer
from pytorch_lightning.examples.new_project_templates.lightning_module_template import LightningTemplateModel
from argparse import Namespace
from test_tube import Experiment
from pytorch_lightning.callbacks import ModelCheckpoint
import numpy as np
import warnings
import torch
import os
import shutil

SEED = 2334
torch.manual_seed(SEED)
np.random.seed(SEED)

# ------------------------------------------------------------------------
# TESTS
# ------------------------------------------------------------------------
def test_cpu_model():
    """
    Make sure model trains on CPU
    :return:
    """
    save_dir = init_save_dir()

    model, hparams = get_model()

    trainer = Trainer(
        progress_bar=False,
        experiment=get_exp(),
        max_nb_epochs=1,
        train_percent_check=0.4,
        val_percent_check=0.4
    )
    result = trainer.fit(model)

    # correct result and ok accuracy
    assert result == 1, 'cpu model failed to complete'
    assert_ok_acc(trainer)

    clear_save_dir()

def test_single_gpu_model():
    """
    Make sure single GPU works (DP mode)
    :return:
    """
    if not torch.cuda.is_available():
        warnings.warn('test_single_gpu_model cannot run. Rerun on a GPU node to run this test')
        return

    trainer_options = dict(
        progress_bar=False,
        max_nb_epochs=1,
        train_percent_check=0.1,
        val_percent_check=0.1,
        gpus=[0]
    )

    run_gpu_model_test(trainer_options)

def test_multi_gpu_model_dp():
    """
    Make sure DP works
    :return:
    """
    if not torch.cuda.is_available():
        warnings.warn('test_multi_gpu_model_dp cannot run. Rerun on a GPU node to run this test')
        return
    if not torch.cuda.device_count() > 1:
        warnings.warn('test_multi_gpu_model_dp cannot run. Rerun on a node with 2+ GPUs to run this test')
        return

    trainer_options = dict(
        progress_bar=False,
        max_nb_epochs=1,
        train_percent_check=0.1,
        val_percent_check=0.1,
        gpus=[0, 1]
    )

    run_gpu_model_test(trainer_options)

def test_amp_gpu_dp():
    """
    Make sure DP + AMP work
    :return:
    """
    if not torch.cuda.is_available():
        warnings.warn('test_amp_gpu_dp cannot run. Rerun on a GPU node to run this test')
        return
    if not torch.cuda.device_count() > 1:
        warnings.warn('test_amp_gpu_dp cannot run. Rerun on a node with 2+ GPUs to run this test')
        return

    trainer_options = dict(
        progress_bar=False,
        max_nb_epochs=1,
        gpus=[0, 1],
        distributed_backend='dp',
        use_amp=True
    )

    run_gpu_model_test(trainer_options)

def test_multi_gpu_model_ddp():
    """
    Make sure DDP works
    :return:
    """
    if not torch.cuda.is_available():
        warnings.warn('test_multi_gpu_model_ddp cannot run. Rerun on a GPU node to run this test')
        return
    if not torch.cuda.device_count() > 1:
        warnings.warn('test_multi_gpu_model_ddp cannot run. Rerun on a node with 2+ GPUs to run this test')
        return

    trainer_options = dict(
        progress_bar=False,
        max_nb_epochs=1,
        train_percent_check=0.1,
        val_percent_check=0.1,
        gpus=[0, 1],
        distributed_backend='ddp'
    )

    run_gpu_model_test(trainer_options)

def test_amp_gpu_ddp():
    """
    Make sure DDP + AMP work
    :return:
    """
    if not torch.cuda.is_available():
        warnings.warn('test_amp_gpu_ddp cannot run. Rerun on a GPU node to run this test')
        return
    if not torch.cuda.device_count() > 1:
        warnings.warn('test_amp_gpu_ddp cannot run. Rerun on a node with 2+ GPUs to run this test')
        return

    trainer_options = dict(
        progress_bar=True,
        max_nb_epochs=1,
        gpus=[0, 1],
        distributed_backend='ddp',
        use_amp=True
    )

    run_gpu_model_test(trainer_options)

# ------------------------------------------------------------------------
# UTILS
# ------------------------------------------------------------------------

def run_gpu_model_test(trainer_options):
    """
        Make sure DDP + AMP work
        :return:
        """
    if not torch.cuda.is_available():
        warnings.warn('test_amp_gpu_ddp cannot run. Rerun on a GPU node to run this test')
        return
    if not torch.cuda.device_count() > 1:
        warnings.warn('test_amp_gpu_ddp cannot run. Rerun on a node with 2+ GPUs to run this test')
        return

    save_dir = init_save_dir()
    model, hparams = get_model()

    # exp file to get meta
    exp = get_exp(False)
    exp.argparse(hparams)
    exp.save()

    # exp file to get weights
    checkpoint = ModelCheckpoint(save_dir)

    # add these to the trainer options
    trainer_options['checkpoint_callback'] = checkpoint
    trainer_options['experiment'] = exp

    # fit model
    trainer = Trainer(**trainer_options)
    result = trainer.fit(model)

    # correct result and ok accuracy
    assert result == 1, 'amp + ddp model failed to complete'

    # test model loading
    pretrained_model = load_model(exp, save_dir)

    # test model preds
    run_prediction(model.test_dataloader, pretrained_model)

    clear_save_dir()

def get_model():
    # set up model with these hyperparams
    root_dir = os.path.dirname(os.path.realpath(__file__))
    hparams = Namespace(**{'drop_prob': 0.2,
                           'batch_size': 32,
                           'in_features': 28*28,
                           'learning_rate': 0.001*8,
                           'optimizer_name': 'adam',
                           'data_root': os.path.join(root_dir, 'mnist'),
                           'out_features': 10,
                           'hidden_dim': 1000})
    model = LightningTemplateModel(hparams)

    return model, hparams

def get_exp(debug=True):
    # set up exp object without actually saving logs
    root_dir = os.path.dirname(os.path.realpath(__file__))
    exp = Experiment(debug=debug, save_dir=root_dir, name='tests_tt_dir')
    return exp

def init_save_dir():
    root_dir = os.path.dirname(os.path.realpath(__file__))
    save_dir = os.path.join(root_dir, 'save_dir')

    if os.path.exists(save_dir):
        shutil.rmtree(save_dir)

    os.makedirs(save_dir, exist_ok=True)

    return save_dir

def clear_save_dir():
    root_dir = os.path.dirname(os.path.realpath(__file__))
    save_dir = os.path.join(root_dir, 'save_dir')
    if os.path.exists(save_dir):
        shutil.rmtree(save_dir)

def load_model(exp, save_dir):

    # load trained model
    tags_path = exp.get_data_path(exp.name, exp.version)
    tags_path = os.path.join(tags_path, 'meta_tags.csv')

    checkpoints = [x for x in os.listdir(save_dir) if '.ckpt' in x]
    weights_dir = os.path.join(save_dir, checkpoints[0])

    trained_model = LightningTemplateModel.load_from_metrics(weights_path=weights_dir, tags_csv=tags_path, on_gpu=True)

    assert trained_model is not None, 'loading model failed'

    return trained_model

def run_prediction(dataloader, trained_model):
    # run prediction on 1 batch
    for batch in dataloader:
        break

    x, y = batch
    x = x.view(x.size(0), -1)

    y_hat = trained_model(x)

    # acc
    labels_hat = torch.argmax(y_hat, dim=1)
    val_acc = torch.sum(y == labels_hat).item() / (len(y) * 1.0)
    val_acc = torch.tensor(val_acc)
    val_acc = val_acc.item()

    print(val_acc)

    assert val_acc > 0.70, f'this model is expected to get > 0.7 in test set (it got {val_acc})'

def assert_ok_acc(trainer):
    # this model should get 0.80+ acc
    acc = trainer.tng_tqdm_dic['val_acc']
    assert acc > 0.70, f'model failed to get expected 0.70 validation accuracy. Got: {acc}'

if __name__ == '__main__':
    pytest.main([__file__])
<commit_msg>refactored model tests<commit_after>import pytest
from pytorch_lightning import Trainer
from pytorch_lightning.examples.new_project_templates.lightning_module_template import LightningTemplateModel
from argparse import Namespace
from test_tube import Experiment
from pytorch_lightning.callbacks import ModelCheckpoint
import numpy as np
import warnings
import torch
import os
import shutil

SEED = 2334
torch.manual_seed(SEED)
np.random.seed(SEED)

# ------------------------------------------------------------------------
# TESTS
# ------------------------------------------------------------------------
def test_cpu_model():
    """
    Make sure model trains on CPU
    :return:
    """
    save_dir = init_save_dir()

    model, hparams = get_model()

    trainer = Trainer(
        progress_bar=False,
        experiment=get_exp(),
        max_nb_epochs=1,
        train_percent_check=0.4,
        val_percent_check=0.4
    )
    result = trainer.fit(model)

    # correct result and ok accuracy
    assert result == 1, 'cpu model failed to complete'
    assert_ok_acc(trainer)

    clear_save_dir()

def test_single_gpu_model():
    """
    Make sure single GPU works (DP mode)
    :return:
    """
    if not torch.cuda.is_available():
        warnings.warn('test_single_gpu_model cannot run. Rerun on a GPU node to run this test')
        return

    trainer_options = dict(
        progress_bar=False,
        max_nb_epochs=1,
        train_percent_check=0.1,
        val_percent_check=0.1,
        gpus=[0]
    )

    run_gpu_model_test(trainer_options)

def test_multi_gpu_model_dp():
    """
    Make sure DP works
    :return:
    """
    if not torch.cuda.is_available():
        warnings.warn('test_multi_gpu_model_dp cannot run. Rerun on a GPU node to run this test')
        return
    if not torch.cuda.device_count() > 1:
        warnings.warn('test_multi_gpu_model_dp cannot run. Rerun on a node with 2+ GPUs to run this test')
        return

    trainer_options = dict(
        progress_bar=False,
        max_nb_epochs=1,
        train_percent_check=0.1,
        val_percent_check=0.1,
        gpus=[0, 1]
    )

    run_gpu_model_test(trainer_options)

def test_amp_gpu_dp():
    """
    Make sure DP + AMP work
    :return:
    """
    if not torch.cuda.is_available():
        warnings.warn('test_amp_gpu_dp cannot run. Rerun on a GPU node to run this test')
        return
    if not torch.cuda.device_count() > 1:
        warnings.warn('test_amp_gpu_dp cannot run. Rerun on a node with 2+ GPUs to run this test')
        return

    trainer_options = dict(
        progress_bar=False,
        max_nb_epochs=1,
        gpus=[0, 1],
        distributed_backend='dp',
        use_amp=True
    )

    run_gpu_model_test(trainer_options)

def test_multi_gpu_model_ddp():
    """
    Make sure DDP works
    :return:
    """
    if not torch.cuda.is_available():
        warnings.warn('test_multi_gpu_model_ddp cannot run. Rerun on a GPU node to run this test')
        return
    if not torch.cuda.device_count() > 1:
        warnings.warn('test_multi_gpu_model_ddp cannot run. Rerun on a node with 2+ GPUs to run this test')
        return

    trainer_options = dict(
        progress_bar=False,
        max_nb_epochs=1,
        train_percent_check=0.1,
        val_percent_check=0.1,
        gpus=[0, 1],
        distributed_backend='ddp'
    )

    run_gpu_model_test(trainer_options)

def test_amp_gpu_ddp():
    """
    Make sure DDP + AMP work
    :return:
    """
    if not torch.cuda.is_available():
        warnings.warn('test_amp_gpu_ddp cannot run. Rerun on a GPU node to run this test')
        return
    if not torch.cuda.device_count() > 1:
        warnings.warn('test_amp_gpu_ddp cannot run. Rerun on a node with 2+ GPUs to run this test')
        return

    trainer_options = dict(
        progress_bar=True,
        max_nb_epochs=1,
        gpus=[0, 1],
        distributed_backend='ddp',
        use_amp=True
    )

    run_gpu_model_test(trainer_options)

# ------------------------------------------------------------------------
# UTILS
# ------------------------------------------------------------------------

def run_gpu_model_test(trainer_options):
    """
        Make sure DDP + AMP work
        :return:
        """
    if not torch.cuda.is_available():
        warnings.warn('test_amp_gpu_ddp cannot run. Rerun on a GPU node to run this test')
        return
    if not torch.cuda.device_count() > 1:
        warnings.warn('test_amp_gpu_ddp cannot run. Rerun on a node with 2+ GPUs to run this test')
        return

    save_dir = init_save_dir()
    model, hparams = get_model()

    # exp file to get meta
    exp = get_exp(False)
    exp.argparse(hparams)
    exp.save()

    # exp file to get weights
    checkpoint = ModelCheckpoint(save_dir)

    # add these to the trainer options
    trainer_options['checkpoint_callback'] = checkpoint
    trainer_options['experiment'] = exp

    # fit model
    trainer = Trainer(**trainer_options)
    result = trainer.fit(model)

    # correct result and ok accuracy
    assert result == 1, 'amp + ddp model failed to complete'

    # test model loading
    pretrained_model = load_model(exp, save_dir)

    # test model preds
    run_prediction(model.test_dataloader, pretrained_model)

    clear_save_dir()

def get_model():
    # set up model with these hyperparams
    root_dir = os.path.dirname(os.path.realpath(__file__))
    hparams = Namespace(**{'drop_prob': 0.2,
                           'batch_size': 32,
                           'in_features': 28*28,
                           'learning_rate': 0.001*8,
                           'optimizer_name': 'adam',
                           'data_root': os.path.join(root_dir, 'mnist'),
                           'out_features': 10,
                           'hidden_dim': 1000})
    model = LightningTemplateModel(hparams)

    return model, hparams

def get_exp(debug=True):
    # set up exp object without actually saving logs
    root_dir = os.path.dirname(os.path.realpath(__file__))
    exp = Experiment(debug=debug, save_dir=root_dir, name='tests_tt_dir')
    return exp

def init_save_dir():
    root_dir = os.path.dirname(os.path.realpath(__file__))
    save_dir = os.path.join(root_dir, 'save_dir')

    if os.path.exists(save_dir):
        shutil.rmtree(save_dir)

    os.makedirs(save_dir, exist_ok=True)

    return save_dir

def clear_save_dir():
    root_dir = os.path.dirname(os.path.realpath(__file__))
    save_dir = os.path.join(root_dir, 'save_dir')
    if os.path.exists(save_dir):
        shutil.rmtree(save_dir)

def load_model(exp, save_dir):

    # load trained model
    tags_path = exp.get_data_path(exp.name, exp.version)
    tags_path = os.path.join(tags_path, 'meta_tags.csv')

    checkpoints = [x for x in os.listdir(save_dir) if '.ckpt' in x]
    weights_dir = os.path.join(save_dir, checkpoints[0])

    trained_model = LightningTemplateModel.load_from_metrics(weights_path=weights_dir, tags_csv=tags_path, on_gpu=True)

    assert trained_model is not None, 'loading model failed'

    return trained_model

def run_prediction(dataloader, trained_model):
    # run prediction on 1 batch
    for batch in dataloader:
        break

    x, y = batch
    x = x.view(x.size(0), -1)

    y_hat = trained_model(x)

    # acc
    labels_hat = torch.argmax(y_hat, dim=1)
    val_acc = torch.sum(y == labels_hat).item() / (len(y) * 1.0)
    val_acc = torch.tensor(val_acc)
    val_acc = val_acc.item()

    print(val_acc)

    assert val_acc > 0.70, f'this model is expected to get > 0.7 in test set (it got {val_acc})'

def assert_ok_acc(trainer):
    # this model should get 0.80+ acc
    acc = trainer.tng_tqdm_dic['val_acc']
    assert acc > 0.60, f'model failed to get expected 0.60 validation accuracy. Got: {acc}'

if __name__ == '__main__':
    pytest.main([__file__])
<|endoftext|>


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

# This program is free software; you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation; either version 2 of the License, or
# (at your option) any later version.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with this program; if not, write to the Free Software
# Foundation, Inc., 59 Temple Place - Suite 330, Boston, MA 02111-1307, USA.
#
# Copyright (C) Albert Kottke, 2013-2015

import os
import pytest

from numpy.testing import assert_almost_equal, assert_equal

from pysra import motion

@pytest.fixture
def tsm():
    '''Setup the default time series for testing.'''
    return motion.TimeSeriesMotion.load_at2_file(
            os.path.join(os.path.dirname(__file__), 'data', 'NIS090.AT2'))

def test_ts_load_at2_file(tsm):
    assert_equal(tsm.accels.size, 4096)
    assert_almost_equal(tsm.time_step, 0.01)

    assert_almost_equal(tsm.accels[0], 0.233833E-06)
    assert_almost_equal(tsm.accels[-1], 0.496963E-04)

def test_ts_freqs(tsm):
    freqs = tsm.freqs
    assert_almost_equal(freqs[0], 0)
    assert_almost_equal(freqs[-1], 50.)

    assert_equal(tsm.freqs.size, tsm.fourier_amps.size)
<commit_msg>Fixed PEP8 issue<commit_after>#!/usr/bin/env python
# encoding: utf-8

# This program is free software; you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation; either version 2 of the License, or
# (at your option) any later version.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with this program; if not, write to the Free Software
# Foundation, Inc., 59 Temple Place - Suite 330, Boston, MA 02111-1307, USA.
#
# Copyright (C) Albert Kottke, 2013-2015

import os
import pytest

from numpy.testing import assert_almost_equal, assert_equal

from pysra import motion

@pytest.fixture
def tsm():
    '''Setup the default time series for testing.'''
    return motion.TimeSeriesMotion.load_at2_file(
        os.path.join(os.path.dirname(__file__), 'data', 'NIS090.AT2'))

def test_ts_load_at2_file(tsm):
    assert_equal(tsm.accels.size, 4096)
    assert_almost_equal(tsm.time_step, 0.01)

    assert_almost_equal(tsm.accels[0], 0.233833E-06)
    assert_almost_equal(tsm.accels[-1], 0.496963E-04)

def test_ts_freqs(tsm):
    freqs = tsm.freqs
    assert_almost_equal(freqs[0], 0)
    assert_almost_equal(freqs[-1], 50.)

    assert_equal(tsm.freqs.size, tsm.fourier_amps.size)
<|endoftext|>


<commit_before>import unittest

import time

from robot.robot import Robot
from robot.tests.mockrobotd import MockRobotD

class MotorBoardTest(unittest.TestCase):
    def setUp(self):
        mock = MockRobotD(root_dir="/tmp/")
        self.mock = mock
        self.robot = Robot(robotd_path="/tmp/robotd")

    def test_insert_motorboards(self):
        self.mock.new_motorboard('ABC')
        self.mock.new_motorboard('DEF')
        # Give it a tiny bit to init the boards
        time.sleep(0.4)

        boards = self.robot.motor_boards

        # Check all the motor boards are initialised and can be indexed
        self.assertTrue(0 in boards)
        self.assertTrue(1 in boards)
        self.assertTrue('ABC' in boards)
        self.assertTrue('DEF' in boards)

    def tearDown(self):
        self.mock.stop()

<commit_msg>fix build error<commit_after>import unittest

import time

from robot.robot import Robot
from robot.tests.mock_robotd import MockRobotD

class MotorBoardTest(unittest.TestCase):
    def setUp(self):
        mock = MockRobotD(root_dir="/tmp/")
        self.mock = mock
        self.robot = Robot(robotd_path="/tmp/robotd")

    def test_insert_motorboards(self):
        self.mock.new_motorboard('ABC')
        self.mock.new_motorboard('DEF')
        # Give it a tiny bit to init the boards
        time.sleep(0.4)

        boards = self.robot.motor_boards

        # Check all the motor boards are initialised and can be indexed
        self.assertTrue(0 in boards)
        self.assertTrue(1 in boards)
        self.assertTrue('ABC' in boards)
        self.assertTrue('DEF' in boards)

    def tearDown(self):
        self.mock.stop()

<|endoftext|>


<commit_before>import numbers
from operator import mul, add
from functools import reduce

# multiplicative compliment to sum
prod = lambda l: reduce(mul, l)

# if a fraction or number
isNumeric = lambda n: isinstance(n, (Fraction, numbers.Number))

# coerce to a polynomial
poly = lambda p: P.make_constant(p) if isNumeric(p) else p

# converts to int if float value is equal to int value
possibleInt = lambda n: int(n) if int(n) == float(n) else n

class Polynomial(object):

    def __init__(self, terms=[], *_terms, var = "x"):
        """
        Polynomial can be called with a single list of multiple points, or
        with multiple arguments, where each argument is a point.

        >>> Polynomial([])
        Traceback (most recent call last):
        ...
        AssertionError: No terms provided.
        >>> Polynomial((3, 2), (2, 0))
        3x^2 + 2
        >>> Polynomial([(3, 2), (2, 0)])
        3x^2 + 2
        """
        assert len(terms) or len(_terms), 'No terms provided.'
        assert all([isinstance(t, tuple) and len(t) == 2 for t in _terms]), \
            "Terms must be tuples of the form (coefficient, exponent)."
        assert isinstance(var, str), "Variable must be a string."
        __terms = terms if hasattr(terms[0], '__iter__') else (terms,) + _terms
        self.termsList = self.unsimplified_terms = __terms
        self.combine_terms()
        self.sort()
        self.var = var

    @property
    def terms(self):
        return self.termsList

    def copy(self):
        return Polynomial(self.terms,self.var)

    def sort(self):
        self.termsList = sorted(self.terms, key = lambda x: x[1],reverse = True)
        return self

    @staticmethod
    def make_constant(c):
        """creates a Polynomial representation of a constant"""
        assert isNumeric(c), 'Non-numerics cannot be coerced to a constant.'
        return Polynomial((c, 0))

    def is_constant(self):
        """tests if polynomial is a constant"""
        return len(self.terms) == 1 and self.terms[0][1] == 0

    def combine_terms(self):
        """Combine terms, by summing all coefficients of the same degree

        >>> P((3, 3), (3, 5), (2, 5)).combine_terms().termsList
        [(3, 3), (5, 5)]
        """
        combined = {}
        for coeff, exp in self.terms:
            combined.setdefault(exp, []).append(coeff)
        self.termsList = [(sum(v), k) for k, v in combined.items()]
        return self

    def clean(self):
        """Removes all terms with coefficient 0

        >>> P((3, 5), (0, 5), (3, 1)).clean().termsList
        [(3, 5), (3, 1)]
        """
        self.termsList = list(filter(lambda term: term[0], self.termsList))
        return self

    def __add__(p1, p2):
        """Implements addition for Polynomial object

        >>> P((1, 0)) + P((2, 1), (2, 0))
        2x + 3
        """
        return Polynomial(poly(p1).terms + poly(p2).terms)

    def __sub__(p1, p2):
        """Implements subtraction for Polynomial object

        >>> P((1, 0)) - P((2, 1), (2, 0))
        -2x - 1
        """
        return p1 + -p2

    def __neg__(p):
        """Implements negation for Polynomial object
        >>> -P((1, 0))
        -1
        """
        return p * -1

    def __mul__(p1, p2):
        """Implements multiplication for Polynomial object. Note: Multiplication
        by non-Polynomial objects, like integers or floats, can only be
        accomplished in one format: Polynomial * non-Polynomial

        >>> P((1, 0)) * -5
        -5
        """
        multiply = lambda p1, p2: (p1[0] * p2[0], p1[1] + p2[1])
        return Polynomial(sum([[multiply(p1, p2) for p1 in poly(p1).terms]
            for p2 in poly(p2).terms], []))

    def __div__(p1, p2):
        raise NotImplemented()

    def __float__(p1):
        if Polynomial.is_constant(p1):
            return float(p1.terms[0][0])
        raise TypeError('Non-constant polynomial cannot be coerced to float.')

    def __repr__(self):
        return str(self)

    def __str__(self):
        """Stringifies polynomials

        >>> P((3, 2))
        3x^2
        >>> P((3, 1)) # no caret (^), if exponent is 1
        3x
        >>> P((3, 0)) # no variable if exponent is 0
        3
        >>> P((1, 1)) # no coefficient if coefficient is 1
        x
        >>> P((0, 2), (1, 1)) # no terms with coefficient 0
        x
        """
        def stringify(term):
            coefficient, exponent = term
            exp = {
                0: '',
                1: self.var
            }.get(exponent, '%s^%d' % (self.var, exponent))
            cof = '' if coefficient == 1 else str(possibleInt(coefficient))
            return '%s%s' % (cof, exp)
        string = ''
        for t in map(stringify, filter(lambda t: t[0], self.terms)):
            if not string:
                string += t
            elif t[0] == '-':
                string += ' - %s' % t[1:]
            else:
                string += ' + %s' % t
        return string

class Fraction(object):
    """Representation of a Fraction, accepts Fractions and Polynomials as
    numerators and denominators, in addition to regular numerical expressions
    in Python.
    """

    def __init__(self, num, den):
        assert den != 0
        assert isNumeric(num) or isinstance(num, Polynomial)
        assert isNumeric(den) or isinstance(den, Polynomial)
#CHECK FOR SIMPLIFIYING FRACTIONS
        self.num = poly(num)
        self.den = poly(den)
        self.reduceConstants()
        self.simplify()

    def copy(self):
        return Fraction(self.num, self.den)

    def __add__(f1, f2):
        assert isNumeric(f1) and isNumeric(f2)
        if not isinstance(f2, Fraction):
            return Fraction(f1.num * f2, f1.den)
        return Fraction(f1.num * f2.den + f2.num * f1.den, f1.den * f2.den)

    def __mul__(f1, f2):
        assert isNumeric(f1) and isNumeric(f2)
        if not isinstance(f2, Fraction):
            return Fraction(f1.num * f2, f1.den)
        return Fraction(f1.num * f2.num, f1.den * f2.den)

    def __div__(f1,f2):
        assert isNumeric(f1) and isNumeric(f2)
        return f1 * f2.copy().invert()

    def reciprocal(self):
        """Returns a new Fraction, reciprocal of the original"""
        return Fraction(self.den, self.num)

    def invert(self):
        """Inverts existing Fraction, in-place"""
        self.num, self.den = self.den, self.num
        return self

    def reduceConstants(self):
        constants = [x[0] for x in self.num.terms] + [x[0] for x in self.den.terms]
        def gcd_help(x,y):
            x,y = map(abs, (x,y))
            if y > x:
                return gcd_help(y,x)
            if y == 0:
                return x
            return gcd_help(y, x % y)
        gcd = reduce(gcd_help, constants)
        self.num = Polynomial([(x[0]*1.0/gcd , x[1]) for x in self.num.terms])
        self.den = Polynomial([(x[0]*1.0/gcd , x[1]) for x in self.den.terms])
        return self

    def simplify(self):
        """Takes several steps to simplify fractions:
        1. Rids of negative denominators
        2. Divides by gcd of terms in numerator and terms in denominator
        """
        if self.den.is_constant() and float(self.den) == -1:
            self.den, self.num = poly(1), -self.num

    def __repr__(self):
        return str(self)

    def __str__(self):
        """String representation of a Fraction"""
        if poly(self.den).is_constant() and float(self.den) == 1:
            return str(self.num)
        return "(%s)/(%s)" % (str(self.num), str(self.den))

def interpolate(*points):
    """Takes in a list of tuples of x,y cordinates.

    >>> ipt = lambda y, d: interpolate(*[(x, y(x)) for x in range(d+1)])
    >>> ipt(lambda x: x**5 + 3*x**3 + x, 5)
    x^5 + 3x^3 + x
    >>> ipt(lambda x: 7*x**4 + 5*x**3 + 9*x**2 + 3, 4)
    7x^4 + 5x^3 + 9x^2 + 3
    """
    points = set(points)
    return reduce(add, [
        Fraction(
            prod(P((1, 1), (-x2, 0)) for x2, y2 in points if (x,y) != (x2,y2)), prod(P((x,0), (-x2, 0)) for x2, y2 in points if (x,y) != (x2,y2))
        ) * y for x, y in points])

# abbreviations
P = Polynomial
<commit_msg>prettier fractions<commit_after>import numbers
from operator import mul, add
from functools import reduce

# multiplicative compliment to sum
prod = lambda l: reduce(mul, l)

# if a fraction or number
isNumeric = lambda n: isinstance(n, (Fraction, numbers.Number))

# coerce to a polynomial
poly = lambda p: P.make_constant(p) if isNumeric(p) else p

# converts to int if float value is equal to int value
pint = lambda n: int(n) if int(n) == float(n) else n

class Polynomial(object):

    def __init__(self, terms=[], *_terms, var = "x"):
        """
        Polynomial can be called with a single list of multiple points, or
        with multiple arguments, where each argument is a point.

        >>> Polynomial([])
        Traceback (most recent call last):
        ...
        AssertionError: No terms provided.
        >>> Polynomial((3, 2), (2, 0))
        3x^2 + 2
        >>> Polynomial([(3, 2), (2, 0)])
        3x^2 + 2
        """
        assert len(terms) or len(_terms), 'No terms provided.'
        assert all([isinstance(t, tuple) and len(t) == 2 for t in _terms]), \
            "Terms must be tuples of the form (coefficient, exponent)."
        assert isinstance(var, str), "Variable must be a string."
        __terms = terms if hasattr(terms[0], '__iter__') else (terms,) + _terms
        self.termsList = self.unsimplified_terms = __terms
        self.combine_terms()
        self.sort()
        self.var = var

    @property
    def terms(self):
        return self.termsList

    def copy(self):
        return Polynomial(self.terms,self.var)

    def sort(self):
        self.termsList = sorted(self.terms, key = lambda x: x[1],reverse = True)
        return self

    @staticmethod
    def make_constant(c):
        """creates a Polynomial representation of a constant"""
        assert isNumeric(c), 'Non-numerics cannot be coerced to a constant.'
        return Polynomial((c, 0))

    def is_constant(self):
        """tests if polynomial is a constant"""
        return len(self.terms) == 1 and self.terms[0][1] == 0

    def combine_terms(self):
        """Combine terms, by summing all coefficients of the same degree

        >>> P((3, 3), (3, 5), (2, 5)).combine_terms().termsList
        [(3, 3), (5, 5)]
        """
        combined = {}
        for coeff, exp in self.terms:
            combined.setdefault(exp, []).append(coeff)
        self.termsList = [(sum(v), k) for k, v in combined.items()]
        return self

    def clean(self):
        """Removes all terms with coefficient 0

        >>> P((3, 5), (0, 5), (3, 1)).clean().termsList
        [(3, 5), (3, 1)]
        """
        self.termsList = list(filter(lambda term: term[0], self.termsList))
        return self

    def __add__(p1, p2):
        """Implements addition for Polynomial object

        >>> P((1, 0)) + P((2, 1), (2, 0))
        2x + 3
        """
        return Polynomial(poly(p1).terms + poly(p2).terms)

    def __sub__(p1, p2):
        """Implements subtraction for Polynomial object

        >>> P((1, 0)) - P((2, 1), (2, 0))
        -2x - 1
        """
        return p1 + -p2

    def __neg__(p):
        """Implements negation for Polynomial object
        >>> -P((1, 0))
        -1
        """
        return p * -1

    def __mul__(p1, p2):
        """Implements multiplication for Polynomial object. Note: Multiplication
        by non-Polynomial objects, like integers or floats, can only be
        accomplished in one format: Polynomial * non-Polynomial

        >>> P((1, 0)) * -5
        -5
        """
        multiply = lambda p1, p2: (p1[0] * p2[0], p1[1] + p2[1])
        return Polynomial(sum([[multiply(p1, p2) for p1 in poly(p1).terms]
            for p2 in poly(p2).terms], []))

    def __div__(p1, p2):
        """Implements division for Polynomial object"""
        raise NotImplemented()

    def __float__(p1):
        """Implements coercion to floating point for Polynomial objects"""
        if Polynomial.is_constant(p1):
            return float(p1.terms[0][0])
        raise TypeError('Non-constant polynomial cannot be coerced to float.')

    def __getitem__(self, i):
        """Implements indexing for Polynomial objects

        >>> P((1, 0), (5, 3), (3, 2))[0]
        5x^3
        """
        self.sort()
        return Polynomial(self.terms[i])

    def __repr__(self):
        return str(self)

    def __str__(self):
        """Stringifies polynomials

        >>> P((3, 2))
        3x^2
        >>> P((3, 1)) # no caret (^), if exponent is 1
        3x
        >>> P((3, 0)) # no variable if exponent is 0
        3
        >>> P((1, 1)) # no coefficient if coefficient is 1
        x
        >>> P((0, 2), (1, 1)) # no terms with coefficient 0
        x
        """
        def stringify(term):
            coefficient, exponent = term
            exp = {
                0: '',
                1: self.var
            }.get(exponent, '%s^%d' % (self.var, exponent))
            cof = '' if coefficient == 1 else str(pint(coefficient))
            return '%s%s' % (cof, exp)
        string = ''
        for t in map(stringify, filter(lambda t: t[0], self.terms)):
            if not string:
                string += t
            elif t[0] == '-':
                string += ' - %s' % t[1:]
            else:
                string += ' + %s' % t
        return string

class Fraction(object):
    """Representation of a Fraction, accepts Fractions and Polynomials as
    numerators and denominators, in addition to regular numerical expressions
    in Python.
    """

    def __init__(self, num, den, sigfigs=2):
        assert den != 0
        assert isNumeric(num) or isinstance(num, Polynomial)
        assert isNumeric(den) or isinstance(den, Polynomial)
#CHECK FOR SIMPLIFIYING FRACTIONS
        self.num = poly(num)
        self.den = poly(den)
        self.sigfigs = sigfigs
        self.reduceConstants()
        self.simplify()

    def copy(self):
        return Fraction(self.num, self.den)

    def __add__(f1, f2):
        """Implements addition for Fraction object"""
        assert isNumeric(f1) and isNumeric(f2)
        if not isinstance(f2, Fraction):
            return Fraction(f1.num * f2, f1.den)
        return Fraction(f1.num * f2.den + f2.num * f1.den, f1.den * f2.den)

    def __mul__(f1, f2):
        """Implements multiplication for Fraction object"""
        assert isNumeric(f1) and isNumeric(f2)
        if not isinstance(f2, Fraction):
            return Fraction(f1.num * f2, f1.den)
        return Fraction(f1.num * f2.num, f1.den * f2.den)

    def __div__(f1,f2):
        """Implements division for Fraction object"""
        assert isNumeric(f1) and isNumeric(f2)
        return f1 * f2.copy().invert()

    def reciprocal(self):
        """Returns a new Fraction, reciprocal of the original"""
        return Fraction(self.den, self.num)

    def invert(self):
        """Inverts existing Fraction, in-place"""
        self.num, self.den = self.den, self.num
        return self

    def reduceConstants(self):
        constants = [x[0] for x in self.num.terms] + [x[0] for x in self.den.terms]
        def gcd_help(x,y):
            x,y = map(abs, (x,y))
            if y > x:
                return gcd_help(y,x)
            if y == 0:
                return x
            return gcd_help(y, x % y)
        gcd = int(reduce(gcd_help, constants)) or 1
        self.num = Polynomial([(x[0]*1.0/gcd , x[1]) for x in self.num.terms])
        self.den = Polynomial([(x[0]*1.0/gcd , x[1]) for x in self.den.terms])
        return self

    def simplify(self):
        """Takes several steps to simplify fractions:
        1. Rids of negative denominators
        2. Divides by gcd of terms in numerator and terms in denominator
        """
        if self.den.is_constant() and float(self.den) == -1:
            self.den, self.num = poly(1), -self.num

    def __repr__(self):
        return str(self)

    def __str__(self):
        """String representation of a Fraction

        >>> Fraction(1, 3)
        1/3
        >>> Fraction(1.4, 2)
        1.4/2
        >>> Fraction(poly(1), poly(3))
        1/3
        >>> Fraction(P((5, 3)), 3)
        5x^3/3
        """
        num, den = self.num, self.den
        pnum, pden = poly(num), poly(den)
        num_is_single, num_is_constant = len(pnum.terms) == 1,pnum.is_constant()
        den_is_single, den_is_constant = len(pden.terms) == 1,pden.is_constant()
        if den_is_constant and float(den) == 1:
            return str(num)
        if num_is_single and den_is_single:
            num = pint(float(num)) if num_is_constant else num[0]
            den = pint(float(den)) if den_is_constant else den[0]
            return "%s/%s" % (str(num), str(den))
        return "(%s)/(%s)" % (str(num), str(den))

def interpolate(*points):
    """Takes in a list of tuples of x,y cordinates.

    >>> ipt = lambda y, d: interpolate(*[(x, y(x)) for x in range(d+1)])
    >>> ipt(lambda x: x**5 + 3*x**3 + x, 5)
    x^5 + 3x^3 + x
    >>> ipt(lambda x: 7*x**4 + 5*x**3 + 9*x**2 + 3, 4)
    7x^4 + 5x^3 + 9x^2 + 3
    """
    points = set(points)
    return reduce(add, [
        Fraction(
            prod(P((1, 1), (-x2, 0)) for x2, y2 in points if (x,y) != (x2,y2)),
            prod(P((x,0), (-x2, 0)) for x2, y2 in points if (x,y) != (x2,y2))
        ) * y for x, y in points])

# abbreviations
P = Polynomial
<|endoftext|>


<commit_before># -*- coding: utf-8 -*-
import pytest
from plexapi.exceptions import BadRequest
from . import conftest as utils

def test_myplex_accounts(account, plex):
    assert account, 'Must specify username, password & resource to run this test.'
    print('MyPlexAccount:')
    print('username: %s' % account.username)
    print('email: %s' % account.email)
    print('home: %s' % account.home)
    print('queueEmail: %s' % account.queueEmail)
    assert account.username, 'Account has no username'
    assert account.authenticationToken, 'Account has no authenticationToken'
    assert account.email, 'Account has no email'
    assert account.home is not None, 'Account has no home'
    assert account.queueEmail, 'Account has no queueEmail'
    account = plex.account()
    print('Local PlexServer.account():')
    print('username: %s' % account.username)
    #print('authToken: %s' % account.authToken)
    print('signInState: %s' % account.signInState)
    assert account.username, 'Account has no username'
    assert account.authToken, 'Account has no authToken'
    assert account.signInState, 'Account has no signInState'

def test_myplex_resources(account):
    assert account, 'Must specify username, password & resource to run this test.'
    resources = account.resources()
    for resource in resources:
        name = resource.name or 'Unknown'
        connections = [c.uri for c in resource.connections]
        connections = ', '.join(connections) if connections else 'None'
        print('%s (%s): %s' % (name, resource.product, connections))
    assert resources, 'No resources found for account: %s' % account.name

def test_myplex_connect_to_resource(plex, account):
    servername = plex.friendlyName
    for resource in account.resources():
        if resource.name == servername:
            break
    assert resource.connect(timeout=10)

def test_myplex_devices(account):
    devices = account.devices()
    for device in devices:
        name = device.name or 'Unknown'
        connections = ', '.join(device.connections) if device.connections else 'None'
        print('%s (%s): %s' % (name, device.product, connections))
    assert devices, 'No devices found for account: %s' % account.name

def test_myplex_device(account):
    assert account.device('pkkid-plexapi')

def _test_myplex_connect_to_device(account):
    devices = account.devices()
    for device in devices:
        if device.name == 'some client name' and len(device.connections):
            break
    client = device.connect()
    assert client, 'Unable to connect to device'

def test_myplex_users(account):
    users = account.users()
    assert users, 'Found no users on account: %s' % account.name
    print('Found %s users.' % len(users))
    user = account.user(users[0].title)
    print('Found user: %s' % user)
    assert user, 'Could not find user %s' % users[0].title

    assert len(users[0].servers[0].sections()) == 7, "Could'nt info about the shared libraries"

def test_myplex_resource(account):
    assert account.resource('pkkid-plexapi')

def test_myplex_webhooks(account):
    with pytest.raises(BadRequest):
        account.webhooks()

def test_myplex_addwebhooks(account):
    with pytest.raises(BadRequest):
        account.addWebhook('http://site.com')

def test_myplex_deletewebhooks(account):
    with pytest.raises(BadRequest):
        account.deleteWebhook('http://site.com')

def test_myplex_optout(account):
    def enabled():
        ele = account.query('https://plex.tv/api/v2/user/privacy')
        lib = ele.attrib.get('optOutLibraryStats')
        play = ele.attrib.get('optOutPlayback')
        return bool(int(lib)), bool(int(play))

    # This should be False False
    library_enabled, playback_enabled = enabled()

    account.optOut(library=True, playback=True)

    assert all(enabled())

    account.optOut(library=False, playback=False)

    assert not all(enabled())

def test_myplex_inviteFriend_remove(account, plex, mocker):
    inv_user = 'hellowlol'
    vid_filter = {'contentRating': ['G'], 'label': ['foo']}
    secs = plex.library.sections()

    ids = account._getSectionIds(plex.machineIdentifier, secs)
    with mocker.patch.object(account, '_getSectionIds', return_value=ids):
        with utils.callable_http_patch(mocker):

            account.inviteFriend(inv_user, plex, secs, allowSync=True, allowCameraUpload=True,
                         allowChannels=False, filterMovies=vid_filter, filterTelevision=vid_filter,
                         filterMusic={'label': ['foo']})

        assert inv_user not in [u.title for u in account.users()]

        with utils.callable_http_patch(mock):
            account.removeFriend(inv_user)
<commit_msg>fix damn typo<commit_after># -*- coding: utf-8 -*-
import pytest
from plexapi.exceptions import BadRequest
from . import conftest as utils

def test_myplex_accounts(account, plex):
    assert account, 'Must specify username, password & resource to run this test.'
    print('MyPlexAccount:')
    print('username: %s' % account.username)
    print('email: %s' % account.email)
    print('home: %s' % account.home)
    print('queueEmail: %s' % account.queueEmail)
    assert account.username, 'Account has no username'
    assert account.authenticationToken, 'Account has no authenticationToken'
    assert account.email, 'Account has no email'
    assert account.home is not None, 'Account has no home'
    assert account.queueEmail, 'Account has no queueEmail'
    account = plex.account()
    print('Local PlexServer.account():')
    print('username: %s' % account.username)
    #print('authToken: %s' % account.authToken)
    print('signInState: %s' % account.signInState)
    assert account.username, 'Account has no username'
    assert account.authToken, 'Account has no authToken'
    assert account.signInState, 'Account has no signInState'

def test_myplex_resources(account):
    assert account, 'Must specify username, password & resource to run this test.'
    resources = account.resources()
    for resource in resources:
        name = resource.name or 'Unknown'
        connections = [c.uri for c in resource.connections]
        connections = ', '.join(connections) if connections else 'None'
        print('%s (%s): %s' % (name, resource.product, connections))
    assert resources, 'No resources found for account: %s' % account.name

def test_myplex_connect_to_resource(plex, account):
    servername = plex.friendlyName
    for resource in account.resources():
        if resource.name == servername:
            break
    assert resource.connect(timeout=10)

def test_myplex_devices(account):
    devices = account.devices()
    for device in devices:
        name = device.name or 'Unknown'
        connections = ', '.join(device.connections) if device.connections else 'None'
        print('%s (%s): %s' % (name, device.product, connections))
    assert devices, 'No devices found for account: %s' % account.name

def test_myplex_device(account):
    assert account.device('pkkid-plexapi')

def _test_myplex_connect_to_device(account):
    devices = account.devices()
    for device in devices:
        if device.name == 'some client name' and len(device.connections):
            break
    client = device.connect()
    assert client, 'Unable to connect to device'

def test_myplex_users(account):
    users = account.users()
    assert users, 'Found no users on account: %s' % account.name
    print('Found %s users.' % len(users))
    user = account.user(users[0].title)
    print('Found user: %s' % user)
    assert user, 'Could not find user %s' % users[0].title

    assert len(users[0].servers[0].sections()) == 7, "Could'nt info about the shared libraries"

def test_myplex_resource(account):
    assert account.resource('pkkid-plexapi')

def test_myplex_webhooks(account):
    with pytest.raises(BadRequest):
        account.webhooks()

def test_myplex_addwebhooks(account):
    with pytest.raises(BadRequest):
        account.addWebhook('http://site.com')

def test_myplex_deletewebhooks(account):
    with pytest.raises(BadRequest):
        account.deleteWebhook('http://site.com')

def test_myplex_optout(account):
    def enabled():
        ele = account.query('https://plex.tv/api/v2/user/privacy')
        lib = ele.attrib.get('optOutLibraryStats')
        play = ele.attrib.get('optOutPlayback')
        return bool(int(lib)), bool(int(play))

    # This should be False False
    library_enabled, playback_enabled = enabled()

    account.optOut(library=True, playback=True)

    assert all(enabled())

    account.optOut(library=False, playback=False)

    assert not all(enabled())

def test_myplex_inviteFriend_remove(account, plex, mocker):
    inv_user = 'hellowlol'
    vid_filter = {'contentRating': ['G'], 'label': ['foo']}
    secs = plex.library.sections()

    ids = account._getSectionIds(plex.machineIdentifier, secs)
    with mocker.patch.object(account, '_getSectionIds', return_value=ids):
        with utils.callable_http_patch(mocker):

            account.inviteFriend(inv_user, plex, secs, allowSync=True, allowCameraUpload=True,
                         allowChannels=False, filterMovies=vid_filter, filterTelevision=vid_filter,
                         filterMusic={'label': ['foo']})

        assert inv_user not in [u.title for u in account.users()]

        with utils.callable_http_patch(mocker):
            account.removeFriend(inv_user)
<|endoftext|>


<commit_before># --------------------------------------------------------------------------- #
#   test_nextid.py                                                            #
#                                                                             #
#   Copyright © 2015-2022, Rajiv Bakulesh Shah, original author.              #
#                                                                             #
#   Licensed under the Apache License, Version 2.0 (the "License");           #
#   you may not use this file except in compliance with the License.          #
#   You may obtain a copy of the License at:                                  #
#       http://www.apache.org/licenses/LICENSE-2.0                            #
#                                                                             #
#   Unless required by applicable law or agreed to in writing, software       #
#   distributed under the License is distributed on an "AS IS" BASIS,         #
#   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.  #
#   See the License for the specific language governing permissions and       #
#   limitations under the License.                                            #
# --------------------------------------------------------------------------- #
'Distributed Redis-powered monotonically increasing ID generator tests.'

import unittest.mock
from typing import Generator

import pytest
from redis import Redis
from redis.commands.core import Script
from redis.exceptions import TimeoutError

from pottery import NextID
from pottery import QuorumIsImpossible
from pottery import QuorumNotAchieved

@pytest.fixture
def ids(redis: Redis) -> Generator[NextID, None, None]:
    redis.unlink('nextid:current')
    yield NextID(masters={redis})
    redis.unlink('nextid:current')

def test_nextid(ids: NextID) -> None:
    for id_ in range(1, 10):
        assert next(ids) == id_

def test_iter(ids: NextID) -> None:
    assert iter(ids) is ids

def test_reset(ids: NextID) -> None:
    ids.reset()
    for redis in ids.masters:
        assert not redis.exists(ids.key)

    assert next(ids) == 1
    for redis in ids.masters:
        assert redis.exists(ids.key)

    ids.reset()
    for redis in ids.masters:
        assert not redis.exists(ids.key)

    assert next(ids) == 1
    for redis in ids.masters:
        assert redis.exists(ids.key)

def test_repr(ids: NextID) -> None:
    assert repr(ids) == '<NextID key=nextid:current>'

def test_slots(ids: NextID) -> None:
    with pytest.raises(AttributeError):
        ids.__dict__

def test_next_quorumnotachieved(ids: NextID) -> None:
    with pytest.raises(QuorumNotAchieved), \
         unittest.mock.patch.object(next(iter(ids.masters)), 'get') as get:
        get.side_effect = TimeoutError
        next(ids)

    with pytest.raises(QuorumNotAchieved), \
         unittest.mock.patch.object(Script, '__call__') as __call__:
        __call__.side_effect = TimeoutError
        next(ids)

def test_next_quorumisimpossible(redis: Redis) -> None:
    ids = NextID(masters={redis}, raise_on_redis_errors=True)
    with pytest.raises(QuorumIsImpossible), \
         unittest.mock.patch.object(next(iter(ids.masters)), 'get') as get:
        get.side_effect = TimeoutError
        next(ids)

    with pytest.raises(QuorumIsImpossible), \
         unittest.mock.patch.object(Script, '__call__') as __call__:
        __call__.side_effect = TimeoutError
        next(ids)

def test_reset_quorumnotachieved(ids: NextID) -> None:
    with pytest.raises(QuorumNotAchieved), \
         unittest.mock.patch.object(next(iter(ids.masters)), 'delete') as delete:
        delete.side_effect = TimeoutError
        ids.reset()

def test_reset_quorumisimpossible(redis: Redis) -> None:
    ids = NextID(masters={redis}, raise_on_redis_errors=True)
    with pytest.raises(QuorumIsImpossible), \
         unittest.mock.patch.object(next(iter(ids.masters)), 'delete') as delete:
        delete.side_effect = TimeoutError
        ids.reset()
<commit_msg>Test NextID contention against 5 Redis masters (#680)<commit_after># --------------------------------------------------------------------------- #
#   test_nextid.py                                                            #
#                                                                             #
#   Copyright © 2015-2022, Rajiv Bakulesh Shah, original author.              #
#                                                                             #
#   Licensed under the Apache License, Version 2.0 (the "License");           #
#   you may not use this file except in compliance with the License.          #
#   You may obtain a copy of the License at:                                  #
#       http://www.apache.org/licenses/LICENSE-2.0                            #
#                                                                             #
#   Unless required by applicable law or agreed to in writing, software       #
#   distributed under the License is distributed on an "AS IS" BASIS,         #
#   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.  #
#   See the License for the specific language governing permissions and       #
#   limitations under the License.                                            #
# --------------------------------------------------------------------------- #
'Distributed Redis-powered monotonically increasing ID generator tests.'

import concurrent.futures
import contextlib
import unittest.mock
from typing import Generator

import pytest
from redis import Redis
from redis.commands.core import Script
from redis.exceptions import TimeoutError

from pottery import NextID
from pottery import QuorumIsImpossible
from pottery import QuorumNotAchieved

@pytest.fixture
def ids(redis: Redis) -> Generator[NextID, None, None]:
    redis.unlink('nextid:current')
    yield NextID(masters={redis})
    redis.unlink('nextid:current')

def test_nextid(ids: NextID) -> None:
    for id_ in range(1, 10):
        assert next(ids) == id_

def test_iter(ids: NextID) -> None:
    assert iter(ids) is ids

def test_reset(ids: NextID) -> None:
    ids.reset()
    for redis in ids.masters:
        assert not redis.exists(ids.key)

    assert next(ids) == 1
    for redis in ids.masters:
        assert redis.exists(ids.key)

    ids.reset()
    for redis in ids.masters:
        assert not redis.exists(ids.key)

    assert next(ids) == 1
    for redis in ids.masters:
        assert redis.exists(ids.key)

@pytest.mark.parametrize('num_ids', range(1, 6))
def test_contention(num_ids: int) -> None:
    dbs = range(1, 6)
    urls = [f'redis://localhost:6379/{db}' for db in dbs]
    masters = [Redis.from_url(url, socket_timeout=1) for url in urls]
    ids = [NextID(key='tweet-ids', masters=masters) for _ in range(num_ids)]

    try:
        results = []
        with concurrent.futures.ThreadPoolExecutor() as executor:
            futures = [executor.submit(next, i) for i in ids]
            for future in concurrent.futures.as_completed(futures):
                with contextlib.suppress(QuorumNotAchieved):
                    result = future.result()
                    results.append(result)
        assert len(results) == len(set(results))
        # To see the following output, issue:
        # $ source venv/bin/activate; pytest -rP tests/test_nextid.py::test_contention; deactivate
        print(f'{num_ids} ids, {results} IDs')

    finally:
        ids[0].reset()

def test_repr(ids: NextID) -> None:
    assert repr(ids) == '<NextID key=nextid:current>'

def test_slots(ids: NextID) -> None:
    with pytest.raises(AttributeError):
        ids.__dict__

def test_next_quorumnotachieved(ids: NextID) -> None:
    with pytest.raises(QuorumNotAchieved), \
         unittest.mock.patch.object(next(iter(ids.masters)), 'get') as get:
        get.side_effect = TimeoutError
        next(ids)

    with pytest.raises(QuorumNotAchieved), \
         unittest.mock.patch.object(Script, '__call__') as __call__:
        __call__.side_effect = TimeoutError
        next(ids)

def test_next_quorumisimpossible(redis: Redis) -> None:
    ids = NextID(masters={redis}, raise_on_redis_errors=True)
    with pytest.raises(QuorumIsImpossible), \
         unittest.mock.patch.object(next(iter(ids.masters)), 'get') as get:
        get.side_effect = TimeoutError
        next(ids)

    with pytest.raises(QuorumIsImpossible), \
         unittest.mock.patch.object(Script, '__call__') as __call__:
        __call__.side_effect = TimeoutError
        next(ids)

def test_reset_quorumnotachieved(ids: NextID) -> None:
    with pytest.raises(QuorumNotAchieved), \
         unittest.mock.patch.object(next(iter(ids.masters)), 'delete') as delete:
        delete.side_effect = TimeoutError
        ids.reset()

def test_reset_quorumisimpossible(redis: Redis) -> None:
    ids = NextID(masters={redis}, raise_on_redis_errors=True)
    with pytest.raises(QuorumIsImpossible), \
         unittest.mock.patch.object(next(iter(ids.masters)), 'delete') as delete:
        delete.side_effect = TimeoutError
        ids.reset()
<|endoftext|>


<commit_before>import operator

import py

from rply import ParserGenerator, Token, ParsingError
from rply.token import SourcePosition
from rply.errors import ParserGeneratorWarning

from .base import BaseTests
from .utils import FakeLexer, RecordingLexer, BoxInt, ParserState

class TestBasic(BaseTests):
    def test_simple(self):
        pg = ParserGenerator(["VALUE"])

        @pg.production("main : VALUE")
        def main(p):
            return p[0]

        parser = pg.build()

        assert parser.parse(FakeLexer([Token("VALUE", "abc")])) == Token("VALUE", "abc")

    def test_arithmetic(self):
        pg = ParserGenerator(["NUMBER", "PLUS"])

        @pg.production("main : expr")
        def main(p):
            return p[0]

        @pg.production("expr : expr PLUS expr")
        def expr_op(p):
            return BoxInt(p[0].getint() + p[2].getint())

        @pg.production("expr : NUMBER")
        def expr_num(p):
            return BoxInt(int(p[0].getstr()))

        with self.assert_warns(ParserGeneratorWarning, "1 shift/reduce conflict"):
            parser = pg.build()
        assert parser.parse(FakeLexer([
            Token("NUMBER", "1"),
            Token("PLUS", "+"),
            Token("NUMBER", "4")
        ])) == BoxInt(5)

    def test_null_production(self):
        pg = ParserGenerator(["VALUE", "SPACE"])

        @pg.production("main : values")
        def main(p):
            return p[0]

        @pg.production("values : none")
        def values_empty(p):
            return []

        @pg.production("values : VALUE")
        def values_value(p):
            return [p[0].getstr()]

        @pg.production("values : values SPACE VALUE")
        def values_values(p):
            return p[0] + [p[2].getstr()]

        @pg.production("none :")
        def none(p):
            return None

        parser = pg.build()
        assert parser.parse(FakeLexer([
            Token("VALUE", "abc"),
            Token("SPACE", " "),
            Token("VALUE", "def"),
            Token("SPACE", " "),
            Token("VALUE", "ghi"),
        ])) == ["abc", "def", "ghi"]

        assert parser.parse(FakeLexer([])) == []

    def test_precedence(self):
        pg = ParserGenerator(["NUMBER", "PLUS", "TIMES"], precedence=[
            ("left", ["PLUS"]),
            ("left", ["TIMES"]),
        ])

        @pg.production("main : expr")
        def main(p):
            return p[0]

        @pg.production("expr : expr PLUS expr")
        @pg.production("expr : expr TIMES expr")
        def expr_binop(p):
            return BoxInt({
                "+": operator.add,
                "*": operator.mul
            }[p[1].getstr()](p[0].getint(), p[2].getint()))

        @pg.production("expr : NUMBER")
        def expr_num(p):
            return BoxInt(int(p[0].getstr()))

        parser = pg.build()

        assert parser.parse(FakeLexer([
            Token("NUMBER", "3"),
            Token("TIMES", "*"),
            Token("NUMBER", "4"),
            Token("PLUS",  "+"),
            Token("NUMBER", "5")
        ])) == BoxInt(17)

    def test_per_rule_precedence(self):
        pg = ParserGenerator(["NUMBER", "MINUS"], precedence=[
            ("right", ["UMINUS"]),
        ])

        @pg.production("main : expr")
        def main_expr(p):
            return p[0]

        @pg.production("expr : expr MINUS expr")
        def expr_minus(p):
            return BoxInt(p[0].getint() - p[2].getint())

        @pg.production("expr : MINUS expr", precedence="UMINUS")
        def expr_uminus(p):
            return BoxInt(-p[1].getint())

        @pg.production("expr : NUMBER")
        def expr_number(p):
            return BoxInt(int(p[0].getstr()))

        with self.assert_warns(ParserGeneratorWarning, "1 shift/reduce conflict"):
            parser = pg.build()

        assert parser.parse(FakeLexer([
            Token("MINUS", "-"),
            Token("NUMBER", "4"),
            Token("MINUS", "-"),
            Token("NUMBER", "5"),
        ])) == BoxInt(-9)

    def test_parse_error(self):
        pg = ParserGenerator(["VALUE"])

        @pg.production("main : VALUE")
        def main(p):
            return p[0]

        parser = pg.build()

        with py.test.raises(ParsingError) as exc_info:
            parser.parse(FakeLexer([
                Token("VALUE", "hello"),
                Token("VALUE", "world", SourcePosition(5, 10, 2)),
            ]))

        assert exc_info.value.getsourcepos().lineno == 10

    def test_parse_error_handler(self):
        pg = ParserGenerator(["VALUE"])

        @pg.production("main : VALUE")
        def main(p):
            return p[0]

        @pg.error
        def error_handler(token):
            raise ValueError(token)

        parser = pg.build()

        token = Token("VALUE", "world")

        with py.test.raises(ValueError) as exc_info:
            parser.parse(FakeLexer([
                Token("VALUE", "hello"),
                token
            ]))

        assert exc_info.value.args[0] is token

    def test_state(self):
        pg = ParserGenerator(["NUMBER", "PLUS"], precedence=[
            ("left", ["PLUS"]),
        ])

        @pg.production("main : expression")
        def main(state, p):
            state.count += 1
            return p[0]

        @pg.production("expression : expression PLUS expression")
        def expression_plus(state, p):
            state.count += 1
            return BoxInt(p[0].getint() + p[2].getint())

        @pg.production("expression : NUMBER")
        def expression_number(state, p):
            state.count += 1
            return BoxInt(int(p[0].getstr()))

        parser = pg.build()

        state = ParserState()
        assert parser.parse(FakeLexer([
            Token("NUMBER", "10"),
            Token("PLUS", "+"),
            Token("NUMBER", "12"),
            Token("PLUS", "+"),
            Token("NUMBER", "-2"),
        ]), state=state) == BoxInt(20)
        assert state.count == 6

    def test_error_handler_state(self):
        pg = ParserGenerator([])

        @pg.production("main :")
        def main(state, p):
            pass

        @pg.error
        def error(state, token):
            raise ValueError(state, token)

        parser = pg.build()

        state = ParserState()
        token = Token("VALUE", "")
        with py.test.raises(ValueError) as exc_info:
            parser.parse(FakeLexer([token]), state=state)

        assert exc_info.value.args[0] is state
        assert exc_info.value.args[1] is token

    def test_default_reductions(self):
        pg = ParserGenerator(["INTEGER_START", "INTEGER_VALUE", "COMPARE"], precedence=[
            ("nonassoc", ["COMPARE"])
        ])
        record = []

        @pg.production("main : expr")
        def main(p):
            record.append("main")
            return p[0]

        @pg.production("expr : expr COMPARE expr")
        def expr_compare(p):
            record.append("expr:compare")
            return BoxInt(p[0].getint() - p[2].getint())

        @pg.production("expr : INTEGER_START INTEGER_VALUE")
        def expr_int(p):
            record.append("expr:int")
            return BoxInt(int(p[1].getstr()))

        parser = pg.build()

        assert parser.parse(RecordingLexer(record, [
            Token("INTEGER_START", ""),
            Token("INTEGER_VALUE", "10"),
            Token("COMPARE", "-"),
            Token("INTEGER_START", ""),
            Token("INTEGER_VALUE", "5")
        ])) == BoxInt(5)

        assert record == [
            "token:INTEGER_START",
            "token:INTEGER_VALUE",
            "expr:int",
            "token:COMPARE",
            "token:INTEGER_START",
            "token:INTEGER_VALUE",
            "expr:int",
            "expr:compare",
            "token:None",
            "main",
        ]
<commit_msg>rename test<commit_after>import operator

import py

from rply import ParserGenerator, Token, ParsingError
from rply.token import SourcePosition
from rply.errors import ParserGeneratorWarning

from .base import BaseTests
from .utils import FakeLexer, RecordingLexer, BoxInt, ParserState

class TestParser(BaseTests):
    def test_simple(self):
        pg = ParserGenerator(["VALUE"])

        @pg.production("main : VALUE")
        def main(p):
            return p[0]

        parser = pg.build()

        assert parser.parse(FakeLexer([Token("VALUE", "abc")])) == Token("VALUE", "abc")

    def test_arithmetic(self):
        pg = ParserGenerator(["NUMBER", "PLUS"])

        @pg.production("main : expr")
        def main(p):
            return p[0]

        @pg.production("expr : expr PLUS expr")
        def expr_op(p):
            return BoxInt(p[0].getint() + p[2].getint())

        @pg.production("expr : NUMBER")
        def expr_num(p):
            return BoxInt(int(p[0].getstr()))

        with self.assert_warns(ParserGeneratorWarning, "1 shift/reduce conflict"):
            parser = pg.build()
        assert parser.parse(FakeLexer([
            Token("NUMBER", "1"),
            Token("PLUS", "+"),
            Token("NUMBER", "4")
        ])) == BoxInt(5)

    def test_null_production(self):
        pg = ParserGenerator(["VALUE", "SPACE"])

        @pg.production("main : values")
        def main(p):
            return p[0]

        @pg.production("values : none")
        def values_empty(p):
            return []

        @pg.production("values : VALUE")
        def values_value(p):
            return [p[0].getstr()]

        @pg.production("values : values SPACE VALUE")
        def values_values(p):
            return p[0] + [p[2].getstr()]

        @pg.production("none :")
        def none(p):
            return None

        parser = pg.build()
        assert parser.parse(FakeLexer([
            Token("VALUE", "abc"),
            Token("SPACE", " "),
            Token("VALUE", "def"),
            Token("SPACE", " "),
            Token("VALUE", "ghi"),
        ])) == ["abc", "def", "ghi"]

        assert parser.parse(FakeLexer([])) == []

    def test_precedence(self):
        pg = ParserGenerator(["NUMBER", "PLUS", "TIMES"], precedence=[
            ("left", ["PLUS"]),
            ("left", ["TIMES"]),
        ])

        @pg.production("main : expr")
        def main(p):
            return p[0]

        @pg.production("expr : expr PLUS expr")
        @pg.production("expr : expr TIMES expr")
        def expr_binop(p):
            return BoxInt({
                "+": operator.add,
                "*": operator.mul
            }[p[1].getstr()](p[0].getint(), p[2].getint()))

        @pg.production("expr : NUMBER")
        def expr_num(p):
            return BoxInt(int(p[0].getstr()))

        parser = pg.build()

        assert parser.parse(FakeLexer([
            Token("NUMBER", "3"),
            Token("TIMES", "*"),
            Token("NUMBER", "4"),
            Token("PLUS",  "+"),
            Token("NUMBER", "5")
        ])) == BoxInt(17)

    def test_per_rule_precedence(self):
        pg = ParserGenerator(["NUMBER", "MINUS"], precedence=[
            ("right", ["UMINUS"]),
        ])

        @pg.production("main : expr")
        def main_expr(p):
            return p[0]

        @pg.production("expr : expr MINUS expr")
        def expr_minus(p):
            return BoxInt(p[0].getint() - p[2].getint())

        @pg.production("expr : MINUS expr", precedence="UMINUS")
        def expr_uminus(p):
            return BoxInt(-p[1].getint())

        @pg.production("expr : NUMBER")
        def expr_number(p):
            return BoxInt(int(p[0].getstr()))

        with self.assert_warns(ParserGeneratorWarning, "1 shift/reduce conflict"):
            parser = pg.build()

        assert parser.parse(FakeLexer([
            Token("MINUS", "-"),
            Token("NUMBER", "4"),
            Token("MINUS", "-"),
            Token("NUMBER", "5"),
        ])) == BoxInt(-9)

    def test_parse_error(self):
        pg = ParserGenerator(["VALUE"])

        @pg.production("main : VALUE")
        def main(p):
            return p[0]

        parser = pg.build()

        with py.test.raises(ParsingError) as exc_info:
            parser.parse(FakeLexer([
                Token("VALUE", "hello"),
                Token("VALUE", "world", SourcePosition(5, 10, 2)),
            ]))

        assert exc_info.value.getsourcepos().lineno == 10

    def test_parse_error_handler(self):
        pg = ParserGenerator(["VALUE"])

        @pg.production("main : VALUE")
        def main(p):
            return p[0]

        @pg.error
        def error_handler(token):
            raise ValueError(token)

        parser = pg.build()

        token = Token("VALUE", "world")

        with py.test.raises(ValueError) as exc_info:
            parser.parse(FakeLexer([
                Token("VALUE", "hello"),
                token
            ]))

        assert exc_info.value.args[0] is token

    def test_state(self):
        pg = ParserGenerator(["NUMBER", "PLUS"], precedence=[
            ("left", ["PLUS"]),
        ])

        @pg.production("main : expression")
        def main(state, p):
            state.count += 1
            return p[0]

        @pg.production("expression : expression PLUS expression")
        def expression_plus(state, p):
            state.count += 1
            return BoxInt(p[0].getint() + p[2].getint())

        @pg.production("expression : NUMBER")
        def expression_number(state, p):
            state.count += 1
            return BoxInt(int(p[0].getstr()))

        parser = pg.build()

        state = ParserState()
        assert parser.parse(FakeLexer([
            Token("NUMBER", "10"),
            Token("PLUS", "+"),
            Token("NUMBER", "12"),
            Token("PLUS", "+"),
            Token("NUMBER", "-2"),
        ]), state=state) == BoxInt(20)
        assert state.count == 6

    def test_error_handler_state(self):
        pg = ParserGenerator([])

        @pg.production("main :")
        def main(state, p):
            pass

        @pg.error
        def error(state, token):
            raise ValueError(state, token)

        parser = pg.build()

        state = ParserState()
        token = Token("VALUE", "")
        with py.test.raises(ValueError) as exc_info:
            parser.parse(FakeLexer([token]), state=state)

        assert exc_info.value.args[0] is state
        assert exc_info.value.args[1] is token

    def test_default_reductions(self):
        pg = ParserGenerator(["INTEGER_START", "INTEGER_VALUE", "COMPARE"], precedence=[
            ("nonassoc", ["COMPARE"])
        ])
        record = []

        @pg.production("main : expr")
        def main(p):
            record.append("main")
            return p[0]

        @pg.production("expr : expr COMPARE expr")
        def expr_compare(p):
            record.append("expr:compare")
            return BoxInt(p[0].getint() - p[2].getint())

        @pg.production("expr : INTEGER_START INTEGER_VALUE")
        def expr_int(p):
            record.append("expr:int")
            return BoxInt(int(p[1].getstr()))

        parser = pg.build()

        assert parser.parse(RecordingLexer(record, [
            Token("INTEGER_START", ""),
            Token("INTEGER_VALUE", "10"),
            Token("COMPARE", "-"),
            Token("INTEGER_START", ""),
            Token("INTEGER_VALUE", "5")
        ])) == BoxInt(5)

        assert record == [
            "token:INTEGER_START",
            "token:INTEGER_VALUE",
            "expr:int",
            "token:COMPARE",
            "token:INTEGER_START",
            "token:INTEGER_VALUE",
            "expr:int",
            "expr:compare",
            "token:None",
            "main",
        ]
<|endoftext|>


<commit_before># -*- coding: utf-8; -*-

import collections
import email.utils
from http import HTTPStatus
import io
import typing

import httpolice
from mitmproxy import ctx
import mitmproxy.flow
import mitmproxy.http
import mitmproxy.net.http
import mitmproxy.types

__version__ = '0.9.0.dev1'

class MitmproxyHTTPolice:

    def __init__(self):
        self.last_report = None

    def load(self, loader):
        loader.add_option(
            name='httpolice_silence',
            # ``typing.Sequence[int]`` would be better, but is not supported.
            typespec=typing.Sequence[str],
            default=[],
            help='Silence these HTTPolice notice IDs.',
        )
        loader.add_option(
            name='httpolice_mark',
            # Could make this a ``typing.Optional[str]``, but
            # that doesn't work well with the interactive editors,
            # so make "disable" an explicit choice.
            typespec=str,
            choices=[''] + [sev.name for sev in httpolice.Severity],
            default='',
            help=
                'Mark flows where HTTPolice found notices of this severity '
                'or higher (empty to disable).'
        )

    def request(self, flow):
        if flow.request.path == '/+httpolice/':
            flow.response = self.serve_report()

    def response(self, flow):
        exch = flow_to_exchange(flow)
        attach_report(exch, flow)
        mark_exchange(exch, flow)
        log_exchange(exch, flow)

    @mitmproxy.command.command('httpolice.report.html')
    def html_report(self,
                    flows: typing.Sequence[mitmproxy.flow.Flow],
                    path: mitmproxy.types.Path) -> None:
        """Produce an HTTPolice report (HTML) on flows."""
        self.report(flows, httpolice.html_report, path)

    @mitmproxy.command.command('httpolice.report.text')
    def text_report(self,
                    flows: typing.Sequence[mitmproxy.flow.Flow],
                    path: mitmproxy.types.Path) -> None:
        """Produce an HTTPolice report (text) on flows."""
        self.report(flows, httpolice.text_report, path)

    def report(self, flows, report_func, path):
        exchanges = (flow_to_exchange(flow) for flow in flows)
        if path == '-':
            buf = io.BytesIO()
            report_func(exchanges, buf)
            self.last_report = buf.getvalue()
            ctx.log.alert(
                f'HTTPolice: saved report on {len(flows)} flows in memory')
        else:
            with open(path, 'wb') as f:
                report_func(exchanges, f)
            ctx.log.alert(
                f'HTTPolice: wrote report on {len(flows)} flows to {path}')

    def serve_report(self):
        if self.last_report is None:
            status_code = HTTPStatus.NOT_FOUND.value
            content = (
                f'<!DOCTYPE html>
No report has been produced yet.
'
            ).encode('utf-8')
        else:
            status_code = HTTPStatus.OK.value
            content = self.last_report
        headers = {
            'Date': email.utils.formatdate(usegmt=True),
            'Content-Type': (
                'text/html; charset=utf-8' if content.startswith(b'<!')
                else 'text/plain; charset=utf-8'
            ),
            'Cache-Control': 'no-store',
        }
        return mitmproxy.http.HTTPResponse.wrap(
            mitmproxy.net.http.Response.make(status_code, content, headers),
        )

def flow_to_exchange(flow):
    req = construct_request(flow)
    resp = construct_response(flow)
    exch = httpolice.Exchange(req, [resp] if resp else [])
    exch.silence([int(id_) for id_ in ctx.options.httpolice_silence])
    httpolice.check_exchange(exch)
    return exch

def construct_request(flow):
    version, headers, body = extract_message_basics(flow.request)
    scheme = decode(flow.request.scheme)
    method = decode(flow.request.method)

    # Authority-form and absolute-form requests in the tunnel
    # are simply rejected as errors by mitmproxy, closing the connection.
    target = decode(flow.request.path)

    if version == 'HTTP/2':
        pseudo_headers = httpolice.helpers.pop_pseudo_headers(headers)
        authority = pseudo_headers.get(':authority')
        has_host = any(k.lower() == 'host' for (k, v) in headers)
        if authority and not has_host and target.startswith('/'):
            # Reconstruct HTTP/2's equivalent of
            # the "absolute form" of request target (RFC 7540 Section 8.1.2.3).
            target = scheme + '://' + decode(authority) + target

    return httpolice.Request(scheme, method, target, version, headers, body)

def construct_response(flow):
    if flow.response is None:
        return None
    version, headers, body = extract_message_basics(flow.response)
    status = flow.response.status_code
    reason = decode(flow.response.reason)
    if version == 'HTTP/2':
        httpolice.helpers.pop_pseudo_headers(headers)
    return httpolice.Response(version, status, reason, headers, body)

def extract_message_basics(msg):
    version = decode(msg.http_version)
    if version == 'HTTP/2.0':
        version = 'HTTP/2'
    headers = [(decode(k), v) for (k, v) in msg.headers.fields]
    body = msg.raw_content
    return version, headers, body

def attach_report(exch, flow):
    buf = io.BytesIO()
    httpolice.text_report([exch], buf)
    report = buf.getvalue().decode('utf-8')
    for_request, for_response = parse_report(report)
    for title, lines in [('HTTPolice: request', for_request),
                         ('HTTPolice: response', for_response)]:
        if lines:
            text = u'\n'.join(lines) + u'\n'
            try:
                # If this script is being run on a flow previously loaded
                # from file, `flow.metadata` might already contain our keys
                # in the wrong order. Reinsert them instead of just updating.
                flow.metadata.pop(title, None)
                flow.metadata[title] = ReprString(text)
            except Exception:
                # `flow.metadata` is not public API,
                # so could theoretically fail.
                pass

def parse_report(report):
    # `report` is a text report as produced by HTTPolice. From it, we want to
    # extract the notices (titles) for the request and for the response.
    # This may sound stupid: why not just make HTTPolice return them
    # in a structured form? But that would have to be a public API
    # (mitmproxy-HTTPolice doesn't use any private APIs from HTTPolice), and I
    # don't want to add public APIs to HTTPolice without a clear picture of
    # how and by whom they will be used. I want some sort of "JSON report"
    # in HTTPolice eventually, but I don't know the details yet. So for now,
    # let's just parse semi-structured text -- a great Unix tradition.
    for_request, for_response = [], []
    target = for_request
    for line in report.splitlines():
        if line.startswith('------------ request:'):
            target = for_request
        elif line.startswith('------------ response:'):
            target = for_response
        else:
            target.append(line)
    return for_request, for_response

def mark_exchange(exch, flow):
    if ctx.options.httpolice_mark:
        mark_severity = httpolice.Severity[ctx.options.httpolice_mark]
        if any(notice.severity >= mark_severity
               for msg in [exch.request] + exch.responses
               for notice in msg.notices):
            flow.marked = True

def log_exchange(exch, flow):
    # Produce lines like "1 errors, 2 comments" without hardcoding severities.
    severities = collections.Counter(notice.severity
                                     for msg in [exch.request] + exch.responses
                                     for notice in msg.notices)
    pieces = ['%d %ss' % (n, severity.name)
              for (severity, n) in sorted(severities.items(), reverse=True)
              if severity > httpolice.Severity.debug]
    if pieces:
        log_func = (ctx.log.warn
                    if max(severities) >= httpolice.Severity.error
                    else ctx.log.info)
        log_func('HTTPolice: %s in: %s %s ← %d' % (
            ', '.join(pieces),
            flow.request.method, ellipsize(flow.request.path),
            flow.response.status_code,
        ))

def decode(s):
    if isinstance(s, bytes):
        return s.decode('iso-8859-1')
    return s

def ellipsize(s, max_length=40):
    if len(s) <= max_length:
        return s
    ellipsis = '...'
    return s[:(max_length - len(ellipsis))] + ellipsis

class ReprString(str):

    # Currently mitmproxy displays ``repr()`` in details view, not ``str()``.
    # See also https://discourse.mitmproxy.org/t/extending-the-ui/359/5

    __slots__ = []

    def __repr__(self):
        return str(self)

addons = [MitmproxyHTTPolice()]

if __name__ == '__main__':
    # Print the path to this script,
    # for substitution into the mitmproxy command.
    print(__file__)
<commit_msg>Stop using %-formatting<commit_after># -*- coding: utf-8; -*-

import collections
import email.utils
from http import HTTPStatus
import io
import typing

import httpolice
from mitmproxy import ctx
import mitmproxy.flow
import mitmproxy.http
import mitmproxy.net.http
import mitmproxy.types

__version__ = '0.9.0.dev1'

class MitmproxyHTTPolice:

    def __init__(self):
        self.last_report = None

    def load(self, loader):
        loader.add_option(
            name='httpolice_silence',
            # ``typing.Sequence[int]`` would be better, but is not supported.
            typespec=typing.Sequence[str],
            default=[],
            help='Silence these HTTPolice notice IDs.',
        )
        loader.add_option(
            name='httpolice_mark',
            # Could make this a ``typing.Optional[str]``, but
            # that doesn't work well with the interactive editors,
            # so make "disable" an explicit choice.
            typespec=str,
            choices=[''] + [sev.name for sev in httpolice.Severity],
            default='',
            help=
                'Mark flows where HTTPolice found notices of this severity '
                'or higher (empty to disable).'
        )

    def request(self, flow):
        if flow.request.path == '/+httpolice/':
            flow.response = self.serve_report()

    def response(self, flow):
        exch = flow_to_exchange(flow)
        attach_report(exch, flow)
        mark_exchange(exch, flow)
        log_exchange(exch, flow)

    @mitmproxy.command.command('httpolice.report.html')
    def html_report(self,
                    flows: typing.Sequence[mitmproxy.flow.Flow],
                    path: mitmproxy.types.Path) -> None:
        """Produce an HTTPolice report (HTML) on flows."""
        self.report(flows, httpolice.html_report, path)

    @mitmproxy.command.command('httpolice.report.text')
    def text_report(self,
                    flows: typing.Sequence[mitmproxy.flow.Flow],
                    path: mitmproxy.types.Path) -> None:
        """Produce an HTTPolice report (text) on flows."""
        self.report(flows, httpolice.text_report, path)

    def report(self, flows, report_func, path):
        exchanges = (flow_to_exchange(flow) for flow in flows)
        if path == '-':
            buf = io.BytesIO()
            report_func(exchanges, buf)
            self.last_report = buf.getvalue()
            ctx.log.alert(
                f'HTTPolice: saved report on {len(flows)} flows in memory')
        else:
            with open(path, 'wb') as f:
                report_func(exchanges, f)
            ctx.log.alert(
                f'HTTPolice: wrote report on {len(flows)} flows to {path}')

    def serve_report(self):
        if self.last_report is None:
            status_code = HTTPStatus.NOT_FOUND.value
            content = (
                f'<!DOCTYPE html>
No report has been produced yet.
'
            ).encode('utf-8')
        else:
            status_code = HTTPStatus.OK.value
            content = self.last_report
        headers = {
            'Date': email.utils.formatdate(usegmt=True),
            'Content-Type': (
                'text/html; charset=utf-8' if content.startswith(b'<!')
                else 'text/plain; charset=utf-8'
            ),
            'Cache-Control': 'no-store',
        }
        return mitmproxy.http.HTTPResponse.wrap(
            mitmproxy.net.http.Response.make(status_code, content, headers),
        )

def flow_to_exchange(flow):
    req = construct_request(flow)
    resp = construct_response(flow)
    exch = httpolice.Exchange(req, [resp] if resp else [])
    exch.silence([int(id_) for id_ in ctx.options.httpolice_silence])
    httpolice.check_exchange(exch)
    return exch

def construct_request(flow):
    version, headers, body = extract_message_basics(flow.request)
    scheme = decode(flow.request.scheme)
    method = decode(flow.request.method)

    # Authority-form and absolute-form requests in the tunnel
    # are simply rejected as errors by mitmproxy, closing the connection.
    target = decode(flow.request.path)

    if version == 'HTTP/2':
        pseudo_headers = httpolice.helpers.pop_pseudo_headers(headers)
        authority = pseudo_headers.get(':authority')
        has_host = any(k.lower() == 'host' for (k, v) in headers)
        if authority and not has_host and target.startswith('/'):
            # Reconstruct HTTP/2's equivalent of
            # the "absolute form" of request target (RFC 7540 Section 8.1.2.3).
            target = scheme + '://' + decode(authority) + target

    return httpolice.Request(scheme, method, target, version, headers, body)

def construct_response(flow):
    if flow.response is None:
        return None
    version, headers, body = extract_message_basics(flow.response)
    status = flow.response.status_code
    reason = decode(flow.response.reason)
    if version == 'HTTP/2':
        httpolice.helpers.pop_pseudo_headers(headers)
    return httpolice.Response(version, status, reason, headers, body)

def extract_message_basics(msg):
    version = decode(msg.http_version)
    if version == 'HTTP/2.0':
        version = 'HTTP/2'
    headers = [(decode(k), v) for (k, v) in msg.headers.fields]
    body = msg.raw_content
    return version, headers, body

def attach_report(exch, flow):
    buf = io.BytesIO()
    httpolice.text_report([exch], buf)
    report = buf.getvalue().decode('utf-8')
    for_request, for_response = parse_report(report)
    for title, lines in [('HTTPolice: request', for_request),
                         ('HTTPolice: response', for_response)]:
        if lines:
            text = u'\n'.join(lines) + u'\n'
            try:
                # If this script is being run on a flow previously loaded
                # from file, `flow.metadata` might already contain our keys
                # in the wrong order. Reinsert them instead of just updating.
                flow.metadata.pop(title, None)
                flow.metadata[title] = ReprString(text)
            except Exception:
                # `flow.metadata` is not public API,
                # so could theoretically fail.
                pass

def parse_report(report):
    # `report` is a text report as produced by HTTPolice. From it, we want to
    # extract the notices (titles) for the request and for the response.
    # This may sound stupid: why not just make HTTPolice return them
    # in a structured form? But that would have to be a public API
    # (mitmproxy-HTTPolice doesn't use any private APIs from HTTPolice), and I
    # don't want to add public APIs to HTTPolice without a clear picture of
    # how and by whom they will be used. I want some sort of "JSON report"
    # in HTTPolice eventually, but I don't know the details yet. So for now,
    # let's just parse semi-structured text -- a great Unix tradition.
    for_request, for_response = [], []
    target = for_request
    for line in report.splitlines():
        if line.startswith('------------ request:'):
            target = for_request
        elif line.startswith('------------ response:'):
            target = for_response
        else:
            target.append(line)
    return for_request, for_response

def mark_exchange(exch, flow):
    if ctx.options.httpolice_mark:
        mark_severity = httpolice.Severity[ctx.options.httpolice_mark]
        if any(notice.severity >= mark_severity
               for msg in [exch.request] + exch.responses
               for notice in msg.notices):
            flow.marked = True

def log_exchange(exch, flow):
    # Produce lines like "1 errors, 2 comments" without hardcoding severities.
    severities = collections.Counter(notice.severity
                                     for msg in [exch.request] + exch.responses
                                     for notice in msg.notices)
    pieces = [f'{n} {severity.name}s'
              for (severity, n) in sorted(severities.items(), reverse=True)
              if severity > httpolice.Severity.debug]
    if pieces:
        log_func = (ctx.log.warn
                    if max(severities) >= httpolice.Severity.error
                    else ctx.log.info)
        log_func('HTTPolice: {0} in: {1} {2} ← {3}'.format(
            ', '.join(pieces),
            flow.request.method, ellipsize(flow.request.path),
            flow.response.status_code,
        ))

def decode(s):
    if isinstance(s, bytes):
        return s.decode('iso-8859-1')
    return s

def ellipsize(s, max_length=40):
    if len(s) <= max_length:
        return s
    ellipsis = '...'
    return s[:(max_length - len(ellipsis))] + ellipsis

class ReprString(str):

    # Currently mitmproxy displays ``repr()`` in details view, not ``str()``.
    # See also https://discourse.mitmproxy.org/t/extending-the-ui/359/5

    __slots__ = []

    def __repr__(self):
        return str(self)

addons = [MitmproxyHTTPolice()]

if __name__ == '__main__':
    # Print the path to this script,
    # for substitution into the mitmproxy command.
    print(__file__)
<|endoftext|>


<commit_before>import pytest
import os
import slug
from slug import ProcessGroup, Process, Pipe
from conftest import runpy, not_in_path

def test_group():
    with ProcessGroup() as pg:
        p1 = Pipe()
        pg.add(Process(runpy("print('spam')"), stdout=p1.side_in))
        p2 = Pipe()
        pg.add(Process(runpy("print('eggs')"), stdout=p2.side_in))
        p3 = Pipe()
        pg.add(Process(runpy("print('vikings')"), stdout=p3.side_in))

    pg.start()
    p1.side_in.close()
    p2.side_in.close()
    p3.side_in.close()
    pg.join()

    assert p1.side_out.read().rstrip() == b'spam'
    assert p2.side_out.read().rstrip() == b'eggs'
    assert p3.side_out.read().rstrip() == b'vikings'

def test_terminate():
    with ProcessGroup() as pg:
        pg.add(Process(runpy("input()")))
        pg.add(Process(runpy("input()")))
        pg.add(Process(runpy("input()")))

    pg.start()
    pg.terminate()
    pg.join()
    # Timing out is failure.

def test_kill():
    with ProcessGroup() as pg:
        pg.add(Process(runpy("input()")))
        pg.add(Process(runpy("input()")))
        pg.add(Process(runpy("input()")))

    pg.start()
    pg.kill()
    pg.join()
    # Timing out is failure.

@pytest.mark.skipif(not hasattr(ProcessGroup, 'pgid'),
                    reason="No Process Group IDs")
def test_pgid():
    with ProcessGroup() as pg:
        pg.add(Process(runpy("input()")))
        pg.add(Process(runpy("input()")))
        pg.add(Process(runpy("input()")))
    pg.start()
    assert pg.pgid is not None
    assert all(pg.pgid == p.pgid for p in pg)
    assert pg.pgid != os.getpgrp()
    pg.kill()

@pytest.mark.skipif(ProcessGroup is slug.base.ProcessGroup,
                    reason="Base Process Group is broken this way")
def test_kill_descendents():
    with ProcessGroup() as pg:
        pin = Pipe()
        pout = Pipe()
        pg.add(Process(
            runpy("""import subprocess
import sys
child = subprocess.Popen([sys.executable, '-c', '''
print("spam")
input()
print("eggs")
'''])
child.wait()
"""),
            stdin=pin.side_out,
            stdout=pout.side_in,
        ))
    pg.start()
    pin.side_out.close()
    pout.side_in.close()
    assert pout.side_out.read(4) == b'spam'
    pg.terminate()
    pg.join()
    try:
        pin.side_in.write(b'\n')
        pin.side_in.close()
    except (BrokenPipeError, OSError):
        # BrokenPipeError on Linux, nothing on Mac, OSError on Windows
        pass
    assert pout.side_out.read() in b'\r\n'

@pytest.mark.skipif(not_in_path('ls', 'wc'),
                    reason="Requires the ls and wc binaries")
def test_lswc_gh10():
    """
    Tests ls|wc
    """

    # From https://github.com/xonsh/slug/issues/10

    with ProcessGroup() as pg:
        pipe = Pipe()
        pg.add(Process(['echo', 'foo'], stdout=pipe.side_in))
        pg.add(Process(['cat'], stdin=pipe.side_out))
    pg.start()
    pipe.side_in.close()
    pipe.side_out.close()
    pg.join()
<commit_msg>Actually test the thing mitnk was testing.<commit_after>import pytest
import os
import slug
from slug import ProcessGroup, Process, Pipe
from conftest import runpy, not_in_path

def test_group():
    with ProcessGroup() as pg:
        p1 = Pipe()
        pg.add(Process(runpy("print('spam')"), stdout=p1.side_in))
        p2 = Pipe()
        pg.add(Process(runpy("print('eggs')"), stdout=p2.side_in))
        p3 = Pipe()
        pg.add(Process(runpy("print('vikings')"), stdout=p3.side_in))

    pg.start()
    p1.side_in.close()
    p2.side_in.close()
    p3.side_in.close()
    pg.join()

    assert p1.side_out.read().rstrip() == b'spam'
    assert p2.side_out.read().rstrip() == b'eggs'
    assert p3.side_out.read().rstrip() == b'vikings'

def test_terminate():
    with ProcessGroup() as pg:
        pg.add(Process(runpy("input()")))
        pg.add(Process(runpy("input()")))
        pg.add(Process(runpy("input()")))

    pg.start()
    pg.terminate()
    pg.join()
    # Timing out is failure.

def test_kill():
    with ProcessGroup() as pg:
        pg.add(Process(runpy("input()")))
        pg.add(Process(runpy("input()")))
        pg.add(Process(runpy("input()")))

    pg.start()
    pg.kill()
    pg.join()
    # Timing out is failure.

@pytest.mark.skipif(not hasattr(ProcessGroup, 'pgid'),
                    reason="No Process Group IDs")
def test_pgid():
    with ProcessGroup() as pg:
        pg.add(Process(runpy("input()")))
        pg.add(Process(runpy("input()")))
        pg.add(Process(runpy("input()")))
    pg.start()
    assert pg.pgid is not None
    assert all(pg.pgid == p.pgid for p in pg)
    assert pg.pgid != os.getpgrp()
    pg.kill()

@pytest.mark.skipif(ProcessGroup is slug.base.ProcessGroup,
                    reason="Base Process Group is broken this way")
def test_kill_descendents():
    with ProcessGroup() as pg:
        pin = Pipe()
        pout = Pipe()
        pg.add(Process(
            runpy("""import subprocess
import sys
child = subprocess.Popen([sys.executable, '-c', '''
print("spam")
input()
print("eggs")
'''])
child.wait()
"""),
            stdin=pin.side_out,
            stdout=pout.side_in,
        ))
    pg.start()
    pin.side_out.close()
    pout.side_in.close()
    assert pout.side_out.read(4) == b'spam'
    pg.terminate()
    pg.join()
    try:
        pin.side_in.write(b'\n')
        pin.side_in.close()
    except (BrokenPipeError, OSError):
        # BrokenPipeError on Linux, nothing on Mac, OSError on Windows
        pass
    assert pout.side_out.read() in b'\r\n'

@pytest.mark.skipif(not_in_path('ls', 'wc'),
                    reason="Requires the ls and wc binaries")
def test_lswc_gh10():
    """
    Tests ls|wc
    """

    # From https://github.com/xonsh/slug/issues/10

    with ProcessGroup() as pg:
        pipe = Pipe()
        pg.add(Process(['ls'], stdout=pipe.side_in))
        pg.add(Process(['wc'], stdin=pipe.side_out))
    pg.start()
    pipe.side_in.close()
    pipe.side_out.close()
    pg.join()
<|endoftext|>