<commit_before>from numba import ir, ir_utils, types, errors, sigutils
from numba.typing.typeof import typeof_impl
from .transforms import find_region_inout_vars

class WithContext(object):
    """A dummy object for use as contextmanager.
    This can be used as a contextmanager.
    """
    is_callable = False

    def __enter__(self):
        pass

    def __exit__(self, typ, val, tb):
        pass

    def mutate_with_body(self, func_ir, blocks, blk_start, blk_end,
                         body_blocks, dispatcher_factory, extra):
        """Mutate the *blocks* to implement this contextmanager.

        Parameters
        ----------
        func_ir : FunctionIR
        blocks : dict[ir.Block]
        blk_start, blk_end : int
            labels of the starting and ending block of the context-maanger.
        body_block: sequence[int]
            A sequence of int's representing labels of the with-body
        dispatcher_factory : callable
            A callable that takes a `FunctionIR` and returns a `Dispatcher`.
        """
        raise NotImplementedError

@typeof_impl.register(WithContext)
def typeof_contextmanager(val, c):
    return types.ContextManager(val)

def _get_var_parent(name):
    """Get parent of the variable given its name
    """
    # If not a temprary variable
    if not name.startswith('$'):
        # Return the base component of the name
        return name.split('.', )[0]

def _clear_blocks(blocks, to_clear):
    """Remove keys in *to_clear* from *blocks*.
    """
    for b in to_clear:
        del blocks[b]

class _ByPassContextType(WithContext):
    """A simple context-manager that tells the compiler to bypass the body
    of the with-block.
    """
    def mutate_with_body(self, func_ir, blocks, blk_start, blk_end,
                         body_blocks, dispatcher_factory, extra):
        assert extra is None
        # Determine variables that need forwarding
        vlt = func_ir.variable_lifetime
        inmap = {_get_var_parent(k): k for k in vlt.livemap[blk_start]}
        outmap = {_get_var_parent(k): k for k in vlt.livemap[blk_end]}
        forwardvars = {inmap[k]: outmap[k] for k in filter(bool, outmap)}
        # Transform the block
        _bypass_with_context(blocks, blk_start, blk_end, forwardvars)
        _clear_blocks(blocks, body_blocks)

bypass_context = _ByPassContextType()

class _CallContextType(WithContext):
    """A simple context-manager that tells the compiler to lift the body of the
    with-block as another function.
    """
    def mutate_with_body(self, func_ir, blocks, blk_start, blk_end,
                         body_blocks, dispatcher_factory, extra):
        assert extra is None
        vlt = func_ir.variable_lifetime

        inputs, outputs = find_region_inout_vars(
            blocks=blocks,
            livemap=vlt.livemap,
            callfrom=blk_start,
            returnto=blk_end,
            body_block_ids=set(body_blocks),
            )

        lifted_blks = {k: blocks[k] for k in body_blocks}
        _mutate_with_block_callee(lifted_blks, blk_start, blk_end,
                                  inputs, outputs)

        # XXX: transform body-blocks to return the output variables
        lifted_ir = func_ir.derive(
            blocks=lifted_blks,
            arg_names=tuple(inputs),
            arg_count=len(inputs),
            force_non_generator=True,
            )

        dispatcher = dispatcher_factory(lifted_ir)

        newblk = _mutate_with_block_caller(
            dispatcher, blocks, blk_start, blk_end, inputs, outputs,
            )

        blocks[blk_start] = newblk
        _clear_blocks(blocks, body_blocks)
        return dispatcher

call_context = _CallContextType()

class _ObjModeContextType(WithContext):
    """Creates a contextmanager to be used inside jitted functions to enter
    *object-mode* for using interpreter features.  The body of the with-context
    is lifted into a function that is compiled into *object-mode*.  This
    transformation process is limitated and cannot process all possible
    Python code.  However, users can wrap complicated logic in another
    Python function, which will then be executed by the interpreter.

    Use this as a function that takes keyword arguments only.
    The argument names must correspond to the output variables from the
    with-block.  Their respective values are strings representing the expected
    types.  When exiting the with-context, the output variables are casted
    to the expected nopython types according to the annotation.  This process
    is the same as passing Python objects into arguments of a nopython
    function.

    Example::

        import numpy as np
        from numba import njit, objmode

        def bar(x):
            # This code is executed by the interpreter.
            return np.asarray(list(reversed(x.tolist()))

        @njit
        def foo():
            x = np.arange(5)
            with objmode(y='intp[:]'):  # annotate return type
                # this region is executed by object-mode.
                y += bar(x)
            return y

    .. note:: Known limitation:

        - with-block cannot use incoming list object.
        - with-block cannot use incoming function object.
        - with-block cannot ``yield``, ``break``, ``return`` or ``raise`` \
          such that the execution will leave the with-block immediately.
        - with-block cannot contain `with` statements.
        - random number generator states do not synchronize; i.e. \
          nopython-mode and object-mode uses different RNG states.

    .. note:: When used outside of no-python mode, the context-manager has no
        effect.

    .. warning:: This feature is experimental.  The supported features may
        change with or without notice.

    """
    is_callable = True

    def _legalize_args(self, extra, loc):
        """
        Legalize arguments to the context-manager
        """
        if extra is None:
            return {}

        if len(extra['args']) != 0:
            raise errors.CompilerError(
                "objectmode context doesn't take any positional arguments",
                )
        callkwargs = extra['kwargs']
        typeanns = {}
        for k, v in callkwargs.items():
            if not isinstance(v, ir.Const) or not isinstance(v.value, str):
                raise errors.CompilerError(
                    "objectmode context requires constants string for "
                    "type annotation",
                )

            typeanns[k] = sigutils._parse_signature_string(v.value)

        return typeanns

    def mutate_with_body(self, func_ir, blocks, blk_start, blk_end,
                         body_blocks, dispatcher_factory, extra):
        typeanns = self._legalize_args(extra, loc=blocks[blk_start].loc)
        vlt = func_ir.variable_lifetime

        inputs, outputs = find_region_inout_vars(
            blocks=blocks,
            livemap=vlt.livemap,
            callfrom=blk_start,
            returnto=blk_end,
            body_block_ids=set(body_blocks),
            )

        # Determine types in the output tuple
        def strip_var_ver(x):
            return x.split('.', 1)[0]

        stripped_outs = list(map(strip_var_ver, outputs))

        # Verify that only outputs are annotated
        extra_annotated = set(typeanns) - set(stripped_outs)
        if extra_annotated:
            msg = (
                'Invalid type annotation on non-outgoing variables: {}.'
                'Suggestion: remove annotation of the listed variables'
            )
            raise errors.TypingError(msg.format(extra_annotated))

        # Verify that all outputs are annotated
        not_annotated = set(stripped_outs) - set(typeanns)
        if not_annotated:
            msg = 'missing type annotation on outgoing variables: {}'
            raise errors.TypingError(msg.format(not_annotated))

        # Get output types
        outtup = types.Tuple([typeanns[v] for v in stripped_outs])

        lifted_blks = {k: blocks[k] for k in body_blocks}
        _mutate_with_block_callee(lifted_blks, blk_start, blk_end,
                                  inputs, outputs)

        lifted_ir = func_ir.derive(
            blocks=lifted_blks,
            arg_names=tuple(inputs),
            arg_count=len(inputs),
            force_non_generator=True,
            )

        dispatcher = dispatcher_factory(lifted_ir, objectmode=True,
                                        output_types=outtup)

        newblk = _mutate_with_block_caller(
            dispatcher, blocks, blk_start, blk_end, inputs, outputs,
            )

        blocks[blk_start] = newblk
        _clear_blocks(blocks, body_blocks)
        return dispatcher

    def __call__(self, *args, **kwargs):
        # No effect when used in pure-python
        return self

objmode_context = _ObjModeContextType()

def _bypass_with_context(blocks, blk_start, blk_end, forwardvars):
    """Given the starting and ending block of the with-context,
    replaces the head block with a new block that jumps to the end.

    *blocks* is modified inplace.
    """
    sblk = blocks[blk_start]
    scope = sblk.scope
    loc = sblk.loc
    newblk = ir.Block(scope=scope, loc=loc)
    for k, v in forwardvars.items():
        newblk.append(ir.Assign(value=scope.get_exact(k),
                                target=scope.get_exact(v),
                                loc=loc))
    newblk.append(ir.Jump(target=blk_end, loc=loc))
    blocks[blk_start] = newblk

def _mutate_with_block_caller(dispatcher, blocks, blk_start, blk_end,
                              inputs, outputs):
    """Make a new block that calls into the lifeted with-context.

    Parameters
    ----------
    dispatcher : Dispatcher
    blocks : dict[ir.Block]
    blk_start, blk_end : int
        labels of the starting and ending block of the context-manager.
    inputs: sequence[str]
        Input variable names
    outputs: sequence[str]
        Output variable names
    """
    sblk = blocks[blk_start]
    scope = sblk.scope
    loc = sblk.loc
    newblock = ir.Block(scope=scope, loc=loc)

    ir_utils.fill_block_with_call(
        newblock=newblock,
        callee=dispatcher,
        label_next=blk_end,
        inputs=inputs,
        outputs=outputs,
        )
    return newblock

def _mutate_with_block_callee(blocks, blk_start, blk_end, inputs, outputs):
    """Mutate *blocks* for the callee of a with-context.

    Parameters
    ----------
    blocks : dict[ir.Block]
    blk_start, blk_end : int
        labels of the starting and ending block of the context-manager.
    inputs: sequence[str]
        Input variable names
    outputs: sequence[str]
        Output variable names
    """
    head_blk = min(blocks)
    temp_blk = blocks[head_blk]
    scope = temp_blk.scope
    loc = temp_blk.loc

    blocks[blk_start] = ir_utils.fill_callee_prologue(
        block=ir.Block(scope=scope, loc=loc),
        inputs=inputs,
        label_next=head_blk,
        )
    blocks[blk_end] = ir_utils.fill_callee_epilogue(
        block=ir.Block(scope=scope, loc=loc),
        outputs=outputs,
    )

<commit_msg>Fix example and docs typos for `objmode` context manager.<commit_after>from numba import ir, ir_utils, types, errors, sigutils
from numba.typing.typeof import typeof_impl
from .transforms import find_region_inout_vars

class WithContext(object):
    """A dummy object for use as contextmanager.
    This can be used as a contextmanager.
    """
    is_callable = False

    def __enter__(self):
        pass

    def __exit__(self, typ, val, tb):
        pass

    def mutate_with_body(self, func_ir, blocks, blk_start, blk_end,
                         body_blocks, dispatcher_factory, extra):
        """Mutate the *blocks* to implement this contextmanager.

        Parameters
        ----------
        func_ir : FunctionIR
        blocks : dict[ir.Block]
        blk_start, blk_end : int
            labels of the starting and ending block of the context-maanger.
        body_block: sequence[int]
            A sequence of int's representing labels of the with-body
        dispatcher_factory : callable
            A callable that takes a `FunctionIR` and returns a `Dispatcher`.
        """
        raise NotImplementedError

@typeof_impl.register(WithContext)
def typeof_contextmanager(val, c):
    return types.ContextManager(val)

def _get_var_parent(name):
    """Get parent of the variable given its name
    """
    # If not a temprary variable
    if not name.startswith('$'):
        # Return the base component of the name
        return name.split('.', )[0]

def _clear_blocks(blocks, to_clear):
    """Remove keys in *to_clear* from *blocks*.
    """
    for b in to_clear:
        del blocks[b]

class _ByPassContextType(WithContext):
    """A simple context-manager that tells the compiler to bypass the body
    of the with-block.
    """
    def mutate_with_body(self, func_ir, blocks, blk_start, blk_end,
                         body_blocks, dispatcher_factory, extra):
        assert extra is None
        # Determine variables that need forwarding
        vlt = func_ir.variable_lifetime
        inmap = {_get_var_parent(k): k for k in vlt.livemap[blk_start]}
        outmap = {_get_var_parent(k): k for k in vlt.livemap[blk_end]}
        forwardvars = {inmap[k]: outmap[k] for k in filter(bool, outmap)}
        # Transform the block
        _bypass_with_context(blocks, blk_start, blk_end, forwardvars)
        _clear_blocks(blocks, body_blocks)

bypass_context = _ByPassContextType()

class _CallContextType(WithContext):
    """A simple context-manager that tells the compiler to lift the body of the
    with-block as another function.
    """
    def mutate_with_body(self, func_ir, blocks, blk_start, blk_end,
                         body_blocks, dispatcher_factory, extra):
        assert extra is None
        vlt = func_ir.variable_lifetime

        inputs, outputs = find_region_inout_vars(
            blocks=blocks,
            livemap=vlt.livemap,
            callfrom=blk_start,
            returnto=blk_end,
            body_block_ids=set(body_blocks),
            )

        lifted_blks = {k: blocks[k] for k in body_blocks}
        _mutate_with_block_callee(lifted_blks, blk_start, blk_end,
                                  inputs, outputs)

        # XXX: transform body-blocks to return the output variables
        lifted_ir = func_ir.derive(
            blocks=lifted_blks,
            arg_names=tuple(inputs),
            arg_count=len(inputs),
            force_non_generator=True,
            )

        dispatcher = dispatcher_factory(lifted_ir)

        newblk = _mutate_with_block_caller(
            dispatcher, blocks, blk_start, blk_end, inputs, outputs,
            )

        blocks[blk_start] = newblk
        _clear_blocks(blocks, body_blocks)
        return dispatcher

call_context = _CallContextType()

class _ObjModeContextType(WithContext):
    """Creates a contextmanager to be used inside jitted functions to enter
    *object-mode* for using interpreter features.  The body of the with-context
    is lifted into a function that is compiled in *object-mode*.  This
    transformation process is limited and cannot process all possible
    Python code.  However, users can wrap complicated logic in another
    Python function, which will then be executed by the interpreter.

    Use this as a function that takes keyword arguments only.
    The argument names must correspond to the output variables from the
    with-block.  Their respective values are strings representing the expected
    types.  When exiting the with-context, the output variables are cast
    to the expected nopython types according to the annotation.  This process
    is the same as passing Python objects into arguments of a nopython
    function.

    Example::

        import numpy as np
        from numba import njit, objmode

        def bar(x):
            # This code is executed by the interpreter.
            return np.asarray(list(reversed(x.tolist())))

        @njit
        def foo():
            y = np.arange(5)
            with objmode(y='intp[:]'):  # annotate return type
                # this region is executed by object-mode.
                y += bar(y)
            return y

    .. note:: Known limitations:

        - with-block cannot use incoming list objects.
        - with-block cannot use incoming function objects.
        - with-block cannot ``yield``, ``break``, ``return`` or ``raise`` \
          such that the execution will leave the with-block immediately.
        - with-block cannot contain `with` statements.
        - random number generator states do not synchronize; i.e. \
          nopython-mode and object-mode uses different RNG states.

    .. note:: When used outside of no-python mode, the context-manager has no
        effect.

    .. warning:: This feature is experimental.  The supported features may
        change with or without notice.

    """
    is_callable = True

    def _legalize_args(self, extra, loc):
        """
        Legalize arguments to the context-manager
        """
        if extra is None:
            return {}

        if len(extra['args']) != 0:
            raise errors.CompilerError(
                "objectmode context doesn't take any positional arguments",
                )
        callkwargs = extra['kwargs']
        typeanns = {}
        for k, v in callkwargs.items():
            if not isinstance(v, ir.Const) or not isinstance(v.value, str):
                raise errors.CompilerError(
                    "objectmode context requires constants string for "
                    "type annotation",
                )

            typeanns[k] = sigutils._parse_signature_string(v.value)

        return typeanns

    def mutate_with_body(self, func_ir, blocks, blk_start, blk_end,
                         body_blocks, dispatcher_factory, extra):
        typeanns = self._legalize_args(extra, loc=blocks[blk_start].loc)
        vlt = func_ir.variable_lifetime

        inputs, outputs = find_region_inout_vars(
            blocks=blocks,
            livemap=vlt.livemap,
            callfrom=blk_start,
            returnto=blk_end,
            body_block_ids=set(body_blocks),
            )

        # Determine types in the output tuple
        def strip_var_ver(x):
            return x.split('.', 1)[0]

        stripped_outs = list(map(strip_var_ver, outputs))

        # Verify that only outputs are annotated
        extra_annotated = set(typeanns) - set(stripped_outs)
        if extra_annotated:
            msg = (
                'Invalid type annotation on non-outgoing variables: {}.'
                'Suggestion: remove annotation of the listed variables'
            )
            raise errors.TypingError(msg.format(extra_annotated))

        # Verify that all outputs are annotated
        not_annotated = set(stripped_outs) - set(typeanns)
        if not_annotated:
            msg = 'missing type annotation on outgoing variables: {}'
            raise errors.TypingError(msg.format(not_annotated))

        # Get output types
        outtup = types.Tuple([typeanns[v] for v in stripped_outs])

        lifted_blks = {k: blocks[k] for k in body_blocks}
        _mutate_with_block_callee(lifted_blks, blk_start, blk_end,
                                  inputs, outputs)

        lifted_ir = func_ir.derive(
            blocks=lifted_blks,
            arg_names=tuple(inputs),
            arg_count=len(inputs),
            force_non_generator=True,
            )

        dispatcher = dispatcher_factory(lifted_ir, objectmode=True,
                                        output_types=outtup)

        newblk = _mutate_with_block_caller(
            dispatcher, blocks, blk_start, blk_end, inputs, outputs,
            )

        blocks[blk_start] = newblk
        _clear_blocks(blocks, body_blocks)
        return dispatcher

    def __call__(self, *args, **kwargs):
        # No effect when used in pure-python
        return self

objmode_context = _ObjModeContextType()

def _bypass_with_context(blocks, blk_start, blk_end, forwardvars):
    """Given the starting and ending block of the with-context,
    replaces the head block with a new block that jumps to the end.

    *blocks* is modified inplace.
    """
    sblk = blocks[blk_start]
    scope = sblk.scope
    loc = sblk.loc
    newblk = ir.Block(scope=scope, loc=loc)
    for k, v in forwardvars.items():
        newblk.append(ir.Assign(value=scope.get_exact(k),
                                target=scope.get_exact(v),
                                loc=loc))
    newblk.append(ir.Jump(target=blk_end, loc=loc))
    blocks[blk_start] = newblk

def _mutate_with_block_caller(dispatcher, blocks, blk_start, blk_end,
                              inputs, outputs):
    """Make a new block that calls into the lifeted with-context.

    Parameters
    ----------
    dispatcher : Dispatcher
    blocks : dict[ir.Block]
    blk_start, blk_end : int
        labels of the starting and ending block of the context-manager.
    inputs: sequence[str]
        Input variable names
    outputs: sequence[str]
        Output variable names
    """
    sblk = blocks[blk_start]
    scope = sblk.scope
    loc = sblk.loc
    newblock = ir.Block(scope=scope, loc=loc)

    ir_utils.fill_block_with_call(
        newblock=newblock,
        callee=dispatcher,
        label_next=blk_end,
        inputs=inputs,
        outputs=outputs,
        )
    return newblock

def _mutate_with_block_callee(blocks, blk_start, blk_end, inputs, outputs):
    """Mutate *blocks* for the callee of a with-context.

    Parameters
    ----------
    blocks : dict[ir.Block]
    blk_start, blk_end : int
        labels of the starting and ending block of the context-manager.
    inputs: sequence[str]
        Input variable names
    outputs: sequence[str]
        Output variable names
    """
    head_blk = min(blocks)
    temp_blk = blocks[head_blk]
    scope = temp_blk.scope
    loc = temp_blk.loc

    blocks[blk_start] = ir_utils.fill_callee_prologue(
        block=ir.Block(scope=scope, loc=loc),
        inputs=inputs,
        label_next=head_blk,
        )
    blocks[blk_end] = ir_utils.fill_callee_epilogue(
        block=ir.Block(scope=scope, loc=loc),
        outputs=outputs,
    )

<|endoftext|>


<commit_before>import functools
from math import ceil
from flask import jsonify
from flask_restful import reqparse

def parse_request(*args, **kwargs):
    """
    Decorator used to parse request arguments.
    This method allows to parse args in different locations then the json body
    (query string, files, form, headers, etc)
    :param args: list of Arguments (flask_restful.reqparse.Argument)
    :return: decorator function
    """
    parser = reqparse.RequestParser(bundle_errors=True)
    for arg in args:
        parser.add_argument(arg)

    if kwargs.get('allow_ordering', None):
        parser_args = [p.name for p in parser.args]
        default_args = [
            reqparse.Argument("sort_by", type=str, store_missing=False),
            reqparse.Argument("order", type=str, store_missing=False, choices=('asc', 'desc')),
        ]
        for p in default_args:
            if p.name not in parser_args:
                parser.add_argument(p)

    def decorator(f):
        @functools.wraps(f)
        def inner(*fargs, **fkwargs):
            fkwargs.update(parser.parse_args())
            return f(*fargs, **fkwargs)

        return inner

    return decorator

def pagination_args_parser(default_page_size=100):
    parser = reqparse.RequestParser()
    parser.add_argument('page', type=int, default=1,
                        help='Page number must be a positive integer.')
    parser.add_argument('count', type=int, default=default_page_size,
                        required=False,
                        help='Items per page must be a positive integer.')
    parser.add_argument('pagination', type=bool, default=True,
                        required=False,
                        help='If "false", return all results in a single page.')
    return parser.parse_args

def paginate_deckbox_results(results, serializer, args_parser=None):
    return get_paginated_response(
        results['cards'],
        serializer,
        results['page'],
        results['count'],
        results['total'],
        results['total_pages']
    )

def paginate_list(results, serializer, args_parser=pagination_args_parser()):
    """Very simple pagination helper for lists.
    This logic could be more elaborated and moved into a class.
    """
    args = args_parser()

    total_items = len(results)

    if not args.pagination:
        return get_paginated_response(results, serializer, 1, total_items, total_items, 1)

    total_pages, page, offset = get_pagination_params(total_items, args.page, args.count)

    return get_paginated_response(
        results[offset:offset + args.count],
        serializer,
        page,
        args.count,
        total_items,
        total_pages
    )

def marshal_with(schema, pagination=False, paginator=paginate_list, args_parser=pagination_args_parser(),
                 success_code=200, **kwargs):
    """Decorator to serialize output using specified schema
    :param kwargs will be passed down to the dump method from marshmallow Schema
    """
    serializer = functools.partial(schema.dump, **kwargs)

    def decorator(f):
        @functools.wraps(f)
        def inner(*fargs, **fkwargs):
            rv = f(*fargs, **fkwargs)
            if pagination:
                if isinstance(rv, list):
                    rv = paginator(rv, serializer, args_parser)
                else:
                    rv = paginator(rv, serializer, args_parser)
            else:
                rv = serializer(rv)
            response = jsonify(rv)
            response.status_code = success_code
            return response

        return inner

    return decorator

def get_pagination_params(total_items, page, count):
    total_pages = int(ceil(total_items / float(count)))
    if page > total_pages:
        page = total_pages

    offset = (page - 1) * count

    if offset < 0:
        offset = 0

    return total_pages, page, offset

def get_paginated_response(items, serializer, page, count, total_items, total_pages):
    return {
        'count': total_items,
        'page': page,
        'items_per_page': count,
        'total_pages': total_pages,
        'items': serializer(items).data
    }
<commit_msg>#10: fixed bug introduced when upgraded marshmallow<commit_after>import functools
from math import ceil
from flask import jsonify
from flask_restful import reqparse

def parse_request(*args, **kwargs):
    """
    Decorator used to parse request arguments.
    This method allows to parse args in different locations then the json body
    (query string, files, form, headers, etc)
    :param args: list of Arguments (flask_restful.reqparse.Argument)
    :return: decorator function
    """
    parser = reqparse.RequestParser(bundle_errors=True)
    for arg in args:
        parser.add_argument(arg)

    if kwargs.get('allow_ordering', None):
        parser_args = [p.name for p in parser.args]
        default_args = [
            reqparse.Argument("sort_by", type=str, store_missing=False),
            reqparse.Argument("order", type=str, store_missing=False, choices=('asc', 'desc')),
        ]
        for p in default_args:
            if p.name not in parser_args:
                parser.add_argument(p)

    def decorator(f):
        @functools.wraps(f)
        def inner(*fargs, **fkwargs):
            fkwargs.update(parser.parse_args())
            return f(*fargs, **fkwargs)

        return inner

    return decorator

def pagination_args_parser(default_page_size=100):
    parser = reqparse.RequestParser()
    parser.add_argument('page', type=int, default=1,
                        help='Page number must be a positive integer.')
    parser.add_argument('count', type=int, default=default_page_size,
                        required=False,
                        help='Items per page must be a positive integer.')
    parser.add_argument('pagination', type=bool, default=True,
                        required=False,
                        help='If "false", return all results in a single page.')
    return parser.parse_args

def paginate_deckbox_results(results, serializer, args_parser=None):
    return get_paginated_response(
        results['cards'],
        serializer,
        results['page'],
        results['count'],
        results['total'],
        results['total_pages']
    )

def paginate_list(results, serializer, args_parser=pagination_args_parser()):
    """Very simple pagination helper for lists.
    This logic could be more elaborated and moved into a class.
    """
    args = args_parser()

    total_items = len(results)

    if not args.pagination:
        return get_paginated_response(results, serializer, 1, total_items, total_items, 1)

    total_pages, page, offset = get_pagination_params(total_items, args.page, args.count)

    return get_paginated_response(
        results[offset:offset + args.count],
        serializer,
        page,
        args.count,
        total_items,
        total_pages
    )

def marshal_with(schema, pagination=False, paginator=paginate_list, args_parser=pagination_args_parser(),
                 success_code=200, **kwargs):
    """Decorator to serialize output using specified schema
    :param kwargs will be passed down to the dump method from marshmallow Schema
    """
    serializer = functools.partial(schema.dump, **kwargs)

    def decorator(f):
        @functools.wraps(f)
        def inner(*fargs, **fkwargs):
            rv = f(*fargs, **fkwargs)
            if pagination:
                if isinstance(rv, list):
                    rv = paginator(rv, serializer, args_parser)
                else:
                    rv = paginator(rv, serializer, args_parser)
            else:
                rv = serializer(rv)
            response = jsonify(rv)
            response.status_code = success_code
            return response

        return inner

    return decorator

def get_pagination_params(total_items, page, count):
    total_pages = int(ceil(total_items / float(count)))
    if page > total_pages:
        page = total_pages

    offset = (page - 1) * count

    if offset < 0:
        offset = 0

    return total_pages, page, offset

def get_paginated_response(items, serializer, page, count, total_items, total_pages):
    return {
        'count': total_items,
        'page': page,
        'items_per_page': count,
        'total_pages': total_pages,
        'items': serializer(items)
    }
<|endoftext|>


<commit_before># This file is the actual code for the custom Python dataset looker-query_run-query

from dataiku.connector import Connector

import logging
import pprint

import json

from looker_sdk import client, models, error

class MyConnector(Connector):

    def __init__(self, config, plugin_config):
        logging.basicConfig(level=logging.INFO, format='dss-plugin-looker-query %(levelname)s - %(message)s')
        self.logger = logging.getLogger()

        Connector.__init__(self, config, plugin_config)  # pass the parameters to the base class

        self.base_url = self.plugin_config["Looker API"]["baseurl"]
        self.client_id = self.plugin_config["Looker API"]["clientid"]
        self.client_secret = self.plugin_config["Looker API"]["clientsecret"]
        self.look_id = int(self.config["lookid"])
        #self.limit = int(self.config["limit"])

        print(self.plugin_config)

        if not (self.base_url and self.client_id and self.client_secret):
            self.logger.error('Connection params: {}'.format(
                {'Client ID:' : self.client_id,
                'Client secret:' : '#' * len(self.client_secret),
                'Base URL:' : self.base_url})
            )
            raise ValueError("Client ID, Client secret and Base URL must be filled")

        if not self.look_id:
            raise ValueError("Look ID was not set")

        #if self.limit is None or self.limit = 0:
        #    self.limit = 500

        file = open("looker.ini", "w")
        line_list = ["[Looker]", "api_version=3.1", "base_url=" + self.base_url, "client_id=" + self.client_id, "client_secret=" + self.client_secret, "verify_ssl=True"]
        file.write('\n'.join(line_list) + '\n')
        file.close()

        self.looker_client = client.setup("looker.ini")
        print(self.looker_client)

    def get_read_schema(self):
        return None

    def generate_rows(self, dataset_schema=None, dataset_partitioning=None,
                            partition_id=None, records_limit = -1):

        response = self.looker_client.run_look(self.look_id, "json", limit=records_limit)
        data = json.loads(response)

        for row in data:
            yield row

    def get_writer(self, dataset_schema=None, dataset_partitioning=None,
                         partition_id=None):
        raise NotImplementedError("Writer")

    def get_partitioning(self):
        raise NotImplementedError("Partition list")

    def list_partitions(self, partitioning):
        return []

    def partition_exists(self, partitioning, partition_id):
        raise NotImplementedError("Partition check")

    def get_records_count(self, partitioning=None, partition_id=None):
        raise NotImplementedError("Count records")
<commit_msg>Forget about limiting record, it's broken<commit_after># This file is the actual code for the custom Python dataset looker-query_run-query

from dataiku.connector import Connector

import logging
import pprint

import json

from looker_sdk import client, models, error

class MyConnector(Connector):

    def __init__(self, config, plugin_config):
        logging.basicConfig(level=logging.INFO, format='dss-plugin-looker-query %(levelname)s - %(message)s')
        self.logger = logging.getLogger()

        Connector.__init__(self, config, plugin_config)  # pass the parameters to the base class

        self.base_url = self.plugin_config["Looker API"]["baseurl"]
        self.client_id = self.plugin_config["Looker API"]["clientid"]
        self.client_secret = self.plugin_config["Looker API"]["clientsecret"]
        self.look_id = int(self.config["lookid"])
        #self.limit = int(self.config["limit"])

        print(self.plugin_config)

        if not (self.base_url and self.client_id and self.client_secret):
            self.logger.error('Connection params: {}'.format(
                {'Client ID:' : self.client_id,
                'Client secret:' : '#' * len(self.client_secret),
                'Base URL:' : self.base_url})
            )
            raise ValueError("Client ID, Client secret and Base URL must be filled")

        if not self.look_id:
            raise ValueError("Look ID was not set")

        #if self.limit is None or self.limit = 0:
        #    self.limit = 500

        file = open("looker.ini", "w")
        line_list = ["[Looker]", "api_version=3.1", "base_url=" + self.base_url, "client_id=" + self.client_id, "client_secret=" + self.client_secret, "verify_ssl=True"]
        file.write('\n'.join(line_list) + '\n')
        file.close()

        self.looker_client = client.setup("looker.ini")
        print(self.looker_client)

    def get_read_schema(self):
        return None

    def generate_rows(self, dataset_schema=None, dataset_partitioning=None,
                            partition_id=None, records_limit = -1):

        response = self.looker_client.run_look(self.look_id, "json")
        data = json.loads(response)

        for row in data:
            yield row

    def get_writer(self, dataset_schema=None, dataset_partitioning=None,
                         partition_id=None):
        raise NotImplementedError("Writer")

    def get_partitioning(self):
        raise NotImplementedError("Partition list")

    def list_partitions(self, partitioning):
        return []

    def partition_exists(self, partitioning, partition_id):
        raise NotImplementedError("Partition check")

    def get_records_count(self, partitioning=None, partition_id=None):
        raise NotImplementedError("Count records")
<|endoftext|>


<commit_before>from django.contrib.auth.models import AnonymousUser
from core.models.group import get_user_group
from core.models.project import Project
from rest_framework import serializers

class ProjectsField(serializers.Field):
    def to_representation(self, project_mgr):
        request_user = self.root.request_user
        if type(request_user) == AnonymousUser:
            return None
        try:
            group = get_user_group(request_user.username)
            projects = project_mgr.filter(owner=group)
            # Modifications to how 'project' should be displayed here:
            return [p.uuid for p in projects]
        except Project.DoesNotExist:
            return None

    def to_internal_value(self, data, files, field_name, into):
        value = data.get(field_name)
        if value is None:
            return
        related_obj = self.root.object
        user = self.root.request_user
        group = get_user_group(user.username)
        # Retrieve the New Project(s)
        if type(value) == list:
            new_projects = value
        else:
            new_projects = [value, ]

        # Remove related_obj from Old Project(s)
        old_projects = related_obj.get_projects(user)
        for old_proj in old_projects:
            related_obj.projects.remove(old_proj)

        # Add Project(s) to related_obj
        for project_id in new_projects:
            # Retrieve/Create the New Project
            #TODO: When projects can be shared,
            #change the qualifier here.
            new_project = Project.objects.get(id=project_id, owner=group)
            # Assign related_obj to New Project
            if not related_obj.projects.filter(id=project_id):
                related_obj.projects.add(new_project)
        # Modifications to how 'project' should be displayed here:
        into[field_name] = new_projects<commit_msg>update naughty DRF internal references to DRF3 api<commit_after>from django.contrib.auth.models import AnonymousUser
from core.models.group import get_user_group
from core.models.project import Project
from rest_framework import serializers

class ProjectsField(serializers.Field):
    def to_representation(self, project_mgr):
        request_user = self.parent.request_user
        if type(request_user) == AnonymousUser:
            return None
        try:
            group = get_user_group(request_user.username)
            projects = project_mgr.filter(owner=group)
            # Modifications to how 'project' should be displayed here:
            return [p.uuid for p in projects]
        except Project.DoesNotExist:
            return None

    def to_internal_value(self, data, files, field_name, into):
        value = data.get(field_name)
        if value is None:
            return
        related_obj = self.parent.instance
        user = self.parent.request_user
        group = get_user_group(user.username)
        # Retrieve the New Project(s)
        if type(value) == list:
            new_projects = value
        else:
            new_projects = [value, ]

        # Remove related_obj from Old Project(s)
        old_projects = related_obj.get_projects(user)
        for old_proj in old_projects:
            related_obj.projects.remove(old_proj)

        # Add Project(s) to related_obj
        for project_id in new_projects:
            # Retrieve/Create the New Project
            #TODO: When projects can be shared,
            #change the qualifier here.
            new_project = Project.objects.get(id=project_id, owner=group)
            # Assign related_obj to New Project
            if not related_obj.projects.filter(id=project_id):
                related_obj.projects.add(new_project)
        # Modifications to how 'project' should be displayed here:
        into[field_name] = new_projects<|endoftext|>


<commit_before>"""
module __GenericGeometry__: Base class to construct pore networks
==================================================================

.. warning:: The classes of this module should be loaded through the 'Geometry.__init__.py' file.

"""

import sys, os
parent_dir = os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
if sys.path[1] != parent_dir:
    sys.path.insert(1, parent_dir)
import OpenPNM
import scipy as sp
from functools import partial

class GenericGeometry(OpenPNM.Utilities.Base):
    r"""
    GenericGeometry - Base class to construct pore networks

    This class contains the interface definition for the construction of networks

    Parameters
    ----------
    network : OpenPNM Network Object

    name : string
        A unique name to apply to the object.  This name will also be used as a
        label to identify where this this geometry applies.

    pnums and tnums : boolean mask or list of indices
        The pore (pnums) and throat (tnums) locations in the network where this
        geometry applies.  By default it will apply everywhere.  To create an
        empty geometry set pnums and tnums to empty lists [].

    loglevel : int
        Level of the logger (10=Debug, 20=Info, 30=Warning, 40=Error, 50=Critical)

    loggername : string (optional)
        Sets a custom name for the logger, to help identify logger messages

    Examples
    --------
    >>> pn = OpenPNM.Network.TestNet()
    >>> loc = pn.get_pore_indices() #Get all pores to define geometry everywhere
    >>> geo = OpenPNM.Geometry.GenericGeometry(name='geo_test',locations=loc,network=pn)
    >>> geo.add_method(prop='pore_seed',model='constant',value=0.123)
    >>> geo.regenerate()
    >>> seeds = pn.get_pore_data(locations='geo_test',prop='seed')
    >>> seeds[0]
    0.123
    """

    def __init__(self, network,name,pnums='all',tnums='all',**kwargs):
        r"""
        Initialize
        """
        super(GenericGeometry,self).__init__(**kwargs)
        self._logger.debug("Method: Constructor")
        for item in network._geometries.keys():
            if item == name:
                raise Exception('A Geometry Object with the supplied name already exists')
        network.set_pore_info(label=name,locations=pnums)
        network.set_throat_info(label=name,locations=tnums)
        network._geometries.update({name:self}) #attach this geometry to network
        self.name = name
        self._net = network #Attach network to self
        self._physics = {} #Create list for physics to append themselves to
        self._prop_list = []

    def regenerate(self, prop_list='',mode=None):
        r'''
        This updates all properties using the selected methods

        Parameters
        ----------
        prop_list : string or list of strings
            The names of the properties that should be updated, defaults to all
        mode : string
            Control how the regeneration occurs.

        Examples
        --------
        >>> pn = OpenPNM.Network.TestNet()
        >>> pind = pn.get_pore_indices()
        >>> geom = OpenPNM.Geometry.Stick_and_Ball(network=pn, name='geo_test', locations=pind)
        >>> geom.regenerate()  # Regenerate all properties at once
        >>> geom.regenerate('pore_seed')  # only one property
        >>> geom.regenerate(['pore_seed', 'pore_diameter'])  # or several
        '''
        if prop_list == '':
            prop_list = self._prop_list
        elif type(prop_list) == str:
            prop_list = [prop_list]
        if mode == 'exclude':
            a = sp.array(self._prop_list)
            b = sp.array(prop_list)
            c = a[sp.where(~sp.in1d(a,b))[0]]
            prop_list = list(c)
        for item in prop_list:
            self._logger.debug('Refreshing: '+item)
            getattr(self,item)()

    def add_method(self,prop='',prop_name='',**kwargs):
        r'''
        Add specified property estimation model to the fluid object.

        Parameters
        ----------
        prop : string
            The name of the fluid property attribute to add.
            This name must correspond with a file in the Fluids folder.
            To add a new property simply add a file with the appropriate name and the necessary methods.

        prop_name : string, optional
            This argument will be used as the method name and the dictionary key
            where data is written by method. This option is provided for occasions
            when multiple properties of the same type are required, such as
            diffusivity coefficients of each species in a multicomponent mixture.

        Examples
        --------
        >>> pn = OpenPNM.Network.TestNet()
        '''
        try:
            function = getattr( getattr(OpenPNM.Geometry, prop), kwargs['model'] ) # this gets the method from the file
            if prop_name: propname = prop = prop_name #overwrite the default prop with user supplied name
            else:
                #remove leading pore_ or throat_ from dictionary key
                propname = prop.split('_')[1]
                element = prop.split('_')[0]
                if len(prop.split('_')) > 2:
                    propname = prop.split(element+'_')[1]
            preloaded_fn = partial(function, geometry=self, network=self._net,propname=propname, **kwargs) #
            setattr(self, prop, preloaded_fn)
            self._logger.info("Successfully loaded {}.".format(prop))
            self._prop_list.append(prop)
        except AttributeError: print('could not find',kwargs['model'])

    def check_consistency(self):
        r'''
        Checks to see if the current geometry conflicts with any other geometry
        '''
        temp = sp.zeros_like(self._net.get_pore_info(label=self.name),dtype=int)
        for item in self._net._geometries:
            temp = temp + sp.array(self._net.get_pore_info(label=item),dtype=int)
        print('Geometry labels overlap in', sp.sum(temp>1),'pores')
        print('Geometry not yet applied to',sp.sum(temp==0),'pores')

        temp = sp.zeros_like(self._net.get_throat_info(label=self.name),dtype=int)
        for item in self._net._geometries:
            temp = temp + sp.array(self._net.get_throat_info(label=item),dtype=int)
        print('Geometry labels overlap in', sp.sum(temp>1),'throats')
        print('Geometry not yet applied to',sp.sum(temp==0),'throats')

if __name__ == '__main__':
    pn = OpenPNM.Network.TestNet()
    loc = pn.get_pore_indices()
    test = OpenPNM.Geometry.GenericGeometry(name='doc_test',locations=loc,network=pn)

<commit_msg>Added the get_pore/throats_indices methods from the GenericNetwork to the Geomety objects using Composition - This also used the partial function to 'lock' in the label as geometry's name attribute<commit_after>"""
module __GenericGeometry__: Base class to construct pore networks
==================================================================

.. warning:: The classes of this module should be loaded through the 'Geometry.__init__.py' file.

"""

import sys, os
parent_dir = os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
if sys.path[1] != parent_dir:
    sys.path.insert(1, parent_dir)
import OpenPNM
import scipy as sp
from functools import partial

class GenericGeometry(OpenPNM.Utilities.Base):
    r"""
    GenericGeometry - Base class to construct pore networks

    This class contains the interface definition for the construction of networks

    Parameters
    ----------
    network : OpenPNM Network Object

    name : string
        A unique name to apply to the object.  This name will also be used as a
        label to identify where this this geometry applies.

    pnums and tnums : boolean mask or list of indices
        The pore (pnums) and throat (tnums) locations in the network where this
        geometry applies.  By default it will apply everywhere.  To create an
        empty geometry set pnums and tnums to empty lists [].

    loglevel : int
        Level of the logger (10=Debug, 20=Info, 30=Warning, 40=Error, 50=Critical)

    loggername : string (optional)
        Sets a custom name for the logger, to help identify logger messages

    Examples
    --------
    >>> pn = OpenPNM.Network.TestNet()
    >>> loc = pn.get_pore_indices() #Get all pores to define geometry everywhere
    >>> geo = OpenPNM.Geometry.GenericGeometry(name='geo_test',locations=loc,network=pn)
    >>> geo.add_method(prop='pore_seed',model='constant',value=0.123)
    >>> geo.regenerate()
    >>> seeds = pn.get_pore_data(locations='geo_test',prop='seed')
    >>> seeds[0]
    0.123
    """

    def __init__(self, network,name,pnums='all',tnums='all',**kwargs):
        r"""
        Initialize
        """
        super(GenericGeometry,self).__init__(**kwargs)
        self._logger.debug("Method: Constructor")
        for item in network._geometries.keys():
            if item == name:
                raise Exception('A Geometry Object with the supplied name already exists')
        network.set_pore_info(label=name,locations=pnums)
        network.set_throat_info(label=name,locations=tnums)
        network._geometries.update({name:self}) #attach this geometry to network
        self.name = name
        self._net = network #Attach network to self
        self._physics = {} #Create list for physics to append themselves to
        self._prop_list = []
        # Add location query methods from GenericNetwork using partial to
        # force the label to be self.name
        preloaded_fn = partial(network.get_pore_indices, labels=name, mode='union')
        setattr(self, 'get_pore_indices', preloaded_fn)
        preloaded_fn = partial(network.get_throat_indices, labels=name, mode='union')
        setattr(self, 'get_throat_indices', preloaded_fn)

    def regenerate(self, prop_list='',mode=None):
        r'''
        This updates all properties using the selected methods

        Parameters
        ----------
        prop_list : string or list of strings
            The names of the properties that should be updated, defaults to all
        mode : string
            Control how the regeneration occurs.

        Examples
        --------
        >>> pn = OpenPNM.Network.TestNet()
        >>> pind = pn.get_pore_indices()
        >>> geom = OpenPNM.Geometry.Stick_and_Ball(network=pn, name='geo_test', locations=pind)
        >>> geom.regenerate()  # Regenerate all properties at once
        >>> geom.regenerate('pore_seed')  # only one property
        >>> geom.regenerate(['pore_seed', 'pore_diameter'])  # or several
        '''
        if prop_list == '':
            prop_list = self._prop_list
        elif type(prop_list) == str:
            prop_list = [prop_list]
        if mode == 'exclude':
            a = sp.array(self._prop_list)
            b = sp.array(prop_list)
            c = a[sp.where(~sp.in1d(a,b))[0]]
            prop_list = list(c)
        for item in prop_list:
            self._logger.debug('Refreshing: '+item)
            getattr(self,item)()

    def add_method(self,prop='',prop_name='',**kwargs):
        r'''
        Add specified property estimation model to the fluid object.

        Parameters
        ----------
        prop : string
            The name of the fluid property attribute to add.
            This name must correspond with a file in the Fluids folder.
            To add a new property simply add a file with the appropriate name and the necessary methods.

        prop_name : string, optional
            This argument will be used as the method name and the dictionary key
            where data is written by method. This option is provided for occasions
            when multiple properties of the same type are required, such as
            diffusivity coefficients of each species in a multicomponent mixture.

        Examples
        --------
        >>> pn = OpenPNM.Network.TestNet()
        '''
        try:
            function = getattr( getattr(OpenPNM.Geometry, prop), kwargs['model'] ) # this gets the method from the file
            if prop_name: propname = prop = prop_name #overwrite the default prop with user supplied name
            else:
                #remove leading pore_ or throat_ from dictionary key
                propname = prop.split('_')[1]
                element = prop.split('_')[0]
                if len(prop.split('_')) > 2:
                    propname = prop.split(element+'_')[1]
            preloaded_fn = partial(function, geometry=self, network=self._net,propname=propname, **kwargs) #
            setattr(self, prop, preloaded_fn)
            self._logger.info("Successfully loaded {}.".format(prop))
            self._prop_list.append(prop)
        except AttributeError: print('could not find',kwargs['model'])

    def check_consistency(self):
        r'''
        Checks to see if the current geometry conflicts with any other geometry
        '''
        temp = sp.zeros_like(self._net.get_pore_info(label=self.name),dtype=int)
        for item in self._net._geometries:
            temp = temp + sp.array(self._net.get_pore_info(label=item),dtype=int)
        print('Geometry labels overlap in', sp.sum(temp>1),'pores')
        print('Geometry not yet applied to',sp.sum(temp==0),'pores')

        temp = sp.zeros_like(self._net.get_throat_info(label=self.name),dtype=int)
        for item in self._net._geometries:
            temp = temp + sp.array(self._net.get_throat_info(label=item),dtype=int)
        print('Geometry labels overlap in', sp.sum(temp>1),'throats')
        print('Geometry not yet applied to',sp.sum(temp==0),'throats')

if __name__ == '__main__':
    pn = OpenPNM.Network.TestNet()
    loc = pn.get_pore_indices()
    test = OpenPNM.Geometry.GenericGeometry(name='doc_test',locations=loc,network=pn)

<|endoftext|>


<commit_before># Copyright (C) 2013-2014 DNAnexus, Inc.
#
# This file is part of dx-toolkit (DNAnexus platform client libraries).
#
#   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.

'''
This submodule gives basic utilities for printing to the terminal.
'''

import textwrap, subprocess, os, sys
from .env import sys_encoding
from ..compat import USING_PYTHON2
from ..exceptions import DXCLIError

if sys.stdout.isatty():
    try:
        tty_rows, tty_cols = map(int, subprocess.check_output(['stty', 'size'], stderr=open(os.devnull, 'w')).split())
        std_width = min(tty_cols - 2, 100)
    except:
        tty_rows, tty_cols = 24, 80
        std_width = 78
    color_state = True
else:
    tty_rows, tty_cols = 24, 80
    std_width = 78
    color_state = False

delimiter = None

def CYAN(message=None):
    if message is None:
        return '\033[36m' if color_state else ''
    else:
        return CYAN() + message + ENDC()

def BLUE(message=None):
    if message is None:
        return '\033[34m' if color_state else ''
    else:
        return BLUE() + message + ENDC()

def YELLOW(message=None):
    if message is None:
        return '\033[33m' if color_state else ''
    else:
        return YELLOW() + message + ENDC()

def GREEN(message=None):
    if message is None:
        return '\033[32m' if color_state else ''
    else:
        return GREEN() + message + ENDC()

def RED(message=None):
    if message is None:
        return '\033[31m' if color_state else ''
    else:
        return RED() + message + ENDC()

def WHITE(message=None):
    if message is None:
        return '\033[37m' if color_state else ''
    else:
        return WHITE() + message + ENDC()

def UNDERLINE(message=None):
    if message is None:
        return '\033[4m' if color_state else ''
    else:
        return UNDERLINE() + message + ENDC()

def BOLD(message=None):
    if message is None:
        return '\033[1m' if color_state else ''
    else:
        return BOLD() + message + ENDC()

def ENDC():
    return '\033[0m' if color_state else ''

def DNANEXUS_LOGO():
    return BOLD() + WHITE() + 'DNAne' + CYAN() + 'x' + WHITE() + 'us' + ENDC()

def DNANEXUS_X():
    return BOLD() + CYAN() + 'x' + WHITE() + ENDC()

def set_colors(state=True):
    global color_state
    color_state = state

def set_delimiter(delim=None):
    global delimiter
    delimiter = delim

def get_delimiter(delim=None):
    global delimiter
    return delimiter

def DELIMITER(alt_delim):
    global delimiter
    return alt_delim if delimiter is None else delimiter

def fill(string, width_adjustment=0, **kwargs):
    if "width" not in kwargs:
        kwargs['width'] = max(std_width + width_adjustment, 20)
    if "break_on_hyphens" not in kwargs:
        kwargs["break_on_hyphens"] = False
    return textwrap.fill(string, **kwargs)

def pager(content, pager=None, file=None):
    if file is None:
        file = sys.stdout

    pager_process = None
    try:
        if file != sys.stdout or not file.isatty():
            raise DXCLIError() # Just print the content, don't use a pager
        content_lines = content.splitlines()
        content_rows = len(content_lines)
        content_cols = max(len(i) for i in content_lines)

        if tty_rows > content_rows and tty_cols > content_cols:
            raise DXCLIError() # Just print the content, don't use a pager

        pager_process = subprocess.Popen(pager or os.environ.get('PAGER', 'less -RS'),
                                         shell=True, stdin=subprocess.PIPE, stdout=file)
        pager_process.stdin.write(content.encode(sys_encoding))
        pager_process.stdin.close()
        pager_process.wait()
    except:
        file.write(content.encode(sys_encoding) if USING_PYTHON2 else content)
    finally:
        try:
            pager_process.terminate()
        except:
            pass

def refill_paragraphs(string, ignored_prefix='    '):
    """Refills the given text, where the text is composed of paragraphs
    separated by blank lines (i.e. '\n\n'). Lines that begin with
    ignored_prefix are not touched; this can be used to keep indented
    code snippets from being incorrectly reformatted.

    """
    paragraphs = string.split('\n\n')
    refilled_paragraphs = [fill(paragraph) if not paragraph.startswith(ignored_prefix) else paragraph for paragraph in paragraphs]
    return '\n\n'.join(refilled_paragraphs).strip('\n')
<commit_msg>dxpy.utils.printing.pager: Detect errors in pager<commit_after># Copyright (C) 2013-2014 DNAnexus, Inc.
#
# This file is part of dx-toolkit (DNAnexus platform client libraries).
#
#   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.

'''
This submodule gives basic utilities for printing to the terminal.
'''

import textwrap, subprocess, os, sys
from .env import sys_encoding
from ..compat import USING_PYTHON2
from ..exceptions import DXCLIError

if sys.stdout.isatty():
    try:
        tty_rows, tty_cols = map(int, subprocess.check_output(['stty', 'size'], stderr=open(os.devnull, 'w')).split())
        std_width = min(tty_cols - 2, 100)
    except:
        tty_rows, tty_cols = 24, 80
        std_width = 78
    color_state = True
else:
    tty_rows, tty_cols = 24, 80
    std_width = 78
    color_state = False

delimiter = None

def CYAN(message=None):
    if message is None:
        return '\033[36m' if color_state else ''
    else:
        return CYAN() + message + ENDC()

def BLUE(message=None):
    if message is None:
        return '\033[34m' if color_state else ''
    else:
        return BLUE() + message + ENDC()

def YELLOW(message=None):
    if message is None:
        return '\033[33m' if color_state else ''
    else:
        return YELLOW() + message + ENDC()

def GREEN(message=None):
    if message is None:
        return '\033[32m' if color_state else ''
    else:
        return GREEN() + message + ENDC()

def RED(message=None):
    if message is None:
        return '\033[31m' if color_state else ''
    else:
        return RED() + message + ENDC()

def WHITE(message=None):
    if message is None:
        return '\033[37m' if color_state else ''
    else:
        return WHITE() + message + ENDC()

def UNDERLINE(message=None):
    if message is None:
        return '\033[4m' if color_state else ''
    else:
        return UNDERLINE() + message + ENDC()

def BOLD(message=None):
    if message is None:
        return '\033[1m' if color_state else ''
    else:
        return BOLD() + message + ENDC()

def ENDC():
    return '\033[0m' if color_state else ''

def DNANEXUS_LOGO():
    return BOLD() + WHITE() + 'DNAne' + CYAN() + 'x' + WHITE() + 'us' + ENDC()

def DNANEXUS_X():
    return BOLD() + CYAN() + 'x' + WHITE() + ENDC()

def set_colors(state=True):
    global color_state
    color_state = state

def set_delimiter(delim=None):
    global delimiter
    delimiter = delim

def get_delimiter(delim=None):
    global delimiter
    return delimiter

def DELIMITER(alt_delim):
    global delimiter
    return alt_delim if delimiter is None else delimiter

def fill(string, width_adjustment=0, **kwargs):
    if "width" not in kwargs:
        kwargs['width'] = max(std_width + width_adjustment, 20)
    if "break_on_hyphens" not in kwargs:
        kwargs["break_on_hyphens"] = False
    return textwrap.fill(string, **kwargs)

def pager(content, pager=None, file=None):
    if file is None:
        file = sys.stdout

    pager_process = None
    try:
        if file != sys.stdout or not file.isatty():
            raise DXCLIError() # Just print the content, don't use a pager
        content_lines = content.splitlines()
        content_rows = len(content_lines)
        content_cols = max(len(i) for i in content_lines)

        if tty_rows > content_rows and tty_cols > content_cols:
            raise DXCLIError() # Just print the content, don't use a pager

        pager_process = subprocess.Popen(pager or os.environ.get('PAGER', 'less -RS'),
                                         shell=True, stdin=subprocess.PIPE, stdout=file)
        pager_process.stdin.write(content.encode(sys_encoding))
        pager_process.stdin.close()
        pager_process.wait()
        if pager_process.returncode != os.EX_OK:
            raise DXCLIError() # Pager had a problem, print the content without it
    except:
        file.write(content.encode(sys_encoding) if USING_PYTHON2 else content)
    finally:
        try:
            pager_process.terminate()
        except:
            pass

def refill_paragraphs(string, ignored_prefix='    '):
    """Refills the given text, where the text is composed of paragraphs
    separated by blank lines (i.e. '\n\n'). Lines that begin with
    ignored_prefix are not touched; this can be used to keep indented
    code snippets from being incorrectly reformatted.

    """
    paragraphs = string.split('\n\n')
    refilled_paragraphs = [fill(paragraph) if not paragraph.startswith(ignored_prefix) else paragraph for paragraph in paragraphs]
    return '\n\n'.join(refilled_paragraphs).strip('\n')
<|endoftext|>


<commit_before>#  Copyright 2008-2010 Nokia Siemens Networks Oyj
#
#  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.

import sys
import os
import tempfile
if sys.platform.startswith('java'):
    from java.awt import Toolkit, Robot, Rectangle
    from javax.imageio import ImageIO
    from java.io import File
else:
    try:
        import wx
        _wx_app_reference = wx.PySimpleApp()
    except ImportError:
        wx = None
    try:
        from gtk import gdk
    except ImportError:
        gdk = None
    try:
        from PIL import ImageGrab  # apparently available only on Windows
    except ImportError:
        ImageGrab = None

from robot.version import get_version
from robot import utils
from robot.libraries.BuiltIn import BuiltIn

class Screenshot:

    """A test library for taking screenshots and embedding them into log files.

    *Using with Jython*

    On Jython this library uses Java AWT APIs. They are always available
    and thus no external modules are needed.

    *Using with Python*

    On Python you need to have one of the following modules installed to be
    able to use this library. The first module that is found will be used.

    - wxPython :: http://wxpython.org :: Required also by RIDE so many Robot
      Framework users already have this module installed.
    - PyGTK :: http://pygtk.org :: This module is available by default on most
      Linux distributions.
    - Python Imaging Library (PIL) :: http://www.pythonware.com/products/pil ::
      This module can take screenshots only on Windows.

    *Where screenshots are saved*

    By default screenshots are saved into the same directory where the Robot
    Framework log file is written. If no log is created, screenshots are saved
    into the directory where the XML output file is written.

    It is possible to specify a custom location for screenshots using
   `screenshot_directory` argument in `importing` and `Set Screenshot Directory`
    keyword during execution. It is also possible to save screenshots using
    an absolute path.

    Note that prior to Robot Framework 2.5.5 the default screenshot location
    was system's temporary directory.

    *Changes in Robot Framework 2.5.5*

    This library was enhanced heavily in Robot Framework 2.5.5 release. The
    changes are listed below and explained more thoroughly in affected places.

    - The support for using this library on Python (see above) was added.
    - The default location where screenshots are saved was changed (see above).
    - New `Take Screenshot` and `Take Screenshot Without Embedding` keywords
      were added. These keywords should be used for taking screenshots in
      the future. Other screenshot taking keywords will be deprecated and
      removed later.
    - `log_file_directory` argument was deprecated everywhere it was used..
    """

    ROBOT_LIBRARY_SCOPE = 'TEST SUITE'
    ROBOT_LIBRARY_VERSION = get_version()

    def __init__(self, screenshot_directory=None, log_file_directory='DEPRECATED'):
        """Optionally save screenshots into a custom directory.

        Examples (use only one of these):

        | *Setting* | *Value*  | *Value*    | *Value* |
        | Library | Screenshot |            | # Default location |
        | Library | Screenshot | ${TEMPDIR} | # System temp (this was default prior to 2.5.5) |

        `log_file_directory` has been deprecated in 2.5.5 release and has no
        effect. The information provided with it earlier is nowadays got
        automatically. This argument will be removed in the 2.6 release.
        """
        if log_file_directory != 'DEPRECATED':
            print '*WARN* TODO'
        self._given_screenshot_dir = self._norm_path(screenshot_directory)
        self._screenshot_taker = ScreenshotTaker()

    def _norm_path(self, path):
        if not path:
            return path
        return os.path.normpath(path.replace('/', os.sep))

    @property
    def _screenshot_dir(self):
        return self._given_screenshot_dir or self._log_dir

    @property
    def _log_dir(self):
        variables = BuiltIn().get_variables()
        outdir = variables['${OUTPUTDIR}']
        log = variables['${LOGFILE}']
        log = os.path.dirname(log) if log != 'NONE' else '.'
        return os.path.join(outdir, log)

    def set_screenshot_directory(self, path):
        """Sets the directory where screenshots are saved.

        The old value is returned.
        """
        path = self._norm_path(path)
        if not os.path.isdir(path):
            raise RuntimeError("Directory '%s' does not exist." % path)
        old = self._given_screenshot_dir
        self._given_screenshot_dir = path
        return old

    def set_screenshot_directories(self, default_directory=None,
                                   log_file_directory='DEPRECATED'):
        """*DEPRECATED* Use `Set Screenshot Directory` keyword instead."""
        self.set_screenshot_directory(default_directory)

    def save_screenshot_to(self, path):
        """Saves a screenshot to the specified file.

        *This keyword is obsolete.* Use `Take Screenshot` or `Take Screenshot
        Without Embedding` instead. This keyword will be deprecated in Robot
        Framework 2.6 and removed later.
        """
        return self._screenshot_to_file(path)

    def _screenshot_to_file(self, path, loglevel='INFO'):
        path = os.path.abspath(path.replace('/', os.sep))
        self._validate_screenshot_path(path)
        print '*DEBUG* Using %s modules for taking screenshot.' \
            % self._screenshot_taker.module
        self._screenshot_taker(path)
        print "*%s* Screenshot saved to '%s'" % (loglevel, path)
        return path

    def _validate_screenshot_path(self, path):
        if not os.path.exists(os.path.dirname(path)):
            raise RuntimeError("Directory '%s' where to save the screenshot does "
                               "not exist" % os.path.dirname(path))

    def save_screenshot(self, basename="screenshot", directory=None):
        """Saves a screenshot with a generated unique name.

        *This keyword is obsolete.* Use `Take Screenshot` or `Take Screenshot
        Without Embedding` instead. This keyword will be deprecated in Robot
        Framework 2.6 and removed later.

        The unique name is derived based on the provided `basename` and
        `directory` passed in as optional arguments. If a `directory`
        is provided, the screenshot is saved under that directory.
        Otherwise, the `default_directory` set during the library
        import or by the keyword `Set Screenshot Directories` is used.
        If a `basename` for the screenshot file is provided, a unique
        filename is determined by appending an underscore and a running
        counter. Otherwise, the `basename` defaults to 'screenshot'.

        The path where the screenshot is saved is returned.

        Examples:
        | Save Screenshot | mypic | /home/user | # (1) |
        | Save Screenshot | mypic |            | # (2) |
        | Save Screenshot |       |            | # (3) |
        =>
        1. /home/user/mypic_1.jpg, /home/user/mypic_2.jpg, ...
        2. /tmp/mypic_1.jpg, /tmp/mypic_2.jpg, ...
        3. /tmp/screenshot_1.jpg, /tmp/screenshot_2.jpg, ...
        """
        return self._save_screenshot(basename, directory)

    def _save_screenshot(self, basename, directory=None):
        path = self._get_screenshot_path(basename, directory)
        return self._screenshot_to_file(path)

    def _get_screenshot_path(self, basename, directory):
        directory = self._norm_path(directory) if directory else self._screenshot_dir
        if basename.endswith('.jpg'):
            return os.path.join(directory, basename)
        index = 0
        while True:
            index += 1
            path = os.path.join(directory, "%s_%d.jpg" % (basename, index))
            if not os.path.exists(path):
                return path

    def log_screenshot(self, basename='screenshot', directory=None,
                       log_file_directory='DEPRECATED', width='100%'):
        """Takes a screenshot and logs it to Robot Framework's log file.

        *This keyword is obsolete.* Use `Take Screenshot` or `Take Screenshot
        Without Embedding` instead. This keyword will be deprecated in Robot
        Framework 2.6 and removed later.

        Saves the files as defined in the keyword `Save Screenshot` and creates
        a picture to Robot Framework's log. `directory` defines the directory
        where the screenshots are saved. By default, its value is
        `screenshot_directory`, which is set at the library import or with the
        keyword `Set Screenshot Directory`.

        `log_file_directory` has been deprecated in 2.5.5 release and has no
        effect. The information provided with it earlier is nowadays got
        automatically. This argument will be removed in the 2.6 release.
        """
        if log_file_directory != 'DEPRECATED':
            print '*WARN* TODO'
        return self._log_screenshot_as_html(basename, width, directory)

    def _log_screenshot_as_html(self, basename, width, directory=None):
        path = self._save_screenshot(basename, directory)
        link = utils.get_link_path(path, self._log_dir)
        print '*HTML* ' \
              % (link, link, width)
        return path

    def take_screenshot(self, basename="screenshot", width="800px"):
        return self._log_screenshot_as_html(basename, width=width)

class ScreenshotTaker(object):

    def __init__(self, module_name=None):
        self._screenshot = self._get_screenshot_taker(module_name)
        self.module = self._screenshot.__name__.split('_')[1]

    def __call__(self, path):
        self._screenshot(path)

    def _get_screenshot_taker(self, module_name):
        if sys.platform.startswith('java'):
            return self._java_screenshot
        if module_name:
            method_name = '_%s_screenshot' % module_name.lower()
            if hasattr(self, method_name):
                return getattr(self, method_name)
        return self._get_default_screenshot_taker()

    def _get_default_screenshot_taker(self):
        for module, screenshot_taker in [(wx, self._wx_screenshot),
                                         (gdk, self._gtk_screenshot),
                                         (ImageGrab, self._pil_screenshot),
                                         (True, self._no_screenshot)]:
            if module:
                return screenshot_taker

    def _java_screenshot(self, path):
        size = Toolkit.getDefaultToolkit().getScreenSize()
        rectangle = Rectangle(0, 0, size.width, size.height)
        image = Robot().createScreenCapture(rectangle)
        ImageIO.write(image, 'jpg', File(path))

    def _wx_screenshot(self, path):
        context = wx.ScreenDC()
        width, height = context.GetSize()
        bitmap = wx.EmptyBitmap(width, height, -1)
        memory = wx.MemoryDC()
        memory.SelectObject(bitmap)
        memory.Blit(0, 0, width, height, context, -1, -1)
        memory.SelectObject(wx.NullBitmap)
        bitmap.SaveFile(path, wx.BITMAP_TYPE_JPEG)

    def _gtk_screenshot(self, path):
        window = gdk.get_default_root_window()
        if not window:
            raise RuntimeError('Taking screenshot failed')
        width, height = window.get_size()
        pb = gdk.Pixbuf(gdk.COLORSPACE_RGB, False, 8, width, height)
        pb = pb.get_from_drawable(window, window.get_colormap(),
                                  0, 0, 0, 0, width, height)
        if not pb:
            raise RuntimeError('Taking screenshot failed')
        pb.save(path, 'jpeg')

    def _pil_screenshot(self, path):
        ImageGrab.grab().save(path, 'JPEG')

    def _no_screenshot(self, path):
        raise RuntimeError('Taking screenshots is not supported on this platform '
                           'by default. See library documentation for details.')

if __name__ == "__main__":
    if len(sys.argv) not in [2, 3]:
        print "Usage: %s path [wx|gtk|pil]" % os.path.basename(sys.argv[0])
        sys.exit(1)
    path = os.path.abspath(sys.argv[1])
    module = sys.argv[2] if len(sys.argv) == 3 else None
    shooter = ScreenshotTaker(module)
    print 'Using %s modules' % shooter.module
    shooter(path)
    print path
<commit_msg>Screenshot: Added deprecation warning when `log_file_directory` is used and doc for Take Screenshot<commit_after>#  Copyright 2008-2010 Nokia Siemens Networks Oyj
#
#  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.

import sys
import os
import tempfile
if sys.platform.startswith('java'):
    from java.awt import Toolkit, Robot, Rectangle
    from javax.imageio import ImageIO
    from java.io import File
else:
    try:
        import wx
        _wx_app_reference = wx.PySimpleApp()
    except ImportError:
        wx = None
    try:
        from gtk import gdk
    except ImportError:
        gdk = None
    try:
        from PIL import ImageGrab  # apparently available only on Windows
    except ImportError:
        ImageGrab = None

from robot.version import get_version
from robot import utils
from robot.libraries.BuiltIn import BuiltIn

class Screenshot:

    """A test library for taking screenshots and embedding them into log files.

    *Using with Jython*

    On Jython this library uses Java AWT APIs. They are always available
    and thus no external modules are needed.

    *Using with Python*

    On Python you need to have one of the following modules installed to be
    able to use this library. The first module that is found will be used.

    - wxPython :: http://wxpython.org :: Required also by RIDE so many Robot
      Framework users already have this module installed.
    - PyGTK :: http://pygtk.org :: This module is available by default on most
      Linux distributions.
    - Python Imaging Library (PIL) :: http://www.pythonware.com/products/pil ::
      This module can take screenshots only on Windows.

    *Where screenshots are saved*

    By default screenshots are saved into the same directory where the Robot
    Framework log file is written. If no log is created, screenshots are saved
    into the directory where the XML output file is written.

    It is possible to specify a custom location for screenshots using
   `screenshot_directory` argument in `importing` and `Set Screenshot Directory`
    keyword during execution. It is also possible to save screenshots using
    an absolute path.

    Note that prior to Robot Framework 2.5.5 the default screenshot location
    was system's temporary directory.

    *Changes in Robot Framework 2.5.5*

    This library was enhanced heavily in Robot Framework 2.5.5 release. The
    changes are listed below and explained more thoroughly in affected places.

    - The support for using this library on Python (see above) was added.
    - The default location where screenshots are saved was changed (see above).
    - New `Take Screenshot` and `Take Screenshot Without Embedding` keywords
      were added. These keywords should be used for taking screenshots in
      the future. Other screenshot taking keywords will be deprecated and
      removed later.
    - `log_file_directory` argument was deprecated everywhere it was used..
    """

    ROBOT_LIBRARY_SCOPE = 'TEST SUITE'
    ROBOT_LIBRARY_VERSION = get_version()

    def __init__(self, screenshot_directory=None, log_file_directory='DEPRECATED'):
        """Optionally save screenshots into a custom directory.

        Examples (use only one of these):

        | *Setting* | *Value*  | *Value*    | *Value* |
        | Library | Screenshot |            | # Default location |
        | Library | Screenshot | ${TEMPDIR} | # System temp (this was default prior to 2.5.5) |

        `log_file_directory` has been deprecated in 2.5.5 release and has no
        effect. The information provided with it earlier is nowadays got
        automatically. This argument will be removed in the 2.6 release.
        """
        self._warn_if_deprecated_argument_was_used(log_file_directory)
        self._given_screenshot_dir = self._norm_path(screenshot_directory)
        self._screenshot_taker = ScreenshotTaker()

    def _norm_path(self, path):
        if not path:
            return path
        return os.path.normpath(path.replace('/', os.sep))

    @property
    def _screenshot_dir(self):
        return self._given_screenshot_dir or self._log_dir

    @property
    def _log_dir(self):
        variables = BuiltIn().get_variables()
        outdir = variables['${OUTPUTDIR}']
        log = variables['${LOGFILE}']
        log = os.path.dirname(log) if log != 'NONE' else '.'
        return os.path.join(outdir, log)

    def set_screenshot_directory(self, path):
        """Sets the directory where screenshots are saved.

        The old value is returned.
        """
        path = self._norm_path(path)
        if not os.path.isdir(path):
            raise RuntimeError("Directory '%s' does not exist." % path)
        old = self._given_screenshot_dir
        self._given_screenshot_dir = path
        return old

    def set_screenshot_directories(self, default_directory=None,
                                   log_file_directory='DEPRECATED'):
        """*DEPRECATED* Use `Set Screenshot Directory` keyword instead."""
        self.set_screenshot_directory(default_directory)

    def save_screenshot_to(self, path):
        """Saves a screenshot to the specified file.

        *This keyword is obsolete.* Use `Take Screenshot` or `Take Screenshot
        Without Embedding` instead. This keyword will be deprecated in Robot
        Framework 2.6 and removed later.
        """
        return self._screenshot_to_file(path)

    def _screenshot_to_file(self, path, loglevel='INFO'):
        path = os.path.abspath(path.replace('/', os.sep))
        self._validate_screenshot_path(path)
        print '*DEBUG* Using %s modules for taking screenshot.' \
            % self._screenshot_taker.module
        self._screenshot_taker(path)
        print "*%s* Screenshot saved to '%s'" % (loglevel, path)
        return path

    def _validate_screenshot_path(self, path):
        if not os.path.exists(os.path.dirname(path)):
            raise RuntimeError("Directory '%s' where to save the screenshot does "
                               "not exist" % os.path.dirname(path))

    def save_screenshot(self, basename="screenshot", directory=None):
        """Saves a screenshot with a generated unique name.

        *This keyword is obsolete.* Use `Take Screenshot` or `Take Screenshot
        Without Embedding` instead. This keyword will be deprecated in Robot
        Framework 2.6 and removed later.

        The unique name is derived based on the provided `basename` and
        `directory` passed in as optional arguments. If a `directory`
        is provided, the screenshot is saved under that directory.
        Otherwise, the `default_directory` set during the library
        import or by the keyword `Set Screenshot Directories` is used.
        If a `basename` for the screenshot file is provided, a unique
        filename is determined by appending an underscore and a running
        counter. Otherwise, the `basename` defaults to 'screenshot'.

        The path where the screenshot is saved is returned.

        Examples:
        | Save Screenshot | mypic | /home/user | # (1) |
        | Save Screenshot | mypic |            | # (2) |
        | Save Screenshot |       |            | # (3) |
        =>
        1. /home/user/mypic_1.jpg, /home/user/mypic_2.jpg, ...
        2. /tmp/mypic_1.jpg, /tmp/mypic_2.jpg, ...
        3. /tmp/screenshot_1.jpg, /tmp/screenshot_2.jpg, ...
        """
        return self._save_screenshot(basename, directory)

    def _save_screenshot(self, basename, directory=None):
        path = self._get_screenshot_path(basename, directory)
        return self._screenshot_to_file(path)

    def _get_screenshot_path(self, basename, directory):
        directory = self._norm_path(directory) if directory else self._screenshot_dir
        if basename.endswith('.jpg'):
            return os.path.join(directory, basename)
        index = 0
        while True:
            index += 1
            path = os.path.join(directory, "%s_%d.jpg" % (basename, index))
            if not os.path.exists(path):
                return path

    def log_screenshot(self, basename='screenshot', directory=None,
                       log_file_directory='DEPRECATED', width='100%'):
        """Takes a screenshot and logs it to Robot Framework's log file.

        *This keyword is obsolete.* Use `Take Screenshot` or `Take Screenshot
        Without Embedding` instead. This keyword will be deprecated in Robot
        Framework 2.6 and removed later.

        Saves the files as defined in the keyword `Save Screenshot` and creates
        a picture to Robot Framework's log. `directory` defines the directory
        where the screenshots are saved. By default, its value is
        `screenshot_directory`, which is set at the library import or with the
        keyword `Set Screenshot Directory`.

        `log_file_directory` has been deprecated in 2.5.5 release and has no
        effect. The information provided with it earlier is nowadays got
        automatically. This argument will be removed in the 2.6 release.
        """
        self._warn_if_deprecated_argument_was_used(log_file_directory)
        return self._log_screenshot_as_html(basename, width, directory)

    def take_screenshot(self, name="screenshot", width="800px"):
        """Takes a screenshot and logs it to Robot Framework's log file.

        Name of the file where the screenshot is stored is derived from `name`.
        If `name` ends with extension `.jpg` the screenshot will be stored with
        that name, otherwise a unique name is derived by adding underscore, a
        running index and extension to the `name`.

        The name will be interpreted to be relative to the directory where
        log file is written.

        `width` can be used to determine the size of the screenshot that is
        visible in the log file.

        The path where the screenshot is saved is returned.

        Examples: (LOGDIR is determined automatically by the library)
        | Save Screenshot | mypic             |  | # (1) |
        | Save Screenshot | ${TEMPDIR}/mypic  |  | # (2) |
        | Save Screenshot | pic.jpg           |  | # (3) |
        =>
        1. LOGDIR/mypic_1.jpg, LOGDIR/mypic_2.jpg, ...
        2. /tmp/mypic_1.jpg, /tmp/mypic_2.jpg, ...
        3. LOGDIR/pic.jpg, LOGDIR/pic.jpg
        """
        return self._log_screenshot_as_html(name, width=width)

    def _warn_if_deprecated_argument_was_used(self, log_file_directory):
        if log_file_directory != 'DEPRECATED':
            print '*WARN* argument `log_file_directory` is deprecated and should not be used.'

class ScreenshotTaker(object):

    def __init__(self, module_name=None):
        self._screenshot = self._get_screenshot_taker(module_name)
        self.module = self._screenshot.__name__.split('_')[1]

    def __call__(self, path):
        self._screenshot(path)

    def _get_screenshot_taker(self, module_name):
        if sys.platform.startswith('java'):
            return self._java_screenshot
        if module_name:
            method_name = '_%s_screenshot' % module_name.lower()
            if hasattr(self, method_name):
                return getattr(self, method_name)
        return self._get_default_screenshot_taker()

    def _get_default_screenshot_taker(self):
        for module, screenshot_taker in [(wx, self._wx_screenshot),
                                         (gdk, self._gtk_screenshot),
                                         (ImageGrab, self._pil_screenshot),
                                         (True, self._no_screenshot)]:
            if module:
                return screenshot_taker

    def _java_screenshot(self, path):
        size = Toolkit.getDefaultToolkit().getScreenSize()
        rectangle = Rectangle(0, 0, size.width, size.height)
        image = Robot().createScreenCapture(rectangle)
        ImageIO.write(image, 'jpg', File(path))

    def _wx_screenshot(self, path):
        context = wx.ScreenDC()
        width, height = context.GetSize()
        bitmap = wx.EmptyBitmap(width, height, -1)
        memory = wx.MemoryDC()
        memory.SelectObject(bitmap)
        memory.Blit(0, 0, width, height, context, -1, -1)
        memory.SelectObject(wx.NullBitmap)
        bitmap.SaveFile(path, wx.BITMAP_TYPE_JPEG)

    def _gtk_screenshot(self, path):
        window = gdk.get_default_root_window()
        if not window:
            raise RuntimeError('Taking screenshot failed')
        width, height = window.get_size()
        pb = gdk.Pixbuf(gdk.COLORSPACE_RGB, False, 8, width, height)
        pb = pb.get_from_drawable(window, window.get_colormap(),
                                  0, 0, 0, 0, width, height)
        if not pb:
            raise RuntimeError('Taking screenshot failed')
        pb.save(path, 'jpeg')

    def _pil_screenshot(self, path):
        ImageGrab.grab().save(path, 'JPEG')

    def _no_screenshot(self, path):
        raise RuntimeError('Taking screenshots is not supported on this platform '
                           'by default. See library documentation for details.')

if __name__ == "__main__":
    if len(sys.argv) not in [2, 3]:
        print "Usage: %s path [wx|gtk|pil]" % os.path.basename(sys.argv[0])
        sys.exit(1)
    path = os.path.abspath(sys.argv[1])
    module = sys.argv[2] if len(sys.argv) == 3 else None
    shooter = ScreenshotTaker(module)
    print 'Using %s modules' % shooter.module
    shooter(path)
    print path
<|endoftext|>