<commit_before>from typing import Iterable, Any, Union, Tuple, Iterator, Sequence, cast, Dict, Optional
import numpy
import itertools
import threading
import random

try:  # pragma: no cover
    import cupy
    from cupy import get_array_module
except ImportError:
    cupy = None
    get_array_module = lambda _: numpy

try:  # pragma: no cover
    import torch
    import torch.tensor
    import torch.utils.dlpack

    has_torch = True
except ImportError:  # pragma: no cover
    has_torch = False

try:  # pragma: no cover
    import tfdlpack

    has_tfdlpack = True
except ImportError:
    has_tfdlpack = False

try:  # pragma: no cover
    import tensorflow as tf

    has_tensorflow = True
except ImportError:  # pragma: no cover
    has_tensorflow = False

from .types import Array, Ragged, Padded, ArgsKwargs, RNNState, IntsNd, FloatsNd

def fix_random_seed(seed: int = 0) -> None:  # pragma: no cover
    """Set the random seed across random, numpy.random and cupy.random."""
    random.seed(seed)
    numpy.random.seed(seed)
    if cupy is not None:
        cupy.random.seed(seed)

def create_thread_local(attrs: Dict[str, Any]):
    obj = threading.local()
    for name, value in attrs.items():
        setattr(obj, name, value)
    return obj

def is_xp_array(obj: Any) -> bool:
    """Check whether an object is a numpy or cupy array."""
    return is_numpy_array(obj) or is_cupy_array(obj)

def is_cupy_array(obj: Any) -> bool:  # pragma: no cover
    """Check whether an object is a cupy array"""
    if cupy is None:
        return False
    elif isinstance(obj, cupy.ndarray):
        return True
    else:
        return False

def is_numpy_array(obj: Any) -> bool:
    """Check whether an object is a numpy array"""
    if isinstance(obj, numpy.ndarray):
        return True
    else:
        return False

def is_torch_array(obj) -> bool:
    if torch is None:
        return False
    elif isinstance(obj, torch.Tensor):
        return True
    else:
        return False

def is_tensorflow_array(obj) -> bool:
    if tf is None:
        return False
    elif isinstance(obj, tf.Tensor):
        return True
    else:
        return False

def set_active_gpu(gpu_id: int) -> "cupy.cuda.Device":  # pragma: no cover
    """Set the current GPU device for cupy and torch (if available)."""
    import cupy.cuda.device

    device = cupy.cuda.device.Device(gpu_id)
    device.use()
    try:
        import torch

        torch.cuda.set_device(gpu_id)
        torch.set_default_tensor_type("torch.cuda.FloatTensor")
    except ImportError:
        pass
    return device

def prefer_gpu(gpu_id: int = 0) -> bool:  # pragma: no cover
    """Use GPU if it's available. Returns True if so, False otherwise."""
    from .backends.cupy_ops import CupyOps

    if CupyOps.xp is None:
        return False
    else:
        require_gpu(gpu_id=gpu_id)
        return True

def require_gpu(gpu_id: int = 0) -> bool:  # pragma: no cover
    from .backends import set_current_ops, CupyOps

    if CupyOps.xp is None:
        raise ValueError("GPU is not accessible. Was the library installed correctly?")

    set_current_ops(CupyOps())
    set_active_gpu(gpu_id)
    return True

def get_shuffled_batches(
    X: Array, Y: Array, batch_size
) -> Iterable[Tuple[Array, Array]]:
    """Iterate over paired batches from two arrays, shuffling the indices."""
    xp = get_array_module(X)
    indices = xp.arange(X.shape[0], dtype="i")
    xp.random.shuffle(indices)
    for index_batch in minibatch(indices, size=batch_size):
        yield X[index_batch], Y[index_batch]

def minibatch(
    items: Iterable[Any], size: Union[int, Iterator[int]] = 8
) -> Iterable[Any]:
    """Iterate over batches of items. `size` may be an iterator,
    so that batch-size can vary on each step.
    """
    if isinstance(size, int):
        size_ = itertools.repeat(size)
    else:
        size_ = size
    if hasattr(items, "__len__") and hasattr(items, "__getitem__"):
        i = 0
        while i < len(items):  # type: ignore
            batch_size = next(size_)
            yield items[i : i + batch_size]  # type: ignore
            i += batch_size
    else:
        items = iter(items)
        while True:
            batch_size = next(size_)
            batch = list(itertools.islice(items, int(batch_size)))
            if len(batch) == 0:
                break
            yield list(batch)

def evaluate_model_on_arrays(
    model, inputs: Array, labels: Array, batch_size: int
) -> float:
    """Helper to evaluate accuracy of a model in the simplest cases, where
    there's one correct output class and the inputs are arrays. Not guaranteed
    to cover all situations – many applications will have to implement their
    own evaluation methods.
    """
    score = 0.0
    total = 0.0
    for i in range(0, inputs.shape[0], batch_size):
        X = inputs[i : i + batch_size]
        Y = labels[i : i + batch_size]
        Yh = model.predict(X)
        score += (Y.argmax(axis=1) == Yh.argmax(axis=1)).sum()
        total += Yh.shape[0]
    return score / total

def copy_array(dst: Array, src: Array) -> None:  # pragma: no cover
    if isinstance(dst, numpy.ndarray) and isinstance(src, numpy.ndarray):
        dst[:] = src
    elif is_cupy_array(dst):
        src = cupy.array(src, copy=False)
        cupy.copyto(dst, src)
    else:
        numpy.copyto(dst, src)

def to_categorical(Y: IntsNd, n_classes: Optional[int] = None) -> FloatsNd:
    # From keras
    xp = get_array_module(Y)
    if xp is cupy:  # pragma: no cover
        Y = Y.get()
    Y = numpy.array(Y, dtype="int").ravel()
    if not n_classes:
        n_classes = numpy.max(Y) + 1
    n = Y.shape[0]
    categorical = numpy.zeros((n, n_classes), dtype="float32")
    categorical[numpy.arange(n), Y] = 1
    return xp.asarray(categorical)

def get_width(
    X: Union[Array, Ragged, Padded, Sequence[Array], RNNState], *, dim: int = -1
) -> int:
    """Infer the 'width' of a batch of data, which could be any of: Array,
    Ragged, Padded or Sequence of Arrays.
    """
    if isinstance(X, Ragged):
        return get_width(X.data, dim=dim)
    elif isinstance(X, Padded):
        return get_width(X.data, dim=dim)
    elif hasattr(X, "shape") and hasattr(X, "ndim"):
        X = cast(Array, X)
        if len(X.shape) == 0:
            return 0
        elif len(X.shape) == 1:
            return int(X.max()) + 1
        else:
            return X.shape[dim]
    elif isinstance(X, (list, tuple)):
        if len(X) == 0:
            return 0
        else:
            return get_width(X[0], dim=dim)
    else:
        err = "Cannot get width of object: has neither shape nor __getitem__"
        raise ValueError(err)

def assert_tensorflow_installed() -> None:  # pragma: no cover
    """Raise an ImportError if TensorFlow is not installed."""
    template = "TensorFlow support requires {pkg}: pip install thinc[tensorflow]"
    if not has_tfdlpack:
        raise ImportError(template.format(pkg="tfdlpack"))
    if not has_tensorflow:
        raise ImportError(template.format(pkg="tensorflow>=2.0.0"))

def assert_pytorch_installed() -> None:  # pragma: no cover
    """Raise an ImportError if PyTorch is not installed."""
    if not has_torch:
        raise ImportError("PyTorch support requires torch: pip install thinc[torch]")

# I think it seems reasonable to leave this untyped?
def convert_recursive(is_match, convert_item, obj):
    """Either convert a single value if it matches a given function, or
    recursively walk over potentially nested lists, tuples and dicts applying
    the conversion, and returns the same type.

    Also supports the ArgsKwargs dataclass.
    """
    if is_match(obj):
        return convert_item(obj)
    elif isinstance(obj, ArgsKwargs):
        converted = convert_recursive(is_match, convert_item, obj.items())
        return ArgsKwargs.from_items(converted)
    elif isinstance(obj, dict):
        converted = {}
        for key, value in obj.items():
            key = convert_recursive(is_match, convert_item, key)
            value = convert_recursive(is_match, convert_item, value)
            converted[key] = value
        return converted
    elif isinstance(obj, list):
        return [convert_recursive(is_match, convert_item, item) for item in obj]
    elif isinstance(obj, tuple):
        return tuple(convert_recursive(is_match, convert_item, item) for item in obj)
    else:
        return obj

def xp2torch(
    xp_tensor: Array, requires_grad: bool = False
) -> "torch.Tensor":  # pragma: no cover
    """Convert a numpy or cupy tensor to a PyTorch tensor."""
    if hasattr(xp_tensor, "toDlpack"):
        dlpack_tensor = xp_tensor.toDlpack()  # type: ignore
        torch_tensor = torch.utils.dlpack.from_dlpack(dlpack_tensor)
    else:
        torch_tensor = torch.from_numpy(xp_tensor)
    if requires_grad:
        torch_tensor.requires_grad_()
    return torch_tensor

def torch2xp(torch_tensor: "torch.Tensor") -> Array:  # pragma: no cover
    """Convert a torch tensor to a numpy or cupy tensor."""
    if torch_tensor.is_cuda:
        return cupy.fromDlpack(torch.utils.dlpack.to_dlpack(torch_tensor))
    else:
        return torch_tensor.detach().numpy()

def xp2tensorflow(
    xp_tensor: Array, requires_grad: bool = False, as_variable: bool = False
) -> "tf.Tensor":  # pragma: no cover
    """Convert a numpy or cupy tensor to a TensorFlow Tensor or Variable"""
    assert_tensorflow_installed()
    if hasattr(xp_tensor, "toDlpack"):
        dlpack_tensor = xp_tensor.toDlpack()  # type: ignore
        tensorflow_tensor = tfdlpack.from_dlpack(dlpack_tensor)
    else:
        tensorflow_tensor = tf.convert_to_tensor(xp_tensor)
    if as_variable:
        # tf.Variable() automatically puts in GPU if available.
        # So we need to control it using the context manager
        with tf.device(tensorflow_tensor.device):
            tensorflow_tensor = tf.Variable(tensorflow_tensor, trainable=requires_grad)
    if requires_grad is False and as_variable is False:
        # tf.stop_gradient() automatically puts in GPU if available.
        # So we need to control it using the context manager
        with tf.device(tensorflow_tensor.device):
            tensorflow_tensor = tf.stop_gradient(tensorflow_tensor)
    return tensorflow_tensor

def tensorflow2xp(tensorflow_tensor: "tf.Tensor") -> Array:  # pragma: no cover
    """Convert a Tensorflow tensor to numpy or cupy tensor."""
    assert_tensorflow_installed()
    if "GPU" in tensorflow_tensor.device:
        return cupy.fromDlpack(tfdlpack.to_dlpack(tensorflow_tensor))
    else:
        return tensorflow_tensor.numpy()

__all__ = [
    "fix_random_seed",
    "create_thread_local",
    "is_cupy_array",
    "is_numpy_array",
    "set_active_gpu",
    "prefer_gpu",
    "require_gpu",
    "copy_array",
    "get_shuffled_batches",
    "minibatch",
    "evaluate_model_on_arrays",
    "to_categorical",
    "get_width",
    "xp2torch",
    "torch2xp",
    "tensorflow2xp",
    "xp2tensorflow",
]
<commit_msg>Tidy up and add types / pragma<commit_after>from typing import Iterable, Any, Union, Tuple, Iterator, Sequence, cast, Dict
from typing import Optional, Callable
import numpy
import itertools
import threading
import random

try:  # pragma: no cover
    import cupy
    from cupy import get_array_module
except ImportError:
    cupy = None
    get_array_module = lambda _: numpy

try:  # pragma: no cover
    import torch
    import torch.tensor
    import torch.utils.dlpack

    has_torch = True
except ImportError:  # pragma: no cover
    has_torch = False

try:  # pragma: no cover
    import tfdlpack

    has_tfdlpack = True
except ImportError:
    has_tfdlpack = False

try:  # pragma: no cover
    import tensorflow as tf

    has_tensorflow = True
except ImportError:  # pragma: no cover
    has_tensorflow = False

from .types import Array, Ragged, Padded, ArgsKwargs, RNNState, IntsNd, FloatsNd

def fix_random_seed(seed: int = 0) -> None:  # pragma: no cover
    """Set the random seed across random, numpy.random and cupy.random."""
    random.seed(seed)
    numpy.random.seed(seed)
    if cupy is not None:
        cupy.random.seed(seed)

def create_thread_local(attrs: Dict[str, Any]):
    obj = threading.local()
    for name, value in attrs.items():
        setattr(obj, name, value)
    return obj

def is_xp_array(obj: Any) -> bool:
    """Check whether an object is a numpy or cupy array."""
    return is_numpy_array(obj) or is_cupy_array(obj)

def is_cupy_array(obj: Any) -> bool:  # pragma: no cover
    """Check whether an object is a cupy array"""
    if cupy is None:
        return False
    elif isinstance(obj, cupy.ndarray):
        return True
    else:
        return False

def is_numpy_array(obj: Any) -> bool:
    """Check whether an object is a numpy array"""
    if isinstance(obj, numpy.ndarray):
        return True
    else:
        return False

def is_torch_array(obj: Any) -> bool:  # pragma: no cover
    if torch is None:
        return False
    elif isinstance(obj, torch.Tensor):
        return True
    else:
        return False

def is_tensorflow_array(obj: Any) -> bool:  # pragma: no cover
    if tf is None:
        return False
    elif isinstance(obj, tf.Tensor):
        return True
    else:
        return False

def set_active_gpu(gpu_id: int) -> "cupy.cuda.Device":  # pragma: no cover
    """Set the current GPU device for cupy and torch (if available)."""
    import cupy.cuda.device

    device = cupy.cuda.device.Device(gpu_id)
    device.use()
    try:
        import torch

        torch.cuda.set_device(gpu_id)
        torch.set_default_tensor_type("torch.cuda.FloatTensor")
    except ImportError:
        pass
    return device

def prefer_gpu(gpu_id: int = 0) -> bool:  # pragma: no cover
    """Use GPU if it's available. Returns True if so, False otherwise."""
    from .backends.cupy_ops import CupyOps

    if CupyOps.xp is None:
        return False
    else:
        require_gpu(gpu_id=gpu_id)
        return True

def require_gpu(gpu_id: int = 0) -> bool:  # pragma: no cover
    from .backends import set_current_ops, CupyOps

    if CupyOps.xp is None:
        raise ValueError("GPU is not accessible. Was the library installed correctly?")

    set_current_ops(CupyOps())
    set_active_gpu(gpu_id)
    return True

def get_shuffled_batches(
    X: Array, Y: Array, batch_size
) -> Iterable[Tuple[Array, Array]]:
    """Iterate over paired batches from two arrays, shuffling the indices."""
    xp = get_array_module(X)
    indices = xp.arange(X.shape[0], dtype="i")
    xp.random.shuffle(indices)
    for index_batch in minibatch(indices, size=batch_size):
        yield X[index_batch], Y[index_batch]

def minibatch(
    items: Iterable[Any], size: Union[int, Iterator[int]] = 8
) -> Iterable[Any]:
    """Iterate over batches of items. `size` may be an iterator,
    so that batch-size can vary on each step.
    """
    if isinstance(size, int):
        size_ = itertools.repeat(size)
    else:
        size_ = size
    if hasattr(items, "__len__") and hasattr(items, "__getitem__"):
        i = 0
        while i < len(items):  # type: ignore
            batch_size = next(size_)
            yield items[i : i + batch_size]  # type: ignore
            i += batch_size
    else:
        items = iter(items)
        while True:
            batch_size = next(size_)
            batch = list(itertools.islice(items, int(batch_size)))
            if len(batch) == 0:
                break
            yield list(batch)

def evaluate_model_on_arrays(
    model, inputs: Array, labels: Array, batch_size: int
) -> float:
    """Helper to evaluate accuracy of a model in the simplest cases, where
    there's one correct output class and the inputs are arrays. Not guaranteed
    to cover all situations – many applications will have to implement their
    own evaluation methods.
    """
    score = 0.0
    total = 0.0
    for i in range(0, inputs.shape[0], batch_size):
        X = inputs[i : i + batch_size]
        Y = labels[i : i + batch_size]
        Yh = model.predict(X)
        score += (Y.argmax(axis=1) == Yh.argmax(axis=1)).sum()
        total += Yh.shape[0]
    return score / total

def copy_array(dst: Array, src: Array) -> None:  # pragma: no cover
    if isinstance(dst, numpy.ndarray) and isinstance(src, numpy.ndarray):
        dst[:] = src
    elif is_cupy_array(dst):
        src = cupy.array(src, copy=False)
        cupy.copyto(dst, src)
    else:
        numpy.copyto(dst, src)

def to_categorical(Y: IntsNd, n_classes: Optional[int] = None) -> FloatsNd:
    # From keras
    xp = get_array_module(Y)
    if xp is cupy:  # pragma: no cover
        Y = Y.get()
    Y = numpy.array(Y, dtype="int").ravel()
    if not n_classes:
        n_classes = numpy.max(Y) + 1
    n = Y.shape[0]
    categorical = numpy.zeros((n, n_classes), dtype="float32")
    categorical[numpy.arange(n), Y] = 1
    return xp.asarray(categorical)

def get_width(
    X: Union[Array, Ragged, Padded, Sequence[Array], RNNState], *, dim: int = -1
) -> int:
    """Infer the 'width' of a batch of data, which could be any of: Array,
    Ragged, Padded or Sequence of Arrays.
    """
    if isinstance(X, Ragged):
        return get_width(X.data, dim=dim)
    elif isinstance(X, Padded):
        return get_width(X.data, dim=dim)
    elif hasattr(X, "shape") and hasattr(X, "ndim"):
        X = cast(Array, X)
        if len(X.shape) == 0:
            return 0
        elif len(X.shape) == 1:
            return int(X.max()) + 1
        else:
            return X.shape[dim]
    elif isinstance(X, (list, tuple)):
        if len(X) == 0:
            return 0
        else:
            return get_width(X[0], dim=dim)
    else:
        err = "Cannot get width of object: has neither shape nor __getitem__"
        raise ValueError(err)

def assert_tensorflow_installed() -> None:  # pragma: no cover
    """Raise an ImportError if TensorFlow is not installed."""
    template = "TensorFlow support requires {pkg}: pip install thinc[tensorflow]"
    if not has_tfdlpack:
        raise ImportError(template.format(pkg="tfdlpack"))
    if not has_tensorflow:
        raise ImportError(template.format(pkg="tensorflow>=2.0.0"))

def assert_pytorch_installed() -> None:  # pragma: no cover
    """Raise an ImportError if PyTorch is not installed."""
    if not has_torch:
        raise ImportError("PyTorch support requires torch: pip install thinc[torch]")

def convert_recursive(
    is_match: Callable[[Any], bool], convert_item: Callable[[Any], Any], obj: Any
) -> Any:
    """Either convert a single value if it matches a given function, or
    recursively walk over potentially nested lists, tuples and dicts applying
    the conversion, and returns the same type. Also supports the ArgsKwargs
    dataclass.
    """
    if is_match(obj):
        return convert_item(obj)
    elif isinstance(obj, ArgsKwargs):
        converted = convert_recursive(is_match, convert_item, obj.items())
        return ArgsKwargs.from_items(converted)
    elif isinstance(obj, dict):
        converted = {}
        for key, value in obj.items():
            key = convert_recursive(is_match, convert_item, key)
            value = convert_recursive(is_match, convert_item, value)
            converted[key] = value
        return converted
    elif isinstance(obj, list):
        return [convert_recursive(is_match, convert_item, item) for item in obj]
    elif isinstance(obj, tuple):
        return tuple(convert_recursive(is_match, convert_item, item) for item in obj)
    else:
        return obj

def xp2torch(
    xp_tensor: Array, requires_grad: bool = False
) -> "torch.Tensor":  # pragma: no cover
    """Convert a numpy or cupy tensor to a PyTorch tensor."""
    if hasattr(xp_tensor, "toDlpack"):
        dlpack_tensor = xp_tensor.toDlpack()  # type: ignore
        torch_tensor = torch.utils.dlpack.from_dlpack(dlpack_tensor)
    else:
        torch_tensor = torch.from_numpy(xp_tensor)
    if requires_grad:
        torch_tensor.requires_grad_()
    return torch_tensor

def torch2xp(torch_tensor: "torch.Tensor") -> Array:  # pragma: no cover
    """Convert a torch tensor to a numpy or cupy tensor."""
    if torch_tensor.is_cuda:
        return cupy.fromDlpack(torch.utils.dlpack.to_dlpack(torch_tensor))
    else:
        return torch_tensor.detach().numpy()

def xp2tensorflow(
    xp_tensor: Array, requires_grad: bool = False, as_variable: bool = False
) -> "tf.Tensor":  # pragma: no cover
    """Convert a numpy or cupy tensor to a TensorFlow Tensor or Variable"""
    assert_tensorflow_installed()
    if hasattr(xp_tensor, "toDlpack"):
        dlpack_tensor = xp_tensor.toDlpack()  # type: ignore
        tensorflow_tensor = tfdlpack.from_dlpack(dlpack_tensor)
    else:
        tensorflow_tensor = tf.convert_to_tensor(xp_tensor)
    if as_variable:
        # tf.Variable() automatically puts in GPU if available.
        # So we need to control it using the context manager
        with tf.device(tensorflow_tensor.device):
            tensorflow_tensor = tf.Variable(tensorflow_tensor, trainable=requires_grad)
    if requires_grad is False and as_variable is False:
        # tf.stop_gradient() automatically puts in GPU if available.
        # So we need to control it using the context manager
        with tf.device(tensorflow_tensor.device):
            tensorflow_tensor = tf.stop_gradient(tensorflow_tensor)
    return tensorflow_tensor

def tensorflow2xp(tensorflow_tensor: "tf.Tensor") -> Array:  # pragma: no cover
    """Convert a Tensorflow tensor to numpy or cupy tensor."""
    assert_tensorflow_installed()
    if "GPU" in tensorflow_tensor.device:
        return cupy.fromDlpack(tfdlpack.to_dlpack(tensorflow_tensor))
    else:
        return tensorflow_tensor.numpy()

__all__ = [
    "fix_random_seed",
    "create_thread_local",
    "is_cupy_array",
    "is_numpy_array",
    "set_active_gpu",
    "prefer_gpu",
    "require_gpu",
    "copy_array",
    "get_shuffled_batches",
    "minibatch",
    "evaluate_model_on_arrays",
    "to_categorical",
    "get_width",
    "xp2torch",
    "torch2xp",
    "tensorflow2xp",
    "xp2tensorflow",
]
<|endoftext|>


<commit_before>from django.conf import settings
from django.conf.urls.defaults import patterns, url
from django.http import HttpResponse
from django.template import Template, Context, loader, TemplateSyntaxError
from .templatetags.spurl import convert_to_boolean
import nose

# This file acts as a urlconf
urlpatterns = patterns('',
    url('^test/$', lambda r: HttpResponse('ok'), name='test')
)

# bootstrap django
settings.configure(
    ROOT_URLCONF='spurl.tests',
    INSTALLED_APPS=['spurl.tests'],
)

# add spurl to builtin tags
loader.add_to_builtins('spurl.templatetags.spurl')

def render(template_string, dictionary=None, autoescape=False):
    """
    Render a template from the supplied string, with optional context data.

    This differs from Django's normal template system in that autoescaping
    is disabled by default. This is simply to make the tests below easier
    to read and write. You can re-enable the default behavior by passing True
    as the value of the autoescape parameter
    """
    context = Context(dictionary, autoescape=autoescape)
    return Template(template_string).render(context)

def test_convert_argument_value_to_boolean():
    assert convert_to_boolean(True) is True
    assert convert_to_boolean(False) is False
    assert convert_to_boolean("True") is True
    assert convert_to_boolean("true") is True
    assert convert_to_boolean("On") is True
    assert convert_to_boolean("on") is True
    assert convert_to_boolean("False") is False
    assert convert_to_boolean("false") is False
    assert convert_to_boolean("Off") is False
    assert convert_to_boolean("off") is False
    assert convert_to_boolean("randomstring") is False

@nose.tools.raises(TemplateSyntaxError)
def test_noargs_raises_exception():
    render("""{% spurl %}""")

@nose.tools.raises(TemplateSyntaxError)
def test_malformed_args_raises_exception():
    render("""{% spurl something %}""")

def test_passthrough():
    template = """{% spurl base="http://www.google.com" %}"""
    rendered = render(template)
    assert rendered == 'http://www.google.com'

def test_url_in_variable():
    template = """{% spurl base=myurl %}"""
    data = {'myurl': 'http://www.google.com'}
    rendered = render(template, data)
    assert rendered == 'http://www.google.com'

def test_make_secure():
    template = """{% spurl base="http://www.google.com" secure="True" %}"""
    rendered = render(template)
    assert rendered == 'https://www.google.com'

def test_make_secure_with_variable():
    template = """{% spurl base=myurl secure=is_secure %}"""
    data = {'myurl': 'http://www.google.com', 'is_secure': True}
    rendered = render(template, data)
    assert rendered == 'https://www.google.com'

def test_make_insecure():
    template = """{% spurl base="https://www.google.com" secure="False" %}"""
    rendered = render(template)
    assert rendered == 'http://www.google.com'

def test_make_insecure_with_variable():
    template = """{% spurl base=myurl secure=is_secure %}"""
    data = {'myurl': 'https://www.google.com', 'is_secure': False}
    rendered = render(template, data)
    assert rendered == 'http://www.google.com'

def test_set_query_from_string():
    template = """{% spurl base="http://www.google.com" query="foo=bar&bar=foo" %}"""
    rendered = render(template)
    assert rendered == 'http://www.google.com?foo=bar&bar=foo'

def test_set_query_from_string_with_variable():
    template = """{% spurl base=myurl query=myquery %}"""
    data = {'myurl': 'http://www.google.com', 'myquery': 'foo=bar&bar=foo'}
    rendered = render(template, data)
    assert rendered == 'http://www.google.com?foo=bar&bar=foo'

def test_set_query_from_dict_with_variable():
    template = """{% spurl base=myurl query=myquery %}"""
    data = {'myurl': 'http://www.google.com', 'myquery': {'foo': 'bar', 'bar': 'foo'}}
    rendered = render(template, data)
    assert rendered == 'http://www.google.com?foo=bar&bar=foo'

def test_set_query_from_template_variables():
    template = """{% spurl base=myurl query="foo={{ first_var }}&bar={{ second_var }}" %}"""
    data = {'myurl': 'http://www.google.com', 'first_var': 'bar', 'second_var': 'baz'}
    rendered = render(template, data)
    assert rendered == 'http://www.google.com?foo=bar&bar=baz'

def test_set_query_from_template_variables_not_double_escaped():
    template = """{% spurl base="http://www.google.com" query="{{ query }}" %}"""
    data = {'query': 'foo=bar&bar=foo'}
    rendered = render(template, data)
    assert rendered == 'http://www.google.com?foo=bar&bar=foo'

def test_set_query_removes_existing_query():
    template = """{% spurl base="http://www.google.com?something=somethingelse" query="foo=bar&bar=foo" %}"""
    rendered = render(template)
    assert rendered == 'http://www.google.com?foo=bar&bar=foo'

def test_add_to_query_from_string():
    template = """{% spurl base="http://www.google.com?something=somethingelse" add_query="foo=bar&bar=foo" %}"""
    rendered = render(template)
    assert rendered == 'http://www.google.com?something=somethingelse&foo=bar&bar=foo'

def test_add_to_query_from_dict_with_variable():
    template = """{% spurl base=myurl add_query=myquery %}"""
    data = {'myurl': 'http://www.google.com?something=somethingelse', 'myquery': {'foo': 'bar', 'bar': 'foo'}}
    rendered = render(template, data)
    assert rendered == 'http://www.google.com?something=somethingelse&foo=bar&bar=foo'

def test_multiple_add_query():
    template = """{% spurl base="http://www.google.com/" add_query="foo=bar" add_query="bar=baz" %}"""
    rendered = render(template)
    assert rendered == 'http://www.google.com/?foo=bar&bar=baz'

def test_add_to_query_from_template_variables():
    template = """{% spurl base="http://www.google.com/?foo=bar" add_query="bar={{ var }}" %}"""
    data = {'var': 'baz'}
    rendered = render(template, data)
    assert rendered == 'http://www.google.com/?foo=bar&bar=baz'

def test_set_query_param_from_string():
    template = """{% spurl base="http://www.google.com?something=somethingelse" set_query="something=foo&somethingelse=bar" %}"""
    rendered = render(template)
    assert rendered == 'http://www.google.com?somethingelse=bar&something=foo'

def test_set_query_param_from_dict_with_variable():
    template = """{% spurl base=myurl set_query=myquery %}"""
    data = {'myurl': 'http://www.google.com?something=somethingelse', 'myquery': {'something': 'foo', 'somethingelse': 'bar'}}
    rendered = render(template, data)
    assert rendered == 'http://www.google.com?somethingelse=bar&something=foo'

def test_multiple_set_query():
    template = """{% spurl base="http://www.google.com/?foo=test" set_query="foo=bar" set_query="bar=baz" %}"""
    rendered = render(template)
    assert rendered == 'http://www.google.com/?foo=bar&bar=baz'

def test_set_query_param_from_template_variables():
    template = """{% spurl base="http://www.google.com/?foo=bar" set_query="foo={{ var }}" %}"""
    data = {'var': 'baz'}
    rendered = render(template, data)
    assert rendered == 'http://www.google.com/?foo=baz'

def test_empty_parameters_preserved():
    template = """{% spurl base="http://www.google.com/?foo=bar" set_query="bar={{ emptyvar }}" %}"""
    data = {} # does not contain and "emptyvar" key
    rendered = render(template, data)
    assert rendered == 'http://www.google.com/?foo=bar&bar='

def test_none_values_are_removed_when_setting_query():
    template = """{% spurl base="http://www.google.com/?foo=bar" set_query="bar={{ nonevar|default_if_none:'' }}" %}"""
    data = {'nonevar': None}
    rendered = render(template, data)
    assert rendered == 'http://www.google.com/?foo=bar&bar='

def test_none_values_are_removed_when_adding_query():
    template = """{% spurl base="http://www.google.com/?foo=bar" add_query="bar={{ nonevar|default_if_none:'' }}" %}"""
    data = {'nonevar': None}
    rendered = render(template, data)
    assert rendered == 'http://www.google.com/?foo=bar&bar='

def test_remove_from_query():
    template = """{% spurl base="http://www.google.com/?foo=bar&bar=baz" remove_query="foo" %}"""
    rendered = render(template)
    assert rendered == 'http://www.google.com/?bar=baz'

def test_remove_multiple_params():
    template = """{% spurl base="http://www.google.com/?foo=bar&bar=baz" remove_query="foo" remove_query="bar" %}"""
    rendered = render(template)
    assert rendered == 'http://www.google.com/'

def test_remove_param_from_template_variable():
    template = """{% spurl base="http://www.google.com/?foo=bar&bar=baz" remove_query=foo remove_query=bar %}"""
    data = {'foo': 'foo', 'bar': 'bar'}
    rendered = render(template, data)
    assert rendered == 'http://www.google.com/'

def test_override_scheme():
    template = """{% spurl base="http://google.com" scheme="ftp" %}"""
    rendered = render(template)
    assert rendered == 'ftp://google.com'

def test_override_host():
    template = """{% spurl base="http://www.google.com/some/path/" host="www.example.com" %}"""
    rendered = render(template)
    assert rendered == 'http://www.example.com/some/path/'

def test_override_path():
    template = """{% spurl base="http://www.google.com/some/path/" path="/another/different/one/" %}"""
    rendered = render(template)
    assert rendered == 'http://www.google.com/another/different/one/'

def test_add_path():
    template = """{% spurl base="http://www.google.com/some/path/" add_path="another/" %}"""
    rendered = render(template)
    assert rendered == 'http://www.google.com/some/path/another/'

def test_multiple_add_path():
    template = """{% spurl base="http://www.google.com/" add_path="some" add_path="another/" %}"""
    rendered = render(template)
    assert rendered == 'http://www.google.com/some/another/'

def test_multiple_add_path_from_template_variables():
    """Usage example for building media urls"""
    template = """{% spurl base=STATIC_URL add_path="javascript" add_path="lib" add_path="jquery.js" %}"""
    data = {'STATIC_URL': 'http://cdn.example.com'}
    rendered = render(template, data)
    assert rendered == 'http://cdn.example.com/javascript/lib/jquery.js'

def test_override_fragment():
    template = """{% spurl base="http://www.google.com/#somefragment" fragment="someotherfragment" %}"""
    rendered = render(template)
    assert rendered == 'http://www.google.com/#someotherfragment'

def test_override_port():
    template = """{% spurl base="http://www.google.com:80" port="8080" %}"""
    rendered = render(template)
    assert rendered == 'http://www.google.com:8080'

def test_build_complete_url():
    template = """{% spurl scheme="http" host="www.google.com" path="/some/path/" port="8080" fragment="somefragment" %}"""
    rendered = render(template)
    assert rendered == 'http://www.google.com:8080/some/path/#somefragment'

def test_sensible_defaults():
    template = """{% spurl path="/some/path/" %}"""
    rendered = render(template)
    assert rendered == '/some/path/'

    template = """{% spurl path="/some/path/" host="www.google.com" %}"""
    rendered = render(template)
    assert rendered == 'http://www.google.com/some/path/'

def test_autoescaping():
    template = """{% spurl base="http://www.google.com" query="a=b" add_query="c=d" add_query="e=f" fragment="frag" %}"""
    rendered = render(template, autoescape=True) # Ordinarily, templates will be autoescaped by default
    assert rendered == 'http://www.google.com?a=b&c=d&e=f#frag'

def test_disable_autoescaping_with_parameter():
    template = """{% spurl base="http://www.google.com" query="a=b" add_query="c=d" autoescape="False" %}"""
    rendered = render(template, autoescape=True)
    assert rendered == 'http://www.google.com?a=b&c=d'

def test_url_as_template_variable():
    template = """{% spurl base="http://www.google.com" as foo %}The url is {{ foo }}"""
    rendered = render(template)
    assert rendered == 'The url is http://www.google.com'

def test_reversing_inside_spurl_tag():
    template = """{% load url from future %}{% spurl base="http://www.google.com/" path='{\% url "test" %\}' %}"""
    rendered = render(template)
    assert rendered == 'http://www.google.com/test/'

    template = """{% load url from future %}{% spurl base="http://www.google.com/" query='next={\% url "test" %\}' %}"""
    rendered = render(template)
    assert rendered == 'http://www.google.com/?next=/test/'
<commit_msg>Add test for extreme nastiness<commit_after>from django.conf import settings
from django.conf.urls.defaults import patterns, url
from django.http import HttpResponse
from django.template import Template, Context, loader, TemplateSyntaxError
from .templatetags.spurl import convert_to_boolean
import nose

# This file acts as a urlconf
urlpatterns = patterns('',
    url('^test/$', lambda r: HttpResponse('ok'), name='test')
)

# bootstrap django
settings.configure(
    ROOT_URLCONF='spurl.tests',
    INSTALLED_APPS=['spurl.tests'],
)

# add spurl to builtin tags
loader.add_to_builtins('spurl.templatetags.spurl')

def render(template_string, dictionary=None, autoescape=False):
    """
    Render a template from the supplied string, with optional context data.

    This differs from Django's normal template system in that autoescaping
    is disabled by default. This is simply to make the tests below easier
    to read and write. You can re-enable the default behavior by passing True
    as the value of the autoescape parameter
    """
    context = Context(dictionary, autoescape=autoescape)
    return Template(template_string).render(context)

def test_convert_argument_value_to_boolean():
    assert convert_to_boolean(True) is True
    assert convert_to_boolean(False) is False
    assert convert_to_boolean("True") is True
    assert convert_to_boolean("true") is True
    assert convert_to_boolean("On") is True
    assert convert_to_boolean("on") is True
    assert convert_to_boolean("False") is False
    assert convert_to_boolean("false") is False
    assert convert_to_boolean("Off") is False
    assert convert_to_boolean("off") is False
    assert convert_to_boolean("randomstring") is False

@nose.tools.raises(TemplateSyntaxError)
def test_noargs_raises_exception():
    render("""{% spurl %}""")

@nose.tools.raises(TemplateSyntaxError)
def test_malformed_args_raises_exception():
    render("""{% spurl something %}""")

def test_passthrough():
    template = """{% spurl base="http://www.google.com" %}"""
    rendered = render(template)
    assert rendered == 'http://www.google.com'

def test_url_in_variable():
    template = """{% spurl base=myurl %}"""
    data = {'myurl': 'http://www.google.com'}
    rendered = render(template, data)
    assert rendered == 'http://www.google.com'

def test_make_secure():
    template = """{% spurl base="http://www.google.com" secure="True" %}"""
    rendered = render(template)
    assert rendered == 'https://www.google.com'

def test_make_secure_with_variable():
    template = """{% spurl base=myurl secure=is_secure %}"""
    data = {'myurl': 'http://www.google.com', 'is_secure': True}
    rendered = render(template, data)
    assert rendered == 'https://www.google.com'

def test_make_insecure():
    template = """{% spurl base="https://www.google.com" secure="False" %}"""
    rendered = render(template)
    assert rendered == 'http://www.google.com'

def test_make_insecure_with_variable():
    template = """{% spurl base=myurl secure=is_secure %}"""
    data = {'myurl': 'https://www.google.com', 'is_secure': False}
    rendered = render(template, data)
    assert rendered == 'http://www.google.com'

def test_set_query_from_string():
    template = """{% spurl base="http://www.google.com" query="foo=bar&bar=foo" %}"""
    rendered = render(template)
    assert rendered == 'http://www.google.com?foo=bar&bar=foo'

def test_set_query_from_string_with_variable():
    template = """{% spurl base=myurl query=myquery %}"""
    data = {'myurl': 'http://www.google.com', 'myquery': 'foo=bar&bar=foo'}
    rendered = render(template, data)
    assert rendered == 'http://www.google.com?foo=bar&bar=foo'

def test_set_query_from_dict_with_variable():
    template = """{% spurl base=myurl query=myquery %}"""
    data = {'myurl': 'http://www.google.com', 'myquery': {'foo': 'bar', 'bar': 'foo'}}
    rendered = render(template, data)
    assert rendered == 'http://www.google.com?foo=bar&bar=foo'

def test_set_query_from_template_variables():
    template = """{% spurl base=myurl query="foo={{ first_var }}&bar={{ second_var }}" %}"""
    data = {'myurl': 'http://www.google.com', 'first_var': 'bar', 'second_var': 'baz'}
    rendered = render(template, data)
    assert rendered == 'http://www.google.com?foo=bar&bar=baz'

def test_set_query_from_template_variables_not_double_escaped():
    template = """{% spurl base="http://www.google.com" query="{{ query }}" %}"""
    data = {'query': 'foo=bar&bar=foo'}
    rendered = render(template, data)
    assert rendered == 'http://www.google.com?foo=bar&bar=foo'

def test_set_query_removes_existing_query():
    template = """{% spurl base="http://www.google.com?something=somethingelse" query="foo=bar&bar=foo" %}"""
    rendered = render(template)
    assert rendered == 'http://www.google.com?foo=bar&bar=foo'

def test_add_to_query_from_string():
    template = """{% spurl base="http://www.google.com?something=somethingelse" add_query="foo=bar&bar=foo" %}"""
    rendered = render(template)
    assert rendered == 'http://www.google.com?something=somethingelse&foo=bar&bar=foo'

def test_add_to_query_from_dict_with_variable():
    template = """{% spurl base=myurl add_query=myquery %}"""
    data = {'myurl': 'http://www.google.com?something=somethingelse', 'myquery': {'foo': 'bar', 'bar': 'foo'}}
    rendered = render(template, data)
    assert rendered == 'http://www.google.com?something=somethingelse&foo=bar&bar=foo'

def test_multiple_add_query():
    template = """{% spurl base="http://www.google.com/" add_query="foo=bar" add_query="bar=baz" %}"""
    rendered = render(template)
    assert rendered == 'http://www.google.com/?foo=bar&bar=baz'

def test_add_to_query_from_template_variables():
    template = """{% spurl base="http://www.google.com/?foo=bar" add_query="bar={{ var }}" %}"""
    data = {'var': 'baz'}
    rendered = render(template, data)
    assert rendered == 'http://www.google.com/?foo=bar&bar=baz'

def test_set_query_param_from_string():
    template = """{% spurl base="http://www.google.com?something=somethingelse" set_query="something=foo&somethingelse=bar" %}"""
    rendered = render(template)
    assert rendered == 'http://www.google.com?somethingelse=bar&something=foo'

def test_set_query_param_from_dict_with_variable():
    template = """{% spurl base=myurl set_query=myquery %}"""
    data = {'myurl': 'http://www.google.com?something=somethingelse', 'myquery': {'something': 'foo', 'somethingelse': 'bar'}}
    rendered = render(template, data)
    assert rendered == 'http://www.google.com?somethingelse=bar&something=foo'

def test_multiple_set_query():
    template = """{% spurl base="http://www.google.com/?foo=test" set_query="foo=bar" set_query="bar=baz" %}"""
    rendered = render(template)
    assert rendered == 'http://www.google.com/?foo=bar&bar=baz'

def test_set_query_param_from_template_variables():
    template = """{% spurl base="http://www.google.com/?foo=bar" set_query="foo={{ var }}" %}"""
    data = {'var': 'baz'}
    rendered = render(template, data)
    assert rendered == 'http://www.google.com/?foo=baz'

def test_empty_parameters_preserved():
    template = """{% spurl base="http://www.google.com/?foo=bar" set_query="bar={{ emptyvar }}" %}"""
    data = {} # does not contain and "emptyvar" key
    rendered = render(template, data)
    assert rendered == 'http://www.google.com/?foo=bar&bar='

def test_none_values_are_removed_when_setting_query():
    template = """{% spurl base="http://www.google.com/?foo=bar" set_query="bar={{ nonevar|default_if_none:'' }}" %}"""
    data = {'nonevar': None}
    rendered = render(template, data)
    assert rendered == 'http://www.google.com/?foo=bar&bar='

def test_none_values_are_removed_when_adding_query():
    template = """{% spurl base="http://www.google.com/?foo=bar" add_query="bar={{ nonevar|default_if_none:'' }}" %}"""
    data = {'nonevar': None}
    rendered = render(template, data)
    assert rendered == 'http://www.google.com/?foo=bar&bar='

def test_remove_from_query():
    template = """{% spurl base="http://www.google.com/?foo=bar&bar=baz" remove_query="foo" %}"""
    rendered = render(template)
    assert rendered == 'http://www.google.com/?bar=baz'

def test_remove_multiple_params():
    template = """{% spurl base="http://www.google.com/?foo=bar&bar=baz" remove_query="foo" remove_query="bar" %}"""
    rendered = render(template)
    assert rendered == 'http://www.google.com/'

def test_remove_param_from_template_variable():
    template = """{% spurl base="http://www.google.com/?foo=bar&bar=baz" remove_query=foo remove_query=bar %}"""
    data = {'foo': 'foo', 'bar': 'bar'}
    rendered = render(template, data)
    assert rendered == 'http://www.google.com/'

def test_override_scheme():
    template = """{% spurl base="http://google.com" scheme="ftp" %}"""
    rendered = render(template)
    assert rendered == 'ftp://google.com'

def test_override_host():
    template = """{% spurl base="http://www.google.com/some/path/" host="www.example.com" %}"""
    rendered = render(template)
    assert rendered == 'http://www.example.com/some/path/'

def test_override_path():
    template = """{% spurl base="http://www.google.com/some/path/" path="/another/different/one/" %}"""
    rendered = render(template)
    assert rendered == 'http://www.google.com/another/different/one/'

def test_add_path():
    template = """{% spurl base="http://www.google.com/some/path/" add_path="another/" %}"""
    rendered = render(template)
    assert rendered == 'http://www.google.com/some/path/another/'

def test_multiple_add_path():
    template = """{% spurl base="http://www.google.com/" add_path="some" add_path="another/" %}"""
    rendered = render(template)
    assert rendered == 'http://www.google.com/some/another/'

def test_multiple_add_path_from_template_variables():
    """Usage example for building media urls"""
    template = """{% spurl base=STATIC_URL add_path="javascript" add_path="lib" add_path="jquery.js" %}"""
    data = {'STATIC_URL': 'http://cdn.example.com'}
    rendered = render(template, data)
    assert rendered == 'http://cdn.example.com/javascript/lib/jquery.js'

def test_override_fragment():
    template = """{% spurl base="http://www.google.com/#somefragment" fragment="someotherfragment" %}"""
    rendered = render(template)
    assert rendered == 'http://www.google.com/#someotherfragment'

def test_override_port():
    template = """{% spurl base="http://www.google.com:80" port="8080" %}"""
    rendered = render(template)
    assert rendered == 'http://www.google.com:8080'

def test_build_complete_url():
    template = """{% spurl scheme="http" host="www.google.com" path="/some/path/" port="8080" fragment="somefragment" %}"""
    rendered = render(template)
    assert rendered == 'http://www.google.com:8080/some/path/#somefragment'

def test_sensible_defaults():
    template = """{% spurl path="/some/path/" %}"""
    rendered = render(template)
    assert rendered == '/some/path/'

    template = """{% spurl path="/some/path/" host="www.google.com" %}"""
    rendered = render(template)
    assert rendered == 'http://www.google.com/some/path/'

def test_autoescaping():
    template = """{% spurl base="http://www.google.com" query="a=b" add_query="c=d" add_query="e=f" fragment="frag" %}"""
    rendered = render(template, autoescape=True) # Ordinarily, templates will be autoescaped by default
    assert rendered == 'http://www.google.com?a=b&c=d&e=f#frag'

def test_disable_autoescaping_with_parameter():
    template = """{% spurl base="http://www.google.com" query="a=b" add_query="c=d" autoescape="False" %}"""
    rendered = render(template, autoescape=True)
    assert rendered == 'http://www.google.com?a=b&c=d'

def test_url_as_template_variable():
    template = """{% spurl base="http://www.google.com" as foo %}The url is {{ foo }}"""
    rendered = render(template)
    assert rendered == 'The url is http://www.google.com'

def test_reversing_inside_spurl_tag():
    template = """{% load url from future %}{% spurl base="http://www.google.com/" path="{\% url 'test' %\}" %}"""
    rendered = render(template)
    assert rendered == 'http://www.google.com/test/'

    template = """{% load url from future %}{% spurl base="http://www.google.com/" query="next={\% url 'test' %\}" %}"""
    rendered = render(template)
    assert rendered == 'http://www.google.com/?next=/test/'

def test_xzibit():
    template = """Yo dawg, the URL is: {% spurl base="http://www.google.com/" query="foo={\% spurl base='http://another.com' secure='true' %\}" %}"""
    rendered = render(template)
    assert rendered == 'Yo dawg, the URL is: http://www.google.com/?foo=https://another.com'
<|endoftext|>


<commit_before>'''
Copyright (c) 2014 Juju. Inc

Code Licensed under MIT License. See LICENSE file.
'''
import rel_imp; rel_imp.init()

from selenium import webdriver
from pyvirtualdisplay import Display
from threading import RLock
from functools import wraps
from selenium.common.exceptions import UnexpectedAlertPresentException
from .base import XpathWdBase, singleton_decorator
from xpathwebdriver.levels import SURVIVE_PROCESS, TEST_ROUND_LIFE, MANAGER_LIFE
from selenium.webdriver.remote.webdriver import WebDriver
import os

def synchronized(lock):
    '''
    Thread synchronization decorator. (for several methods with same lock, use
    re-entrant lock)
    :param lock: lock object (from threading package for example)
    '''
    def wrap_method(method):
        @wraps(method)
        def newFunction(*args, **kw):
            with lock:
                return method(*args, **kw)
        return newFunction
    return wrap_method

@singleton_decorator
class WebdriverManager(XpathWdBase):
    '''
    This is a "Context Manager" for the Webdriver's instances available. (used
    for running tests in general)
    '''
    # Thread lock for methods in this class
    _methods_lock = RLock()

    def __init__(self):
        # Set of locked webdrivers
        self._locked = set()
        # Set of released webdrivers
        self._released = set()
        # Pool of drivers {wdriver:(browser_name, driver_life_level)}
        self._wdriver_pool = {}
        # Virtual display object where we start webdrivers
        self._virtual_display = None
        # Any browser without specified level
        self._current_context_level = -1
        if (self.global_settings.get('webdriver_browser_keep_open')
        or (self.global_settings.get('webdriver_remote_command_executor')
            and self.global_settings.get('webdriver_remote_session_id'))):
            # Start a never killed process
            self.init_level(SURVIVE_PROCESS)

    @property
    def enabled(self):
        return self.global_settings.get('webdriver_enabled')

    @synchronized(_methods_lock)
    def enter_level(self, level=None, base_url=None, name='', browser_name=None):
        if not level:
            # There is no level declared, so it will be only valid for the life
            # of the context (assuming "with manager.enter_level() as browser: ..." etc)
            level = self._current_context_level
            self._current_context_level -= 1
        self.init_level(level, browser_name)
        return BrowserContextManager(self, level, base_url, name)

    @synchronized(_methods_lock)
    def exit_level(self, level):
        def quit_wdriver(wdriver, container):
            _, drv_level = self._wdriver_pool[wdriver]
            if self._quit_failed_webdriver(wdriver):
                container.remove(wdriver)
                self._wdriver_pool.pop(wdriver)
            elif drv_level <= level:
                container.remove(wdriver)
                self._quit_webdriver(wdriver)
                self._wdriver_pool.pop(wdriver)
            else:
                self.log.d('Keeping webdriver %s', wdriver)
        # Make copies of sets, we are going to modify them
        for wdriver in self._released.copy():
            quit_wdriver(wdriver, self._released)
        for wdriver in self._locked.copy():
            quit_wdriver(wdriver, self._locked)

    @synchronized(_methods_lock)
    def init_level(self, level, browser_name=None):
        '''
        Initialize webdriver for a specific Life Level.
        There are two possibilities:
          1- no webdriver was initialized on previous level => init it
          2- there is a webdriver for the browser selected, no action is performed.
        :param level: webdriver's level of life we are entering
        '''
        if not self.enabled:
            return
        # level is mandatory
        assert level and level <= SURVIVE_PROCESS, 'Not a browser level'
        # Get rid of non-responding browsers
        self.quit_all_failed_webdrivers()
        # Get the set of released webdrivers for the selected browser
        released = self.get_available_set(browser_name)
        # If no webdriver available, then create a new one
        if not released:
            # Create webdriver if needed
            browser = self.get_browser_name(browser_name)
            wdriver = self._new_webdriver(browser, level)
            self._wdriver_pool[wdriver] = browser, level
            self._released.add(wdriver)

    @synchronized(_methods_lock)
    def get_available_set(self, browser_name):
        '''
        Return the set of avialable webdrivers for the current Browser selected
        (in global configuration)
        '''
        browser = self.get_browser_name(browser_name)
        browser_set = set([wdriver
                           for wdriver,(brws,_) in self._wdriver_pool.items()
                           if brws == browser])
        return (self._released & browser_set)

    def _new_webdriver(self, browser, level, args=None, kwargs=None):
        args = args or []
        kwargs = kwargs or {}
        browser = self.expand_browser_name(browser)
        # Setup display before creating the browser
        self.setup_display()
        if (level == SURVIVE_PROCESS
        and self.global_settings.get('webdriver_remote_command_executor')
        and self.global_settings.get('webdriver_remote_session_id')):
            return self._build_remote(self.global_settings.get('webdriver_remote_command_executor'),
                                      self.global_settings.get('webdriver_remote_session_id'))
        if browser == 'PhantomJS':
            self._append_service_arg('--ignore-ssl-errors=true', kwargs)
        if (browser == 'Firefox'
        and self.global_settings.get('webdriver_firefox_profile')
        and not args and not kwargs.has_key('firefox_profile')):
            # Update with profile specified from config
            fp = webdriver.FirefoxProfile(self.global_settings.get('webdriver_firefox_profile'))
            kwargs['firefox_profile'] = fp
        if (browser == 'Chrome' and os.name == 'posix' and os.geteuid() == 0):
            self.log.w('Passing --no-sandbox flag to Chrome (running as root)')
            from selenium.webdriver.chrome.options import Options
            chrome_options = Options()
            chrome_options.add_argument('--no-sandbox')
            #self.log.w('Adding --disable-application-cache')
            #chrome_options.add_argument('--disable-application-cache')
            #chrome_options.add_argument('--incognito')
            kwargs['chrome_options'] = chrome_options
        driver = getattr(webdriver, browser)(*args, **kwargs)
        return driver

    def _build_remote(self, executor_url, session_id):
        original_execute = WebDriver.execute
        non_local = dict(first_run=True)
        def _patched_execute(self, command, params=None):
            if command == "newSession" and non_local['first_run']:
                non_local['first_run'] = False
                # Mock the response
                return {'success': 0, 'value': None, 'sessionId': session_id}
            else:
                return original_execute(self, command, params)
        # Patch the function before creating the driver object
        WebDriver.execute = _patched_execute
        driver = webdriver.Remote(command_executor=executor_url)
        driver.session_id = session_id
        # Replace the patched function with original function
        WebDriver.execute = original_execute
        return driver

    def _append_service_arg(self, arg, kwargs):
        service_args = kwargs.get('service_args', [])
        service_args.append(arg)
        kwargs['service_args'] = service_args

    @synchronized(_methods_lock)
    def acquire_driver(self, level, browser_name=None):
        if not self.enabled:
            return None
        self.init_level(level, browser_name)
        wdriver = self.get_available_set(browser_name).pop()
        # Keep track of acquired webdrivers in case we need to close them
        self._locked.add(wdriver)
        self._released.remove(wdriver)
        return wdriver

    @synchronized(_methods_lock)
    def release_driver(self, wdriver, level):
        '''
        Release passed webdriver instance, so it becomes available to be used
        by another test. (this function is not generally used directly)
        :param wdriver: webdriver instance to be released
        :param level: level at wich we are releasing the webdriver
        '''
        if not wdriver:
            return
        assert wdriver in self._locked, 'Webdriver %r was never locked' % wdriver
        self._locked.remove(wdriver)
        _, drv_level = self._wdriver_pool[wdriver]
        # Make sure webdriver is healthy and from the right level
        # before reusing it
        if drv_level >= level and not self._quit_failed_webdriver(wdriver):
            # Keep webdriver with higher level of life
            self._released.add(wdriver)
        else:
            # Remove no longer needed webdriver
            self._quit_webdriver(wdriver)

    @synchronized(_methods_lock)
    def quit_all_webdrivers(self):
        for wdriver in self._wdriver_pool:
            self._quit_webdriver(wdriver)
        self._wdriver_pool.clear()
        self._locked.clear()
        self._released.clear()

    def _quit_webdriver(self, wdriver):
        # quit a webdriver, catch any exception and report it
        try:
            self.log.d('Quitting webdriver %s', wdriver)
            wdriver.quit()
        except Exception as e:
            self.log.w('Ignoring %r:%s' % (e,e))

    @synchronized(_methods_lock)
    def quit_all_failed_webdrivers(self):
        remove = lambda s,e: e in s and s.remove(e)
        for wdriver in list(self._wdriver_pool.keys()):
            if self._quit_failed_webdriver(wdriver):
                self._wdriver_pool.pop(wdriver)
                remove(self._locked, wdriver)
                remove(self._released, wdriver)

    def _quit_failed_webdriver(self, wdriver, tested_once=False):
        # Test wether a webdriver is responding, if not quit it and
        # unregister it to avoid further usage.
        try:
            wdriver.current_url
            return False
        except UnexpectedAlertPresentException as e:
            # An alert window will cause an exception, handle it only once
            alert = wdriver.switch_to_alert()
            alert.accept()
            if not tested_once:
                return self._quit_failed_webdriver(wdriver, tested_once=True)
        except Exception as e:
            self.log.e('Quitting webdriver %r due to exception %r:%s' %
                       (wdriver, e, e))
        self._quit_webdriver(wdriver)
        return True

    @synchronized(_methods_lock)
    def setup_display(self):
        '''
        Create virtual display if set by configuration
        '''
        def get(name, default=None):
            return self.global_settings.get('virtual_display_' + name, default)
        if not get('enabled'):
            if self._virtual_display:
                self.log.w('There is a display enabled although config says'
                           ' different')
            return
        elif self._virtual_display:
            # There is a display configured
            return
        # We need to setup a new virtual display
        d = Display(backend=get('backend'), size=get('size', (800, 600)),
                visible=get('visible'))
        d.start()
        self._virtual_display = d

    @synchronized(_methods_lock)
    def stop_display(self):
        '''
        Convenient function to stop the virtual display
        '''
        # Nice alias
        display = self._virtual_display
        if ((not self.global_settings.get('virtual_display_keep_open')
             or not self.global_settings.get('virtual_display_visible'))
                and display):
            self.log.d('Stopping virtual display %r' % display)
            display.stop()
            WebdriverManager._virtual_display = None

    @synchronized(_methods_lock)
    def list_webdrivers(self, which='all'):
        '''
        Make a report of all, released or locked webdrivers.
        :param which: string specifiying the type of report. May be:'all', 'released', 'locked'
        '''
        wdrivers_report = {}
        possible = ('all', 'released', 'locked')
        assert which in possible, 'Which must be one of %r' % possible
        # Choose the set to report
        if which =='all':
            in_report = self._wdriver_pool
        elif which == 'released':
            in_report = self._released
        elif which == 'locked':
            in_report = self._locked
        # Build report
        for wdriver in in_report:
            browser, level = self._wdriver_pool[wdriver]
            wdrivers_report[wdriver] = browser, level
        return wdrivers_report

    def get_browser_name(self, browser_name=None):
        '''
        Get the current browser name in usage
        If not set, we use the one specified in global settings
        or PhantomJS as the final default
        '''
        browser = (browser_name
                   or self.global_settings.get('webdriver_browser'))
        return self.expand_browser_name(browser)

    def expand_browser_name(self, browser=None):
        '''
        Expand a partial or lower case Browser's name into its full name.
        Eg:
            firefo -> Firefox
            chrom -> Chrome
        :param browser: partial browser's name string
        '''
        # TODO: add IE and Opera
        char_browser = dict(f='Firefox',
                            c='Chrome',
                            p='PhantomJS',
                            )
        char = browser.lower()[0]
        assert char in char_browser, 'Could not find browser %r' % browser
        return char_browser.get(char)

    def __del__(self):
        self.exit_level(MANAGER_LIFE)
        self.stop_display()

class BrowserContextManager(XpathWdBase):
    def __init__(self, parent, level, base_url=None, name='', browser_name=None):
        self.parent = parent
        self.level = level
        self.base_url = base_url
        self.name = name
        self.webdriver = None
        self.browser_name = browser_name

    def __enter__(self):
        return self.get_xpathbrowser()

    def __exit__(self, type=None, value=None, traceback=None):
        self.release_driver()
        self.parent.exit_level(self.level)

    def acquire_driver(self):
        if not self.webdriver:
            self.webdriver = self.parent.acquire_driver(self.level, self.browser_name)
        return self.webdriver

    def get_xpathbrowser(self, base_url=None, name=''):
        from .logger import Logger
        from .xpath_browser import XpathBrowser
        base_url = base_url or self.base_url or self.global_settings.get('base_url')
        name = name or self.name
        # Initialize the XpathBrowser class
        return XpathBrowser(self.acquire_driver(), base_url,
                            Logger(name), settings=self.global_settings)

    def release_driver(self):
        if self.webdriver:
            self.parent.release_driver(self.webdriver, self.level)
            self.webdriver = None

def smoke_test_module():
    from .logger import log_test
    mngr = WebdriverManager()
#    import ipdb; ipdb.set_trace()
    lvl = mngr.enter_level(level=TEST_ROUND_LIFE)
    wd = lvl.acquire_driver()
    log_test(wd)
    lvl.release_driver()
    log_test(mngr.list_webdrivers())
#    import ipdb; ipdb.set_trace()
    xb = lvl.get_xpathbrowser()
    log_test(xb)
    lvl.exit_level()
#    import ipdb; ipdb.set_trace()
    mngr.stop_display()
    mngr.quit_all_webdrivers()
    # Wait some time for browser to quit
    import time
    time.sleep(.5)

if __name__ == "__main__":
    smoke_test_module()
<commit_msg>Add docs<commit_after>'''
Copyright (c) 2014 Juju. Inc

Code Licensed under MIT License. See LICENSE file.
'''
import rel_imp; rel_imp.init()

from selenium import webdriver
from pyvirtualdisplay import Display
from threading import RLock
from functools import wraps
from selenium.common.exceptions import UnexpectedAlertPresentException
from .base import XpathWdBase, singleton_decorator
from xpathwebdriver.levels import SURVIVE_PROCESS, TEST_ROUND_LIFE, MANAGER_LIFE
from selenium.webdriver.remote.webdriver import WebDriver
import os

def synchronized(lock):
    '''
    Thread synchronization decorator. (for several methods with same lock, use
    re-entrant lock)
    :param lock: lock object (from threading package for example)
    '''
    def wrap_method(method):
        @wraps(method)
        def newFunction(*args, **kw):
            with lock:
                return method(*args, **kw)
        return newFunction
    return wrap_method

@singleton_decorator
class WebdriverManager(XpathWdBase):
    '''
    This is a "Context Manager" for the Webdriver's instances available. (used
    for running tests in general)
    '''
    # Thread lock for methods in this class
    _methods_lock = RLock()

    def __init__(self):
        # Set of locked webdrivers
        self._locked = set()
        # Set of released webdrivers
        self._released = set()
        # Pool of drivers {wdriver:(browser_name, driver_life_level)}
        self._wdriver_pool = {}
        # Virtual display object where we start webdrivers
        self._virtual_display = None
        # Any browser without specified level
        self._current_context_level = -1
        if (self.global_settings.get('webdriver_browser_keep_open')
        or (self.global_settings.get('webdriver_remote_command_executor')
            and self.global_settings.get('webdriver_remote_session_id'))):
            # Start a never killed process
            self.init_level(SURVIVE_PROCESS)

    @property
    def enabled(self):
        return self.global_settings.get('webdriver_enabled')

    @synchronized(_methods_lock)
    def enter_level(self, level=None, base_url=None, name='', browser_name=None):
        '''
        :param level: webdriver's level of life we are entering
        :param base_url: optional set a base url for the browser (to specify paths)
        :param name: optional level's name (mainly for logging purposes)
        :param browser_name: optional browser name string (eg: 'Firefox', 'Chrome', 'PhantomJs')
        '''
        if not level:
            # There is no level declared, so it will be only valid for the life
            # of the context (assuming "with manager.enter_level() as browser: ..." etc)
            level = self._current_context_level
            self._current_context_level -= 1
        self.init_level(level, browser_name)
        return BrowserContextManager(self, level, base_url, name)

    @synchronized(_methods_lock)
    def exit_level(self, level):
        '''
        :param level: webdriver's level of life we are exiting
        '''
        def quit_wdriver(wdriver, container):
            _, drv_level = self._wdriver_pool[wdriver]
            if self._quit_failed_webdriver(wdriver):
                container.remove(wdriver)
                self._wdriver_pool.pop(wdriver)
            elif drv_level <= level:
                container.remove(wdriver)
                self._quit_webdriver(wdriver)
                self._wdriver_pool.pop(wdriver)
            else:
                self.log.d('Keeping webdriver %s', wdriver)
        # Make copies of sets, we are going to modify them
        for wdriver in self._released.copy():
            quit_wdriver(wdriver, self._released)
        for wdriver in self._locked.copy():
            quit_wdriver(wdriver, self._locked)

    @synchronized(_methods_lock)
    def init_level(self, level, browser_name=None):
        '''
        Initialize webdriver for a specific Life Level.
        There are two possibilities:
          1- no webdriver was initialized on previous level => init it
          2- there is a webdriver for the browser selected, no action is performed.
        :param level: webdriver's level of life we are entering
        '''
        if not self.enabled:
            return
        # level is mandatory
        assert level and level <= SURVIVE_PROCESS, 'Not a browser level'
        # Get rid of non-responding browsers
        self.quit_all_failed_webdrivers()
        # Get the set of released webdrivers for the selected browser
        released = self.get_available_set(browser_name)
        # If no webdriver available, then create a new one
        if not released:
            # Create webdriver if needed
            browser = self.get_browser_name(browser_name)
            wdriver = self._new_webdriver(browser, level)
            self._wdriver_pool[wdriver] = browser, level
            self._released.add(wdriver)

    @synchronized(_methods_lock)
    def get_available_set(self, browser_name):
        '''
        Return the set of avialable webdrivers for the current Browser selected
        (in global configuration)
        '''
        browser = self.get_browser_name(browser_name)
        browser_set = set([wdriver
                           for wdriver,(brws,_) in self._wdriver_pool.items()
                           if brws == browser])
        return (self._released & browser_set)

    def _new_webdriver(self, browser, level, args=None, kwargs=None):
        args = args or []
        kwargs = kwargs or {}
        browser = self.expand_browser_name(browser)
        # Setup display before creating the browser
        self.setup_display()
        if (level == SURVIVE_PROCESS
        and self.global_settings.get('webdriver_remote_command_executor')
        and self.global_settings.get('webdriver_remote_session_id')):
            return self._build_remote(self.global_settings.get('webdriver_remote_command_executor'),
                                      self.global_settings.get('webdriver_remote_session_id'))
        if browser == 'PhantomJS':
            self._append_service_arg('--ignore-ssl-errors=true', kwargs)
        if (browser == 'Firefox'
        and self.global_settings.get('webdriver_firefox_profile')
        and not args and not kwargs.has_key('firefox_profile')):
            # Update with profile specified from config
            fp = webdriver.FirefoxProfile(self.global_settings.get('webdriver_firefox_profile'))
            kwargs['firefox_profile'] = fp
        if (browser == 'Chrome' and os.name == 'posix' and os.geteuid() == 0):
            self.log.w('Passing --no-sandbox flag to Chrome (running as root)')
            from selenium.webdriver.chrome.options import Options
            chrome_options = Options()
            chrome_options.add_argument('--no-sandbox')
            #self.log.w('Adding --disable-application-cache')
            #chrome_options.add_argument('--disable-application-cache')
            #chrome_options.add_argument('--incognito')
            kwargs['chrome_options'] = chrome_options
        driver = getattr(webdriver, browser)(*args, **kwargs)
        return driver

    def _build_remote(self, executor_url, session_id):
        original_execute = WebDriver.execute
        non_local = dict(first_run=True)
        def _patched_execute(self, command, params=None):
            if command == "newSession" and non_local['first_run']:
                non_local['first_run'] = False
                # Mock the response
                return {'success': 0, 'value': None, 'sessionId': session_id}
            else:
                return original_execute(self, command, params)
        # Patch the function before creating the driver object
        WebDriver.execute = _patched_execute
        driver = webdriver.Remote(command_executor=executor_url)
        driver.session_id = session_id
        # Replace the patched function with original function
        WebDriver.execute = original_execute
        return driver

    def _append_service_arg(self, arg, kwargs):
        service_args = kwargs.get('service_args', [])
        service_args.append(arg)
        kwargs['service_args'] = service_args

    @synchronized(_methods_lock)
    def acquire_driver(self, level, browser_name=None):
        if not self.enabled:
            return None
        self.init_level(level, browser_name)
        wdriver = self.get_available_set(browser_name).pop()
        # Keep track of acquired webdrivers in case we need to close them
        self._locked.add(wdriver)
        self._released.remove(wdriver)
        return wdriver

    @synchronized(_methods_lock)
    def release_driver(self, wdriver, level):
        '''
        Release passed webdriver instance, so it becomes available to be used
        by another test. (this function is not generally used directly)
        :param wdriver: webdriver instance to be released
        :param level: level at wich we are releasing the webdriver
        '''
        if not wdriver:
            return
        assert wdriver in self._locked, 'Webdriver %r was never locked' % wdriver
        self._locked.remove(wdriver)
        _, drv_level = self._wdriver_pool[wdriver]
        # Make sure webdriver is healthy and from the right level
        # before reusing it
        if drv_level >= level and not self._quit_failed_webdriver(wdriver):
            # Keep webdriver with higher level of life
            self._released.add(wdriver)
        else:
            # Remove no longer needed webdriver
            self._quit_webdriver(wdriver)

    @synchronized(_methods_lock)
    def quit_all_webdrivers(self):
        for wdriver in self._wdriver_pool:
            self._quit_webdriver(wdriver)
        self._wdriver_pool.clear()
        self._locked.clear()
        self._released.clear()

    def _quit_webdriver(self, wdriver):
        # quit a webdriver, catch any exception and report it
        try:
            self.log.d('Quitting webdriver %s', wdriver)
            wdriver.quit()
        except Exception as e:
            self.log.w('Ignoring %r:%s' % (e,e))

    @synchronized(_methods_lock)
    def quit_all_failed_webdrivers(self):
        remove = lambda s,e: e in s and s.remove(e)
        for wdriver in list(self._wdriver_pool.keys()):
            if self._quit_failed_webdriver(wdriver):
                self._wdriver_pool.pop(wdriver)
                remove(self._locked, wdriver)
                remove(self._released, wdriver)

    def _quit_failed_webdriver(self, wdriver, tested_once=False):
        # Test wether a webdriver is responding, if not quit it and
        # unregister it to avoid further usage.
        try:
            wdriver.current_url
            return False
        except UnexpectedAlertPresentException as e:
            # An alert window will cause an exception, handle it only once
            alert = wdriver.switch_to_alert()
            alert.accept()
            if not tested_once:
                return self._quit_failed_webdriver(wdriver, tested_once=True)
        except Exception as e:
            self.log.e('Quitting webdriver %r due to exception %r:%s' %
                       (wdriver, e, e))
        self._quit_webdriver(wdriver)
        return True

    @synchronized(_methods_lock)
    def setup_display(self):
        '''
        Create virtual display if set by configuration
        '''
        def get(name, default=None):
            return self.global_settings.get('virtual_display_' + name, default)
        if not get('enabled'):
            if self._virtual_display:
                self.log.w('There is a display enabled although config says'
                           ' different')
            return
        elif self._virtual_display:
            # There is a display configured
            return
        # We need to setup a new virtual display
        d = Display(backend=get('backend'), size=get('size', (800, 600)),
                visible=get('visible'))
        d.start()
        self._virtual_display = d

    @synchronized(_methods_lock)
    def stop_display(self):
        '''
        Convenient function to stop the virtual display
        '''
        # Nice alias
        display = self._virtual_display
        if ((not self.global_settings.get('virtual_display_keep_open')
             or not self.global_settings.get('virtual_display_visible'))
                and display):
            self.log.d('Stopping virtual display %r' % display)
            display.stop()
            WebdriverManager._virtual_display = None

    @synchronized(_methods_lock)
    def list_webdrivers(self, which='all'):
        '''
        Make a report of all, released or locked webdrivers.
        :param which: string specifiying the type of report. May be:'all', 'released', 'locked'
        '''
        wdrivers_report = {}
        possible = ('all', 'released', 'locked')
        assert which in possible, 'Which must be one of %r' % possible
        # Choose the set to report
        if which =='all':
            in_report = self._wdriver_pool
        elif which == 'released':
            in_report = self._released
        elif which == 'locked':
            in_report = self._locked
        # Build report
        for wdriver in in_report:
            browser, level = self._wdriver_pool[wdriver]
            wdrivers_report[wdriver] = browser, level
        return wdrivers_report

    def get_browser_name(self, browser_name=None):
        '''
        Get the current browser name in usage
        If not set, we use the one specified in global settings
        or PhantomJS as the final default
        '''
        browser = (browser_name
                   or self.global_settings.get('webdriver_browser'))
        return self.expand_browser_name(browser)

    def expand_browser_name(self, browser=None):
        '''
        Expand a partial or lower case Browser's name into its full name.
        Eg:
            firefo -> Firefox
            chrom -> Chrome
        :param browser: partial browser's name string
        '''
        # TODO: add IE and Opera
        char_browser = dict(f='Firefox',
                            c='Chrome',
                            p='PhantomJS',
                            )
        char = browser.lower()[0]
        assert char in char_browser, 'Could not find browser %r' % browser
        return char_browser.get(char)

    def __del__(self):
        self.exit_level(MANAGER_LIFE)
        self.stop_display()

class BrowserContextManager(XpathWdBase):
    def __init__(self, parent, level, base_url=None, name='', browser_name=None):
        self.parent = parent
        self.level = level
        self.base_url = base_url
        self.name = name
        self.webdriver = None
        self.browser_name = browser_name

    def __enter__(self):
        return self.get_xpathbrowser()

    def __exit__(self, type=None, value=None, traceback=None):
        self.release_driver()
        self.parent.exit_level(self.level)

    def acquire_driver(self):
        if not self.webdriver:
            self.webdriver = self.parent.acquire_driver(self.level, self.browser_name)
        return self.webdriver

    def get_xpathbrowser(self, base_url=None, name=''):
        from .logger import Logger
        from .xpath_browser import XpathBrowser
        base_url = base_url or self.base_url or self.global_settings.get('base_url')
        name = name or self.name
        # Initialize the XpathBrowser class
        return XpathBrowser(self.acquire_driver(), base_url,
                            Logger(name), settings=self.global_settings)

    def release_driver(self):
        if self.webdriver:
            self.parent.release_driver(self.webdriver, self.level)
            self.webdriver = None

def smoke_test_module():
    from .logger import log_test
    mngr = WebdriverManager()
#    import ipdb; ipdb.set_trace()
    lvl = mngr.enter_level(level=TEST_ROUND_LIFE)
    wd = lvl.acquire_driver()
    log_test(wd)
    lvl.release_driver()
    log_test(mngr.list_webdrivers())
#    import ipdb; ipdb.set_trace()
    xb = lvl.get_xpathbrowser()
    log_test(xb)
    lvl.exit_level()
#    import ipdb; ipdb.set_trace()
    mngr.stop_display()
    mngr.quit_all_webdrivers()
    # Wait some time for browser to quit
    import time
    time.sleep(.5)

if __name__ == "__main__":
    smoke_test_module()
<|endoftext|>