import json

from opt.appearance import make_colors

from os import listdir
from os.path import join, isfile, splitext

from difflib import SequenceMatcher

from terminaltables import SingleTable, DoubleTable

EMPTY_VALUE = '-'

NAME_FIELD  = 'name'
DEF_FIELD   = 'definition'
ID_FIELD    = 'id'

NO_MATCHES_ERROR = '<red,bold>[-]</> No matches. Try again!'

class SysHandler:
    def __init__(self):
        self.tables         = {}
        self.req_similarity = 0.75

        self.tcache = {}

    def get_printable_table(self, arch, pattern, colored=False, verbose=False):
        table = self.get_table(arch, pattern, colored, verbose)
        return table.table if table is not None else NO_MATCHES_ERROR

    def get_table(self, arch, pattern, colored=False, verbose=False):
        '''
        This function is used in sys command (when user want to find a specific syscall)

        :param Architecture for syscall table;
        :param Searching pattern;
        :param Flag for verbose output
        :return Return a printable table of matched syscalls
        '''

        rawtable = self.search(arch, pattern)
        if len(rawtable) == 0:
            return None

        used_hd = self.__fetch_used_headers(rawtable, verbose)
        table   = [self.__make_colored_row(used_hd, 'yellow,bold', upper=True) if colored else used_hd]

        for command in rawtable:
            cur_tb_field = []
            for hd in used_hd:
                value = command[hd]
                cur_tb_field.append(self.__make_colored_field(value, hd, verbose=verbose))
            table.append(cur_tb_field)
        return DoubleTable(table)

    def __make_colored_field(self, field, hd, verbose=False):
        if hd == NAME_FIELD:
            return make_colors('<red>{}</>'.format(field))
        elif self.__is_number(field):
            return make_colors('<green>{}</>'.format(field))
        else:
            parts = field.split(' ')
            if verbose:
                return make_colors('<magenta>{}</> <white,bold>{}</>'.format(' '.join(parts[:-1]), parts[-1]))
            else:
                return make_colors('<white,bold>{}</>'.format(parts[-1]))

    def __is_number(self, s, base=16):
        try:
            int(s, base)
            return True
        except Exception:
            return False

    def __make_colored_row(self, row, pcolor, upper=False):
        return [make_colors('<{}>{}</>'.format(pcolor, val.upper() if upper else val)) for val in row]

    def __fetch_used_headers(self, table, verbose=False):
        def hdkey(hd):
            return -len(hd), hd.upper()

        used_hd = set()
        for command in table:
            for header, value in command.items():
                if value != EMPTY_VALUE:
                    used_hd.add(header)
        used_hd.remove(NAME_FIELD)
        used_hd.remove(DEF_FIELD)
        used_hd.remove(ID_FIELD)
        return [NAME_FIELD] + sorted(list(used_hd), key=hdkey) + ([DEF_FIELD] if verbose else [])

    def search(self, arch, pattern):
        try:
            table   = self.tables[arch]
            similar = []
            for command in table:
                if self.__similar(command['name'], pattern) >= self.req_similarity:
                    similar.append(command)
            return similar
        except KeyError:
            return []

    def load_tables(self):
        for filename in listdir(self.dir):
            fileloc = join(self.dir, filename)
            name    = splitext(filename)[0]
            if isfile(fileloc):
                with open(fileloc, 'r') as tmpfile:
                    self.tables[name] = json.loads(tmpfile.read())

    def __similar(self, s, f):
        return SequenceMatcher(None, s, f).ratio()


from opt.appearance import cprint

from asms.baseexc import BaseExec, BaseExecWrapper, hex2bytes
from archsconf import *

from keystone import *

from binascii import hexlify, unhexlify

NULLBYTE = '\\x00'

class Assembler(BaseExec):
    def __init__(self, parch):
        super().__init__()

        self.last_shellcode = None
        self.setarch(parch)

    def get_last_shellcode(self):
        return self.last_shellcode

    def avail_archs(self):
        ''' Initialize the dictionary of architectures for assembling via keystone'''

        return {
            ARM32:   (KS_ARCH_ARM,     KS_MODE_ARM),
            ARM64:   (KS_ARCH_ARM64,   KS_MODE_LITTLE_ENDIAN),
            ARM_TB:  (KS_ARCH_ARM,     KS_MODE_THUMB),
            HEXAGON: (KS_ARCH_HEXAGON, KS_MODE_BIG_ENDIAN),
            MIPS32:  (KS_ARCH_MIPS,    KS_MODE_MIPS32),
            MIPS64:  (KS_ARCH_MIPS,    KS_MODE_MIPS64),
            PPC32:   (KS_ARCH_PPC,     KS_MODE_PPC32),
            PPC64:   (KS_ARCH_PPC,     KS_MODE_PPC64),
            SPARC32: (KS_ARCH_SPARC,   KS_MODE_SPARC32),
            SPARC64: (KS_ARCH_SPARC,   KS_MODE_SPARC64),
            SYSTEMZ: (KS_ARCH_SYSTEMZ, KS_MODE_BIG_ENDIAN),
            X86_16:  (KS_ARCH_X86,     KS_MODE_16),
            X86_32:  (KS_ARCH_X86,     KS_MODE_32),
            X86_64:  (KS_ARCH_X86,     KS_MODE_64),
        }

    def update_engine(self):
        self.__ks = Ks(*self.arch)

    def execv(self, data):
        shellcode, num_instructions = self.__ks.asm(data)
        self.last_shellcode = bytes(shellcode)
        return (shellcode, num_instructions)

class AssemblerWrapper(BaseExecWrapper):
    def __init__(self, arch):
        super().__init__(arch)

        self.executor = Assembler(self.arch)

    def __decorate_shellcode(self, sc, pbytes=True):
        cld_sc = ''
        offset = 4 if pbytes else 2
        for pos in range(0, len(sc), offset):
            _byte = sc[pos:pos+offset]
            if _byte == (NULLBYTE if pbytes else NULLBYTE[2:]):
                cld_sc += '<red,bold>' + _byte + '</>'
            else:
                cld_sc += '<white,bold>' + _byte + '</>'
        return cld_sc

    def print_res(self, res):
        encoding, count = res

        raw_hex   = hexlify(bytearray(encoding)).decode('utf-8')
        raw_bytes = hex2bytes(raw_hex)

        dec_sc_bytes = self.__decorate_shellcode(raw_bytes)
        dec_sc_hex   = self.__decorate_shellcode(raw_hex, False)

        is_zeroed = NULLBYTE in raw_bytes

        prefix = ''
        if is_zeroed:
            prefix += '   <yellow,bold>[!]</> Warning! Your shellcode contains <white,underline>null bytes</>!\n'
        cprint(prefix + (
            '   <green,bold>[+]</> Bytes count: <white,bold>{}</>\n'
            '       Raw bytes:  "{}"\n'
            '       Hex string: "{}"\n'
        ).format(len(encoding), dec_sc_bytes, dec_sc_hex))


from opt.appearance import cprint

from asms.baseexc import BaseExec, BaseExecWrapper, hex2bytes
from archsconf import *

from capstone import *

from binascii import hexlify, unhexlify

class Disassembler(BaseExec):
    def __init__(self, parch):
        super().__init__()

        self.setarch(parch)

        self.baseaddr = 0x00080000
        self.last_shellcode = None

        self.update_engine()

    def get_last_shellcode(self):
        return self.last_shellcode

    def avail_archs(self):
        ''' Initialize the dictionary of architectures for disassembling via capstone'''

        return {
            ARM32:   (CS_ARCH_ARM,   CS_MODE_ARM),
            ARM64:   (CS_ARCH_ARM64, CS_MODE_LITTLE_ENDIAN),
            ARM_TB:  (CS_ARCH_ARM,   CS_MODE_THUMB),
            MIPS32:  (CS_ARCH_MIPS,  CS_MODE_MIPS32),
            MIPS64:  (CS_ARCH_MIPS,  CS_MODE_MIPS64),
            SPARC32: (CS_ARCH_SPARC, CS_MODE_BIG_ENDIAN),
            SPARC64: (CS_ARCH_SPARC, CS_MODE_V9),
            SYSTEMZ: (CS_ARCH_SYSZ,  CS_MODE_BIG_ENDIAN),
            X86_16:  (CS_ARCH_X86,   CS_MODE_16),
            X86_32:  (CS_ARCH_X86,   CS_MODE_32),
            X86_64:  (CS_ARCH_X86,   CS_MODE_64),
        }

    def update_engine(self):
        self.__cs = Cs(*self.arch)

    def __parse_bytes(self, data):
        data = data.replace('\\x', '')
        if data[0] == data[-1] and data[0] in ["'", '"']:
            data = data[1:-1]
        return data

    def execv(self, data):
        unhexed = unhexlify(self.__parse_bytes(data))
        disassembled = self.__cs.disasm(unhexed, self.baseaddr)
        self.last_shellcode = bytes(unhexed)
        return disassembled

class DisassemblerWrapper(BaseExecWrapper):
    def __init__(self, arch):
        super().__init__(arch)

        self.executor = Disassembler(self.arch)

    def print_res(self, res):
        for line in res:
            cprint("\t<cyan>0x{:08X}</>:\t<yellow,bold>{:<8}</><white,bold>{}</>".format(line.address, line.mnemonic, line.op_str))
        cprint('\n')


# originally inspired by "six" by Benjamin Peterson

import inspect
import sys

if sys.version_info < (3,0):
    text_type = unicode
    binary_type = str

    import StringIO
    StringIO = BytesIO = StringIO.StringIO
else:
    text_type = str
    binary_type = bytes

    import io
    StringIO = io.StringIO
    BytesIO = io.BytesIO

def getargspec_permissive(func):
    """
    An `inspect.getargspec` with a relaxed sanity check to support Cython.

    Motivation:

        A Cython-compiled function is *not* an instance of Python's
        types.FunctionType.  That is the sanity check the standard Py2
        library uses in `inspect.getargspec()`.  So, an exception is raised
        when calling `argh.dispatch_command(cythonCompiledFunc)`.  However,
        the CyFunctions do have perfectly usable `.func_code` and
        `.func_defaults` which is all `inspect.getargspec` needs.

        This function just copies `inspect.getargspec()` from the standard
        library but relaxes the test to a more duck-typing one of having
        both `.func_code` and `.func_defaults` attributes.
    """
    if inspect.ismethod(func):
        func = func.im_func

    # Py2 Stdlib uses isfunction(func) which is too strict for Cython-compiled
    # functions though such have perfectly usable func_code, func_defaults.
    if not (hasattr(func, "func_code") and hasattr(func, "func_defaults")):
        raise TypeError('{!r} missing func_code or func_defaults'.format(func))

    args, varargs, varkw = inspect.getargs(func.func_code)
    return inspect.ArgSpec(args, varargs, varkw, func.func_defaults)

if sys.version_info < (3,0):
    getargspec = getargspec_permissive
else:
    # in Python 3 the basic getargspec doesn't support keyword-only arguments
    # and annotations and raises ValueError if they are discovered
    getargspec = inspect.getfullargspec


# coding: utf-8
#
#  Copyright © 2010—2014 Andrey Mikhaylenko and contributors
#
#  This file is part of Argh.
#
#  Argh is free software under terms of the GNU Lesser
#  General Public License version 3 (LGPLv3) as published by the Free
#  Software Foundation. See the file README.rst for copying conditions.
#
"""
Dispatching
~~~~~~~~~~~
"""
import argparse
import sys
from types import GeneratorType

from argh import compat, io
from argh.constants import (
    ATTR_WRAPPED_EXCEPTIONS,
    ATTR_WRAPPED_EXCEPTIONS_PROCESSOR,
    ATTR_EXPECTS_NAMESPACE_OBJECT,
    PARSER_FORMATTER,
    DEST_FUNCTION,
)
from argh.completion import autocomplete
from argh.assembling import add_commands, set_default_command
from argh.exceptions import DispatchingError, CommandError
from argh.utils import get_arg_spec

__all__ = ['dispatch', 'dispatch_command', 'dispatch_commands',
           'PARSER_FORMATTER', 'EntryPoint']

class ArghNamespace(argparse.Namespace):
    """
    A namespace object which collects the stack of functions (the
    :attr:`~argh.constants.DEST_FUNCTION` arguments passed to it via
    parser's defaults).
    """
    def __init__(self, *args, **kw):
        super(ArghNamespace, self).__init__(*args, **kw)
        self._functions_stack = []

    def __setattr__(self, k, v):
        if k == DEST_FUNCTION:
            # don't register the function under DEST_FUNCTION name.
            # If `ArgumentParser.parse_known_args()` sees that we already have
            # such attribute, it skips it.  However, it goes from the topmost
            # parser to subparsers.  We need the function mapped to the
            # subparser.  So we fool the `ArgumentParser` and pretend that we
            # didn't get a DEST_FUNCTION attribute; however, in fact we collect
            # all its values in a stack.  The last item in the stack would be
            # the function mapped to the innermost parser — the one we need.
            self._functions_stack.append(v)
        else:
            super(ArghNamespace, self).__setattr__(k, v)

    def get_function(self):
        return self._functions_stack[-1]

def dispatch(parser, argv=None, add_help_command=True,
             completion=True, pre_call=None,
             output_file=sys.stdout, errors_file=sys.stderr,
             raw_output=False, namespace=None,
             skip_unknown_args=False):
    """
    Parses given list of arguments using given parser, calls the relevant
    function and prints the result.

    The target function should expect one positional argument: the
    :class:`argparse.Namespace` object. However, if the function is decorated with
    :func:`~argh.decorators.plain_signature`, the positional and named
    arguments from the namespace object are passed to the function instead
    of the object itself.

    :param parser:

        the ArgumentParser instance.

    :param argv:

        a list of strings representing the arguments. If `None`, ``sys.argv``
        is used instead. Default is `None`.

    :param add_help_command:

        if `True`, converts first positional argument "help" to a keyword
        argument so that ``help foo`` becomes ``foo --help`` and displays usage
        information for "foo". Default is `True`.

    :param output_file:

        A file-like object for output. If `None`, the resulting lines are
        collected and returned as a string. Default is ``sys.stdout``.

    :param errors_file:

        Same as `output_file` but for ``sys.stderr``.

    :param raw_output:

        If `True`, results are written to the output file raw, without adding
        whitespaces or newlines between yielded strings. Default is `False`.

    :param completion:

        If `True`, shell tab completion is enabled. Default is `True`. (You
        will also need to install it.)  See :mod:`argh.completion`.

    :param skip_unknown_args:

        If `True`, unknown arguments do not cause an error
        (`ArgumentParser.parse_known_args` is used).

    :param namespace:

        An `argparse.Namespace`-like object.  By default an
        :class:`ArghNamespace` object is used.  Please note that support for
        combined default and nested functions may be broken if a different
        type of object is forced.

    By default the exceptions are not wrapped and will propagate. The only
    exception that is always wrapped is :class:`~argh.exceptions.CommandError`
    which is interpreted as an expected event so the traceback is hidden.
    You can also mark arbitrary exceptions as "wrappable" by using the
    :func:`~argh.decorators.wrap_errors` decorator.
    """
    if completion:
        autocomplete(parser)

    if argv is None:
        argv = sys.argv[1:]

    if add_help_command:
        if argv and argv[0] == 'help':
            argv.pop(0)
            argv.append('--help')

    if skip_unknown_args:
        parse_args = parser.parse_known_args
    else:
        parse_args = parser.parse_args

    if not namespace:
        namespace = ArghNamespace()

    # this will raise SystemExit if parsing fails
    namespace_obj = parse_args(argv, namespace=namespace)

    function = _get_function_from_namespace_obj(namespace_obj)

    if function:
        lines = _execute_command(function, namespace_obj, errors_file,
                                 pre_call=pre_call)
    else:
        # no commands declared, can't dispatch; display help message
        lines = [parser.format_usage()]

    if output_file is None:
        # user wants a string; we create an internal temporary file-like object
        # and will return its contents as a string
        if sys.version_info < (3,0):
            f = compat.BytesIO()
        else:
            f = compat.StringIO()
    else:
        # normally this is stdout; can be any file
        f = output_file

    for line in lines:
        # print the line as soon as it is generated to ensure that it is
        # displayed to the user before anything else happens, e.g.
        # raw_input() is called

        io.dump(line, f)
        if not raw_output:
            # in most cases user wants one message per line
            io.dump('\n', f)

    if output_file is None:
        # user wanted a string; return contents of our temporary file-like obj
        f.seek(0)
        return f.read()

def _get_function_from_namespace_obj(namespace_obj):
    if isinstance(namespace_obj, ArghNamespace):
        # our special namespace object keeps the stack of assigned functions
        try:
            function = namespace_obj.get_function()
        except (AttributeError, IndexError):
            return None
    else:
        # a custom (probably vanilla) namespace object keeps the last assigned
        # function; this may be wrong but at least something may work
        if not hasattr(namespace_obj, DEST_FUNCTION):
            return None
        function = getattr(namespace_obj, DEST_FUNCTION)

    if not function or not hasattr(function, '__call__'):
        return None

    return function

def _execute_command(function, namespace_obj, errors_file, pre_call=None):
    """
    Assumes that `function` is a callable.  Tries different approaches
    to call it (with `namespace_obj` or with ordinary signature).
    Yields the results line by line.

    If :class:`~argh.exceptions.CommandError` is raised, its message is
    appended to the results (i.e. yielded by the generator as a string).
    All other exceptions propagate unless marked as wrappable
    by :func:`wrap_errors`.
    """
    if pre_call:  # XXX undocumented because I'm unsure if it's OK
        # Actually used in real projects:
        # * https://google.com/search?q=argh+dispatch+pre_call
        # * https://github.com/neithere/argh/issues/63
        pre_call(namespace_obj)

    # the function is nested to catch certain exceptions (see below)
    def _call():
        # Actually call the function
        if getattr(function, ATTR_EXPECTS_NAMESPACE_OBJECT, False):
            result = function(namespace_obj)
        else:
            # namespace -> dictionary
            _flat_key = lambda key: key.replace('-', '_')
            all_input = dict((_flat_key(k), v)
                             for k,v in vars(namespace_obj).items())

            # filter the namespace variables so that only those expected
            # by the actual function will pass

            spec = get_arg_spec(function)

            positional = [all_input[k] for k in spec.args]
            kwonly = getattr(spec, 'kwonlyargs', [])
            keywords = dict((k, all_input[k]) for k in kwonly)

            # *args
            if spec.varargs:
                positional += getattr(namespace_obj, spec.varargs)

            # **kwargs
            varkw = getattr(spec, 'varkw', getattr(spec, 'keywords', []))
            if varkw:
                not_kwargs = [DEST_FUNCTION] + spec.args + [spec.varargs] + kwonly
                for k in vars(namespace_obj):
                    if k.startswith('_') or k in not_kwargs:
                        continue
                    keywords[k] = getattr(namespace_obj, k)

            result = function(*positional, **keywords)

        # Yield the results
        if isinstance(result, (GeneratorType, list, tuple)):
            # yield each line ASAP, convert CommandError message to a line
            for line in result:
                yield line
        else:
            # yield non-empty non-iterable result as a single line
            if result is not None:
                yield result

    wrappable_exceptions = [CommandError]
    wrappable_exceptions += getattr(function, ATTR_WRAPPED_EXCEPTIONS, [])

    try:
        result = _call()
        for line in result:
            yield line
    except tuple(wrappable_exceptions) as e:
        processor = getattr(function, ATTR_WRAPPED_EXCEPTIONS_PROCESSOR,
                            lambda e: '{0.__class__.__name__}: {0}'.format(e))

        errors_file.write(compat.text_type(processor(e)))
        errors_file.write('\n')

def dispatch_command(function, *args, **kwargs):
    """
    A wrapper for :func:`dispatch` that creates a one-command parser.
    Uses :attr:`PARSER_FORMATTER`.

    This::

        dispatch_command(foo)

    ...is a shortcut for::

        parser = ArgumentParser()
        set_default_command(parser, foo)
        dispatch(parser)

    This function can be also used as a decorator.
    """
    parser = argparse.ArgumentParser(formatter_class=PARSER_FORMATTER)
    set_default_command(parser, function)
    dispatch(parser, *args, **kwargs)

def dispatch_commands(functions, *args, **kwargs):
    """
    A wrapper for :func:`dispatch` that creates a parser, adds commands to
    the parser and dispatches them.
    Uses :attr:`PARSER_FORMATTER`.

    This::

        dispatch_commands([foo, bar])

    ...is a shortcut for::

        parser = ArgumentParser()
        add_commands(parser, [foo, bar])
        dispatch(parser)

    """
    parser = argparse.ArgumentParser(formatter_class=PARSER_FORMATTER)
    add_commands(parser, functions)
    dispatch(parser, *args, **kwargs)

class EntryPoint(object):
    """
    An object to which functions can be attached and then dispatched.

    When called with an argument, the argument (a function) is registered
    at this entry point as a command.

    When called without an argument, dispatching is triggered with all
    previously registered commands.

    Usage::

        from argh import EntryPoint

        app = EntryPoint('main', dict(description='This is a cool app'))

        @app
        def ls():
            for i in range(10):
                print i

        @app
        def greet():
            print 'hello'

        if __name__ == '__main__':
            app()

    """
    def __init__(self, name=None, parser_kwargs=None):
        self.name = name or 'unnamed'
        self.commands = []
        self.parser_kwargs = parser_kwargs or {}

    def __call__(self, f=None):
        if f:
            self._register_command(f)
            return f

        return self._dispatch()

    def _register_command(self, f):
        self.commands.append(f)

    def _dispatch(self):
        if not self.commands:
            raise DispatchingError('no commands for entry point "{0}"'
                                   .format(self.name))

        parser = argparse.ArgumentParser(**self.parser_kwargs)
        add_commands(parser, self.commands)
        dispatch(parser)


# coding: utf-8
#
#  Copyright © 2010—2014 Andrey Mikhaylenko and contributors
#
#  This file is part of Argh.
#
#  Argh is free software under terms of the GNU Lesser
#  General Public License version 3 (LGPLv3) as published by the Free
#  Software Foundation. See the file README.rst for copying conditions.
#
"""
Output Processing
~~~~~~~~~~~~~~~~~
"""
import locale
import sys

from argh import compat

__all__ = ['dump', 'encode_output', 'safe_input']

def _input(prompt):
    # this function can be mocked up in tests
    if sys.version_info < (3,0):
        return raw_input(prompt)
    else:
        return input(prompt)

def safe_input(prompt):
    """
    Prompts user for input. Correctly handles prompt message encoding.
    """

    if sys.version_info < (3,0):
        if isinstance(prompt, compat.text_type):
            # Python 2.x: unicode →  bytes
            encoding = locale.getpreferredencoding() or 'utf-8'
            prompt = prompt.encode(encoding)
    else:
        if not isinstance(prompt, compat.text_type):
            # Python 3.x: bytes →  unicode
            prompt = prompt.decode()

    return _input(prompt)

def encode_output(value, output_file):
    """
    Encodes given value so it can be written to given file object.

    Value may be Unicode, binary string or any other data type.

    The exact behaviour depends on the Python version:

    Python 3.x

        `sys.stdout` is a `_io.TextIOWrapper` instance that accepts `str`
        (unicode) and breaks on `bytes`.

        It is OK to simply assume that everything is Unicode unless special
        handling is introduced in the client code.

        Thus, no additional processing is performed.

    Python 2.x

        `sys.stdout` is a file-like object that accepts `str` (bytes)
        and breaks when `unicode` is passed to `sys.stdout.write()`.

        We can expect both Unicode and bytes. They need to be encoded so as
        to match the file object encoding.

        The output is binary if the object doesn't explicitly require Unicode.

    """
    if sys.version_info > (3,0):
        # Python 3:  whatever → unicode
        return compat.text_type(value)
    else:
        # Python 2:  handle special cases
        stream_encoding = getattr(output_file, 'encoding', None)
        if stream_encoding:
            if stream_encoding.upper() == 'UTF-8':
                return compat.text_type(value)
            else:
                return value.encode(stream_encoding, 'ignore')
        else:
            # no explicit encoding requirements; force binary
            if isinstance(value, compat.text_type):
                # unicode → binary
                return value.encode('utf-8')
            else:
                return str(value)

def dump(raw_data, output_file):
    """
    Writes given line to given output file.
    See :func:`encode_output` for details.
    """
    data = encode_output(raw_data, output_file)
    output_file.write(data)


# coding: utf-8
#
#  Copyright © 2010—2014 Andrey Mikhaylenko and contributors
#
#  This file is part of Argh.
#
#  Argh is free software under terms of the GNU Lesser
#  General Public License version 3 (LGPLv3) as published by the Free
#  Software Foundation. See the file README.rst for copying conditions.
#
"""
Shell completion
~~~~~~~~~~~~~~~~

Command and argument completion is a great way to reduce the number of
keystrokes and improve user experience.

To display suggestions when you press :kbd:`tab`, a shell must obtain choices
from your program.  It calls the program in a specific environment and expects
it to return a list of relevant choices.

`Argparse` does not support completion out of the box.  However, there are
3rd-party apps that do the job, such as argcomplete_ and
python-selfcompletion_.

`Argh` supports only argcomplete_ which doesn't require subclassing
the parser and monkey-patches it instead.  Combining `Argh`
with python-selfcompletion_ isn't much harder though: simply use
`SelfCompletingArgumentParser` instead of vanilla `ArgumentParser`.

See installation details and gotchas in the documentation of the 3rd-party app
you've chosen for the completion backend.

`Argh` automatically enables completion if argcomplete_ is available
(see :attr:`COMPLETION_ENABLED`).  If completion is undesirable in given app by
design, it can be turned off by setting ``completion=False``
in :func:`argh.dispatching.dispatch`.

Note that you don't *have* to add completion via `Argh`; it doesn't matter
whether you let it do it for you or use the underlying API.

.. _argcomplete: https://github.com/kislyuk/argcomplete
.. _python-selfcompletion: https://github.com/dbarnett/python-selfcompletion

Argument-level completion
-------------------------

Argcomplete_ supports custom "completers".  The documentation suggests adding
the completer as an attribute of the argument parser action::

    parser.add_argument("--env-var1").completer = EnvironCompleter

However, this doesn't fit the normal `Argh`-assisted workflow.
It is recommended to use the :func:`~argh.decorators.arg` decorator::

    @arg('--env-var1', completer=EnvironCompleter)
    def func(...):
        ...

"""
import logging
import os

COMPLETION_ENABLED = False
"""
Dynamically set to `True` on load if argcomplete_ was successfully imported.
"""

try:
    import argcomplete
except ImportError:
    pass
else:
    COMPLETION_ENABLED = True

__all__ = ['autocomplete', 'COMPLETION_ENABLED']

logger = logging.getLogger(__package__)

def autocomplete(parser):
    """
    Adds support for shell completion via argcomplete_ by patching given
    `argparse.ArgumentParser` (sub)class.

    If completion is not enabled, logs a debug-level message.
    """
    if COMPLETION_ENABLED:
        argcomplete.autocomplete(parser)
    elif 'bash' in os.getenv('SHELL', ''):
        logger.debug('Bash completion not available. Install argcomplete.')


# coding: utf-8
#
#  Copyright © 2010—2014 Andrey Mikhaylenko and contributors
#
#  This file is part of Argh.
#
#  Argh is free software under terms of the GNU Lesser
#  General Public License version 3 (LGPLv3) as published by the Free
#  Software Foundation. See the file README.rst for copying conditions.
#
"""
Helpers
~~~~~~~
"""
import argparse

from argh.completion import autocomplete
from argh.assembling import add_commands, set_default_command
from argh.dispatching import PARSER_FORMATTER, ArghNamespace, dispatch

__all__ = ['ArghParser']

class ArghParser(argparse.ArgumentParser):
    """
    A subclass of :class:`ArgumentParser` with support for and a couple
    of convenience methods.

    All methods are but wrappers for stand-alone functions
    :func:`~argh.assembling.add_commands`,
    :func:`~argh.completion.autocomplete` and
    :func:`~argh.dispatching.dispatch`.

    Uses :attr:`~argh.dispatching.PARSER_FORMATTER`.
    """
    def __init__(self, *args, **kwargs):
        kwargs.setdefault('formatter_class', PARSER_FORMATTER)
        super(ArghParser, self).__init__(*args, **kwargs)

    def set_default_command(self, *args, **kwargs):
        "Wrapper for :func:`~argh.assembling.set_default_command`."
        return set_default_command(self, *args, **kwargs)

    def add_commands(self, *args, **kwargs):
        "Wrapper for :func:`~argh.assembling.add_commands`."
        return add_commands(self, *args, **kwargs)

    def autocomplete(self):
        "Wrapper for :func:`~argh.completion.autocomplete`."
        return autocomplete(self)

    def dispatch(self, *args, **kwargs):
        "Wrapper for :func:`~argh.dispatching.dispatch`."
        return dispatch(self, *args, **kwargs)

    def parse_args(self, args=None, namespace=None):
        """
        Wrapper for :meth:`argparse.ArgumentParser.parse_args`.  If `namespace`
        is not defined, :class:`argh.dispatching.ArghNamespace` is used.
        This is required for functions to be properly used as commands.
        """
        namespace = namespace or ArghNamespace()
        return super(ArghParser, self).parse_args(args, namespace)


# coding: utf-8
#
#  Copyright © 2010—2014 Andrey Mikhaylenko and contributors
#
#  This file is part of Argh.
#
#  Argh is free software under terms of the GNU Lesser
#  General Public License version 3 (LGPLv3) as published by the Free
#  Software Foundation. See the file README.rst for copying conditions.
#
import argparse

__all__ = (
    'ATTR_NAME', 'ATTR_ALIASES', 'ATTR_ARGS', 'ATTR_WRAPPED_EXCEPTIONS',
    'ATTR_WRAPPED_EXCEPTIONS_PROCESSOR', 'ATTR_EXPECTS_NAMESPACE_OBJECT',
    'PARSER_FORMATTER', 'DEFAULT_ARGUMENT_TEMPLATE', 'DEST_FUNCTION',
)

#
# Names of function attributes where Argh stores command behaviour
#

#: explicit command name (differing from function name)
ATTR_NAME = 'argh_name'

#: alternative command names
ATTR_ALIASES = 'argh_aliases'

#: declared arguments
ATTR_ARGS = 'argh_args'

#: list of exception classes that should be wrapped and printed as results
ATTR_WRAPPED_EXCEPTIONS = 'argh_wrap_errors'

#: a function to preprocess the exception object when it is wrapped
ATTR_WRAPPED_EXCEPTIONS_PROCESSOR = 'argh_wrap_errors_processor'

#: forcing argparse.Namespace object instead of signature introspection
ATTR_EXPECTS_NAMESPACE_OBJECT = 'argh_expects_namespace_object'

#
# Dest names in parser defaults
#

#: dest name for a function mapped to given endpoint (goes to Namespace obj)
DEST_FUNCTION = 'function'

#
# Other library-wide stuff
#

class CustomFormatter(argparse.ArgumentDefaultsHelpFormatter,
                      argparse.RawDescriptionHelpFormatter):
    def _expand_help(self, action):
        """
        This method is copied verbatim from ArgumentDefaultsHelpFormatter with
        a couple of lines added just before the end.  Reason: we need to
        `repr()` default values instead of simply inserting them as is.
        This helps notice, for example, an empty string as the default value;
        moreover, it prevents breaking argparse due to logical quirks inside
        of its formatters.

        Ideally this could be achieved by simply defining
        :attr:`DEFAULT_ARGUMENT_TEMPLATE` as ``{default!r}`` but unfortunately
        argparse only supports the old printf syntax.
        """
        params = dict(vars(action), prog=self._prog)
        for name in list(params):
            if params[name] is argparse.SUPPRESS:
                del params[name]
        for name in list(params):
            if hasattr(params[name], '__name__'):
                params[name] = params[name].__name__
        if params.get('choices') is not None:
            choices_str = ', '.join([str(c) for c in params['choices']])
            params['choices'] = choices_str

        # XXX this is added in Argh vs. argparse.ArgumentDefaultsHelpFormatter
        #     (avoiding empty strings, otherwise Argparse would die with
        #     an IndexError in _format_action)
        #
        if 'default' in params:
            if params['default'] is None:
                params['default'] = '-'
            else:
                params['default'] = repr(params['default'])
        #
        # /

        return self._get_help_string(action) % params

#: Default formatter to be used in implicitly instantiated ArgumentParser.
PARSER_FORMATTER = CustomFormatter

DEFAULT_ARGUMENT_TEMPLATE = '%(default)s'
"""
Default template of argument help message (see issue #64).
The template ``%(default)s`` is used by `argparse` to display the argument's
default value.
"""


# coding: utf-8
#
#  Copyright © 2010—2014 Andrey Mikhaylenko and contributors
#
#  This file is part of Argh.
#
#  Argh is free software under terms of the GNU Lesser
#  General Public License version 3 (LGPLv3) as published by the Free
#  Software Foundation. See the file README.rst for copying conditions.
#
"""
Assembling
~~~~~~~~~~

Functions and classes to properly assemble your commands in a parser.
"""
import argparse
from collections import OrderedDict
import sys
import warnings

from argh.completion import COMPLETION_ENABLED
from argh.constants import (
    ATTR_ALIASES,
    ATTR_ARGS,
    ATTR_NAME,
    ATTR_EXPECTS_NAMESPACE_OBJECT,
    PARSER_FORMATTER,
    DEFAULT_ARGUMENT_TEMPLATE,
    DEST_FUNCTION,
)
from argh.utils import get_subparsers, get_arg_spec
from argh.exceptions import AssemblingError

__all__ = [
    'SUPPORTS_ALIASES',
    'set_default_command',
    'add_commands',
    'add_subcommands',
]

def _check_support_aliases():
    p = argparse.ArgumentParser()
    s = p.add_subparsers()
    try:
        s.add_parser('x', aliases=[])
    except TypeError:
        return False
    else:
        return True

SUPPORTS_ALIASES = _check_support_aliases()
"""
Calculated on load. If `True`, current version of argparse supports
alternative command names (can be set via :func:`~argh.decorators.aliases`).
"""

def _get_args_from_signature(function):
    if getattr(function, ATTR_EXPECTS_NAMESPACE_OBJECT, False):
        return

    spec = get_arg_spec(function)

    defaults = dict(zip(*[reversed(x) for x in (spec.args,
                                                spec.defaults or [])]))
    defaults.update(getattr(spec, 'kwonlydefaults', None) or {})

    kwonly = getattr(spec, 'kwonlyargs', [])

    if sys.version_info < (3,0):
        annotations = {}
    else:
        annotations = dict((k,v) for k,v in function.__annotations__.items()
                           if isinstance(v, str))

    # define the list of conflicting option strings
    # (short forms, i.e. single-character ones)
    chars = [a[0] for a in spec.args + kwonly]
    char_counts = dict((char, chars.count(char)) for char in set(chars))
    conflicting_opts = tuple(char for char in char_counts
                             if 1 < char_counts[char])

    for name in spec.args + kwonly:
        flags = []    # name_or_flags
        akwargs = {}  # keyword arguments for add_argument()

        if name in annotations:
            # help message:  func(a : "b")  ->  add_argument("a", help="b")
            akwargs.update(help=annotations.get(name))

        if name in defaults or name in kwonly:
            if name in defaults:
                akwargs.update(default=defaults.get(name))
            else:
                akwargs.update(required=True)
            flags = ('-{0}'.format(name[0]), '--{0}'.format(name))
            if name.startswith(conflicting_opts):
                # remove short name
                flags = flags[1:]

        else:
            # positional argument
            flags = (name,)

        # cmd(foo_bar)  ->  add_argument('foo-bar')
        flags = tuple(x.replace('_', '-') if x.startswith('-') else x
                      for x in flags)

        yield dict(option_strings=flags, **akwargs)

    if spec.varargs:
        # *args
        yield dict(option_strings=[spec.varargs], nargs='*')

def _guess(kwargs):
    """
    Adds types, actions, etc. to given argument specification.
    For example, ``default=3`` implies ``type=int``.

    :param arg: a :class:`argh.utils.Arg` instance
    """
    guessed = {}

    # Parser actions that accept argument 'type'
    TYPE_AWARE_ACTIONS = 'store', 'append'

    # guess type/action from default value
    value = kwargs.get('default')
    if value is not None:
        if isinstance(value, bool):
            if kwargs.get('action') is None:
                # infer action from default value
                guessed['action'] = 'store_false' if value else 'store_true'
        elif kwargs.get('type') is None:
            # infer type from default value
            # (make sure that action handler supports this keyword)
            if kwargs.get('action', 'store') in TYPE_AWARE_ACTIONS:
                guessed['type'] = type(value)

    # guess type from choices (first item)
    if kwargs.get('choices') and 'type' not in list(guessed) + list(kwargs):
        guessed['type'] = type(kwargs['choices'][0])

    return dict(kwargs, **guessed)

def _is_positional(args, prefix_chars='-'):
    assert args
    if 1 < len(args) or args[0][0].startswith(tuple(prefix_chars)):
        return False
    else:
        return True

def _get_parser_param_kwargs(parser, argspec):
    argspec = argspec.copy()    # parser methods modify source data
    args = argspec['option_strings']

    if _is_positional(args, prefix_chars=parser.prefix_chars):
        get_kwargs = parser._get_positional_kwargs
    else:
        get_kwargs = parser._get_optional_kwargs

    kwargs = get_kwargs(*args, **argspec)

    kwargs['dest'] = kwargs['dest'].replace('-', '_')

    return kwargs

def _get_dest(parser, argspec):
    kwargs = _get_parser_param_kwargs(parser, argspec)
    return kwargs['dest']

def _require_support_for_default_command_with_subparsers():
    if sys.version_info < (3,4):
        raise AssemblingError(
            'Argparse library bundled with this version of Python '
            'does not support combining a default command with nested ones.')

def set_default_command(parser, function):
    """
    Sets default command (i.e. a function) for given parser.

    If `parser.description` is empty and the function has a docstring,
    it is used as the description.

    .. note::

       An attempt to set default command to a parser which already has
       subparsers (e.g. added with :func:`~argh.assembling.add_commands`)
       results in a `AssemblingError`.

    .. note::

       If there are both explicitly declared arguments (e.g. via
       :func:`~argh.decorators.arg`) and ones inferred from the function
       signature (e.g. via :func:`~argh.decorators.command`), declared ones
       will be merged into inferred ones. If an argument does not conform
       function signature, `AssemblingError` is raised.

    .. note::

       If the parser was created with ``add_help=True`` (which is by default),
       option name ``-h`` is silently removed from any argument.

    """
    if parser._subparsers:
        _require_support_for_default_command_with_subparsers()

    spec = get_arg_spec(function)

    declared_args = getattr(function, ATTR_ARGS, [])
    inferred_args = list(_get_args_from_signature(function))

    if inferred_args and declared_args:
        # We've got a mixture of declared and inferred arguments

        # a mapping of "dest" strings to argument declarations.
        #
        # * a "dest" string is a normalized form of argument name, i.e.:
        #
        #     '-f', '--foo' → 'foo'
        #     'foo-bar'     → 'foo_bar'
        #
        # * argument declaration is a dictionary representing an argument;
        #   it is obtained either from _get_args_from_signature() or from
        #   an @arg decorator (as is).
        #
        dests = OrderedDict()

        for argspec in inferred_args:
            dest = _get_parser_param_kwargs(parser, argspec)['dest']
            dests[dest] = argspec

        for declared_kw in declared_args:
            # an argument is declared via decorator
            dest = _get_dest(parser, declared_kw)
            if dest in dests:
                # the argument is already known from function signature
                #
                # now make sure that this declared arg conforms to the function
                # signature and therefore only refines an inferred arg:
                #
                #      @arg('my-foo')    maps to  func(my_foo)
                #      @arg('--my-bar')  maps to  func(my_bar=...)

                # either both arguments are positional or both are optional
                decl_positional = _is_positional(declared_kw['option_strings'])
                infr_positional = _is_positional(dests[dest]['option_strings'])
                if decl_positional != infr_positional:
                    kinds = {True: 'positional', False: 'optional'}
                    raise AssemblingError(
                        '{func}: argument "{dest}" declared as {kind_i} '
                        '(in function signature) and {kind_d} (via decorator)'
                        .format(
                            func=function.__name__,
                            dest=dest,
                            kind_i=kinds[infr_positional],
                            kind_d=kinds[decl_positional],
                        ))

                # merge explicit argument declaration into the inferred one
                # (e.g. `help=...`)
                dests[dest].update(**declared_kw)
            else:
                # the argument is not in function signature
                varkw = getattr(spec, 'varkw', getattr(spec, 'keywords', []))
                if varkw:
                    # function accepts **kwargs; the argument goes into it
                    dests[dest] = declared_kw
                else:
                    # there's no way we can map the argument declaration
                    # to function signature
                    xs = (dests[x]['option_strings'] for x in dests)
                    raise AssemblingError(
                        '{func}: argument {flags} does not fit '
                        'function signature: {sig}'.format(
                            flags=', '.join(declared_kw['option_strings']),
                            func=function.__name__,
                            sig=', '.join('/'.join(x) for x in xs)))

        # pack the modified data back into a list
        inferred_args = dests.values()

    command_args = inferred_args or declared_args

    # add types, actions, etc. (e.g. default=3 implies type=int)
    command_args = [_guess(x) for x in command_args]

    for draft in command_args:
        draft = draft.copy()
        if 'help' not in draft:
            draft.update(help=DEFAULT_ARGUMENT_TEMPLATE)
        dest_or_opt_strings = draft.pop('option_strings')
        if parser.add_help and '-h' in dest_or_opt_strings:
            dest_or_opt_strings = [x for x in dest_or_opt_strings if x != '-h']
        completer = draft.pop('completer', None)
        try:
            action = parser.add_argument(*dest_or_opt_strings, **draft)
            if COMPLETION_ENABLED and completer:
                action.completer = completer
        except Exception as e:
            raise type(e)('{func}: cannot add arg {args}: {msg}'.format(
                args='/'.join(dest_or_opt_strings), func=function.__name__, msg=e))

    if function.__doc__ and not parser.description:
        parser.description = function.__doc__
    parser.set_defaults(**{
        DEST_FUNCTION: function,
    })

def add_commands(parser, functions, namespace=None, namespace_kwargs=None,
                 func_kwargs=None,
                 # deprecated args:
                 title=None, description=None, help=None):
    """
    Adds given functions as commands to given parser.

    :param parser:

        an :class:`argparse.ArgumentParser` instance.

    :param functions:

        a list of functions. A subparser is created for each of them.
        If the function is decorated with :func:`~argh.decorators.arg`, the
        arguments are passed to :class:`argparse.ArgumentParser.add_argument`.
        See also :func:`~argh.dispatching.dispatch` for requirements
        concerning function signatures. The command name is inferred from the
        function name. Note that the underscores in the name are replaced with
        hyphens, i.e. function name "foo_bar" becomes command name "foo-bar".

    :param namespace:

        an optional string representing the group of commands. For example, if
        a command named "hello" is added without the namespace, it will be
        available as "prog.py hello"; if the namespace if specified as "greet",
        then the command will be accessible as "prog.py greet hello". The
        namespace itself is not callable, so "prog.py greet" will fail and only
        display a help message.

    :param func_kwargs:

        a `dict` of keyword arguments to be passed to each nested ArgumentParser
        instance created per command (i.e. per function).  Members of this
        dictionary have the highest priority, so a function's docstring is
        overridden by a `help` in `func_kwargs` (if present).

    :param namespace_kwargs:

        a `dict` of keyword arguments to be passed to the nested ArgumentParser
        instance under given `namespace`.

    Deprecated params that should be moved into `namespace_kwargs`:

    :param title:

        passed to :meth:`argparse.ArgumentParser.add_subparsers` as `title`.

        .. deprecated:: 0.26.0

           Please use `namespace_kwargs` instead.

    :param description:

        passed to :meth:`argparse.ArgumentParser.add_subparsers` as
        `description`.

        .. deprecated:: 0.26.0

           Please use `namespace_kwargs` instead.

    :param help:

        passed to :meth:`argparse.ArgumentParser.add_subparsers` as `help`.

        .. deprecated:: 0.26.0

           Please use `namespace_kwargs` instead.

    .. note::

        This function modifies the parser object. Generally side effects are
        bad practice but we don't seem to have any choice as ArgumentParser is
        pretty opaque.
        You may prefer :class:`~argh.helpers.ArghParser.add_commands` for a bit
        more predictable API.

    .. note::

       An attempt to add commands to a parser which already has a default
       function (e.g. added with :func:`~argh.assembling.set_default_command`)
       results in `AssemblingError`.

    """
    # FIXME "namespace" is a correct name but it clashes with the "namespace"
    # that represents arguments (argparse.Namespace and our ArghNamespace).
    # We should rename the argument here.

    if DEST_FUNCTION in parser._defaults:
        _require_support_for_default_command_with_subparsers()

    namespace_kwargs = namespace_kwargs or {}

    # FIXME remove this by 1.0
    #
    if title:
        warnings.warn('argument `title` is deprecated in add_commands(),'
                      ' use `parser_kwargs` instead', DeprecationWarning)
        namespace_kwargs['description'] = title
    if help:
        warnings.warn('argument `help` is deprecated in add_commands(),'
                      ' use `parser_kwargs` instead', DeprecationWarning)
        namespace_kwargs['help'] = help
    if description:
        warnings.warn('argument `description` is deprecated in add_commands(),'
                      ' use `parser_kwargs` instead', DeprecationWarning)
        namespace_kwargs['description'] = description
    #
    # /

    subparsers_action = get_subparsers(parser, create=True)

    if namespace:
        # Make a nested parser and init a deeper _SubParsersAction under it.

        # Create a named group of commands.  It will be listed along with
        # root-level commands in ``app.py --help``; in that context its `title`
        # can be used as a short description on the right side of its name.
        # Normally `title` is shown above the list of commands
        # in ``app.py my-namespace --help``.
        subsubparser_kw = {
            'help': namespace_kwargs.get('title'),
        }
        subsubparser = subparsers_action.add_parser(namespace, **subsubparser_kw)
        subparsers_action = subsubparser.add_subparsers(**namespace_kwargs)
    else:
        assert not namespace_kwargs, ('`parser_kwargs` only makes sense '
                                      'with `namespace`.')

    for func in functions:
        cmd_name, func_parser_kwargs = _extract_command_meta_from_func(func)

        # override any computed kwargs by manually supplied ones
        if func_kwargs:
            func_parser_kwargs.update(func_kwargs)

        # create and set up the parser for this command
        command_parser = subparsers_action.add_parser(cmd_name, **func_parser_kwargs)
        set_default_command(command_parser, func)

def _extract_command_meta_from_func(func):
    # use explicitly defined name; if none, use function name (a_b → a-b)
    cmd_name = getattr(func, ATTR_NAME,
                       func.__name__.replace('_','-'))

    func_parser_kwargs = {

        # add command help from function's docstring
        'help': func.__doc__,

        # set default formatter
        'formatter_class': PARSER_FORMATTER,

    }

    # try adding aliases for command name
    if SUPPORTS_ALIASES:
        func_parser_kwargs['aliases'] = getattr(func, ATTR_ALIASES, [])

    return cmd_name, func_parser_kwargs

def add_subcommands(parser, namespace, functions, **namespace_kwargs):
    """
    A wrapper for :func:`add_commands`.

    These examples are equivalent::

        add_commands(parser, [get, put], namespace='db',
                     namespace_kwargs={
                         'title': 'database commands',
                         'help': 'CRUD for our silly database'
                     })

        add_subcommands(parser, 'db', [get, put],
                        title='database commands',
                        help='CRUD for our silly database')

    """
    add_commands(parser, functions, namespace=namespace,
                 namespace_kwargs=namespace_kwargs)


# coding: utf-8
#
#  Copyright © 2010—2014 Andrey Mikhaylenko and contributors
#
#  This file is part of Argh.
#
#  Argh is free software under terms of the GNU Lesser
#  General Public License version 3 (LGPLv3) as published by the Free
#  Software Foundation. See the file README.rst for copying conditions.
#
"""
Utilities
~~~~~~~~~
"""
import argparse
import inspect

from argh import compat

def get_subparsers(parser, create=False):
    """
    Returns the :class:`argparse._SubParsersAction` instance for given
    :class:`ArgumentParser` instance as would have been returned by
    :meth:`ArgumentParser.add_subparsers`. The problem with the latter is that
    it only works once and raises an exception on the second attempt, and the
    public API seems to lack a method to get *existing* subparsers.

    :param create:
        If `True`, creates the subparser if it does not exist. Default if
        `False`.

    """
    # note that ArgumentParser._subparsers is *not* what is returned by
    # ArgumentParser.add_subparsers().
    if parser._subparsers:
        actions = [a for a in parser._actions
                   if isinstance(a, argparse._SubParsersAction)]
        assert len(actions) == 1
        return actions[0]
    else:
        if create:
            return parser.add_subparsers()

def get_arg_spec(function):
    """
    Returns argument specification for given function.  Omits special
    arguments of instance methods (`self`) and static methods (usually `cls`
    or something like this).
    """
    while hasattr(function, '__wrapped__'):
        function = function.__wrapped__
    spec = compat.getargspec(function)
    if inspect.ismethod(function):
        spec = spec._replace(args=spec.args[1:])
    return spec


# coding: utf-8
#
#  Copyright © 2010—2014 Andrey Mikhaylenko and contributors
#
#  This file is part of Argh.
#
#  Argh is free software under terms of the GNU Lesser
#  General Public License version 3 (LGPLv3) as published by the Free
#  Software Foundation. See the file README.rst for copying conditions.
#
"""
Command decorators
~~~~~~~~~~~~~~~~~~
"""
from argh.constants import (ATTR_ALIASES, ATTR_ARGS, ATTR_NAME,
                            ATTR_WRAPPED_EXCEPTIONS,
                            ATTR_WRAPPED_EXCEPTIONS_PROCESSOR,
                            ATTR_EXPECTS_NAMESPACE_OBJECT)

__all__ = ['aliases', 'named', 'arg', 'wrap_errors', 'expects_obj']

def named(new_name):
    """
    Sets given string as command name instead of the function name.
    The string is used verbatim without further processing.

    Usage::

        @named('load')
        def do_load_some_stuff_and_keep_the_original_function_name(args):
            ...

    The resulting command will be available only as ``load``.  To add aliases
    without renaming the command, check :func:`aliases`.

    .. versionadded:: 0.19
    """
    def wrapper(func):
        setattr(func, ATTR_NAME, new_name)
        return func
    return wrapper

def aliases(*names):
    """
    Defines alternative command name(s) for given function (along with its
    original name). Usage::

        @aliases('co', 'check')
        def checkout(args):
            ...

    The resulting command will be available as ``checkout``, ``check`` and ``co``.

    .. note::

       This decorator only works with a recent version of argparse (see `Python
       issue 9324`_ and `Python rev 4c0426`_).  Such version ships with
       **Python 3.2+** and may be available in other environments as a separate
       package.  Argh does not issue warnings and simply ignores aliases if
       they are not supported.  See :attr:`~argh.assembling.SUPPORTS_ALIASES`.

       .. _Python issue 9324: http://bugs.python.org/issue9324
       .. _Python rev 4c0426: http://hg.python.org/cpython/rev/4c0426261148/

    .. versionadded:: 0.19
    """
    def wrapper(func):
        setattr(func, ATTR_ALIASES, names)
        return func
    return wrapper

def arg(*args, **kwargs):
    """
    Declares an argument for given function. Does not register the function
    anywhere, nor does it modify the function in any way.

    The signature of the decorator matches that of
    :meth:`argparse.ArgumentParser.add_argument`, only some keywords are not
    required if they can be easily guessed (e.g. you don't have to specify type
    or action when an `int` or `bool` default value is supplied).

    Typical use cases:

    - In combination with :func:`expects_obj` (which is not recommended);
    - in combination with ordinary function signatures to add details that
      cannot be expressed with that syntax (e.g. help message).

    Usage::

        from argh import arg

        @arg('path', help='path to the file to load')
        @arg('--format', choices=['yaml','json'])
        @arg('-v', '--verbosity', choices=range(0,3), default=2)
        def load(path, something=None, format='json', dry_run=False, verbosity=1):
            loaders = {'json': json.load, 'yaml': yaml.load}
            loader = loaders[args.format]
            data = loader(args.path)
            if not args.dry_run:
                if verbosity < 1:
                    print('saving to the database')
                put_to_database(data)

    In this example:

    - `path` declaration is extended with `help`;
    - `format` declaration is extended with `choices`;
    - `dry_run` declaration is not duplicated;
    - `verbosity` is extended with `choices` and the default value is
      overridden.  (If both function signature and `@arg` define a default
      value for an argument, `@arg` wins.)

    .. note::

        It is recommended to avoid using this decorator unless there's no way
        to tune the argument's behaviour or presentation using ordinary
        function signatures.  Readability counts, don't repeat yourself.

    """
    def wrapper(func):
        declared_args = getattr(func, ATTR_ARGS, [])
        # The innermost decorator is called first but appears last in the code.
        # We need to preserve the expected order of positional arguments, so
        # the outermost decorator inserts its value before the innermost's:
        declared_args.insert(0, dict(option_strings=args, **kwargs))
        setattr(func, ATTR_ARGS, declared_args)
        return func
    return wrapper

def wrap_errors(errors=None, processor=None, *args):
    """
    Decorator. Wraps given exceptions into
    :class:`~argh.exceptions.CommandError`. Usage::

        @wrap_errors([AssertionError])
        def foo(x=None, y=None):
            assert x or y, 'x or y must be specified'

    If the assertion fails, its message will be correctly printed and the
    stack hidden. This helps to avoid boilerplate code.

    :param errors:
        A list of exception classes to catch.
    :param processor:
        A callable that expects the exception object and returns a string.
        For example, this renders all wrapped errors in red colour::

            from termcolor import colored

            def failure(err):
                return colored(str(err), 'red')

            @wrap_errors(processor=failure)
            def my_command(...):
                ...

    """

    def wrapper(func):
        if errors:
            setattr(func, ATTR_WRAPPED_EXCEPTIONS, errors)

        if processor:
            setattr(func, ATTR_WRAPPED_EXCEPTIONS_PROCESSOR, processor)

        return func
    return wrapper

def expects_obj(func):
    """
    Marks given function as expecting a namespace object.

    Usage::

        @arg('bar')
        @arg('--quux', default=123)
        @expects_obj
        def foo(args):
            yield args.bar, args.quux

    This is equivalent to::

        def foo(bar, quux=123):
            yield bar, quux

    In most cases you don't need this decorator.
    """
    setattr(func, ATTR_EXPECTS_NAMESPACE_OBJECT, True)
    return func


# coding: utf-8
#
#  Copyright © 2010—2014 Andrey Mikhaylenko and contributors
#
#  This file is part of Argh.
#
#  Argh is free software under terms of the GNU Lesser
#  General Public License version 3 (LGPLv3) as published by the Free
#  Software Foundation. See the file README.rst for copying conditions.
#
"""
Interaction
~~~~~~~~~~~
"""
from argh.compat import text_type
from argh.io import safe_input

__all__ = ['confirm', 'safe_input']

def confirm(action, default=None, skip=False):
    """
    A shortcut for typical confirmation prompt.

    :param action:

        a string describing the action, e.g. "Apply changes". A question mark
        will be appended.

    :param default:

        `bool` or `None`. Determines what happens when user hits :kbd:`Enter`
        without typing in a choice. If `True`, default choice is "yes". If
        `False`, it is "no". If `None` the prompt keeps reappearing until user
        types in a choice (not necessarily acceptable) or until the number of
        iteration reaches the limit. Default is `None`.

    :param skip:

        `bool`; if `True`, no interactive prompt is used and default choice is
        returned (useful for batch mode). Default is `False`.

    Usage::

        def delete(key, silent=False):
            item = db.get(Item, args.key)
            if confirm('Delete '+item.title, default=True, skip=silent):
                item.delete()
                print('Item deleted.')
            else:
                print('Operation cancelled.')

    Returns `None` on `KeyboardInterrupt` event.
    """
    MAX_ITERATIONS = 3
    if skip:
        return default
    else:
        defaults = {
            None: ('y','n'),
            True: ('Y','n'),
            False: ('y','N'),
        }
        y, n = defaults[default]
        prompt = text_type('{action}? ({y}/{n})').format(**locals())
        choice = None
        try:
            if default is None:
                cnt = 1
                while not choice and cnt < MAX_ITERATIONS:
                    choice = safe_input(prompt)
                    cnt += 1
            else:
                choice = safe_input(prompt)
        except KeyboardInterrupt:
            return None
    if choice in ('yes', 'y', 'Y'):
        return True
    if choice in ('no', 'n', 'N'):
        return False
    if default is not None:
        return default
    return None


'''
Rom - the Redis object mapper for Python

Copyright 2013-2016 Josiah Carlson

Released under the LGPL license version 2.1 and version 3 (you can choose
which you'd like to be bound under).
'''

from collections import namedtuple
from datetime import datetime, date, time as dtime
from decimal import Decimal as _Decimal
import json
import warnings
import uuid

import six

from .exceptions import QueryError
from .index import Geofilter, Pattern, Prefix, Suffix, _ts
from .util import (_connect, session, dt2ts, t2ts, _script_load,
    STRING_SORT_KEYGENS, STRING_SORT_KEYGENS_STR)

_skip = None
_skip = set(globals()) - set(['__doc__'])

NUMERIC_TYPES = six.integer_types + (float, _Decimal, datetime, date, dtime)

NOT_NULL = (None, None)
_STRING_SORT_KEYGENS = [ss.__name__ for ss in STRING_SORT_KEYGENS]
ALLOWED_DIST = ('m', 'km', 'mi', 'ft')

def _dict_data_factory(columns):
    _dict = dict
    _zip = zip
    def make(data):
        return _dict(_zip(columns, data))
    return make

def _namedtuple_data_factory(columns):
    # note: named tuples don't like lowerscore prefix attributes
    nt = namedtuple('_'.join(columns), [c.lstrip('_') for c in columns])
    def make(data):
        return nt(*data)
    return make

def _tuple_data_factory(columns):
    def make(data):
        return tuple(data)
    return make

def _list_data_factory(columns):
    def make(data):
        # Lists are the representation we get from the json.decode and/or the
        # list comprehension. :P
        return data
    return make

_LT = six.string_types + (six.binary_type,)

class Query(object):
    '''
    This is a query object. It behaves a lot like other query objects. Every
    operation performed on Query objects returns a new Query object. The old
    Query object *does not* have any updated filters.
    '''
    __slots__ = '_model _filters _order_by _limit _select'.split()
    def __init__(self, model, filters=(), order_by=None, limit=None, select=None):
        self._model = model
        self._filters = filters
        self._order_by = order_by
        self._limit = limit
        self._select = select

    def _check(self, column, value=None, which='order_by'):
        column = column.strip('-').partition(':')[0]
        col = self._model._columns.get(column)
        if not col:
            raise QueryError("Cannot use '%s' clause on a non-existent column %r"%(which, column))

        if which == 'filter' and column not in self._model._index:
            raise QueryError("Cannot use 'filter' clause on a column defined with 'index=False'")

        if which == 'startswith' and column not in self._model._prefix:
            raise QueryError("Cannot use 'startswith' clause on a column defined with 'prefix=False'")

        if which == 'like' and column not in self._model._prefix:
            raise QueryError("Cannot use 'like' clause on a column defined with 'prefix=False'")

        if which == 'endswith' and column not in self._model._suffix:
            raise QueryError("Cannot use 'endswith' clause on a column defined with 'suffix=False'")

        if value is not None:
            if isinstance(value, bool):
                value = str(bool(value))

            if col._keygen.__name__ in ('FULL_TEXT', 'SIMPLE_CI', 'CASE_INSENSITIVE', 'IDENTITY_CI'):
                if isinstance(value, _LT):
                    value = value.lower()
                if isinstance(value, list):
                    value = [v.lower() for v in value]

            return value
        return col

    def select(self, *column_names, **kwargs):
        '''
        Select the provided column names from the model, do not return an entity,
        do not involve the rom session, just get the raw and/or processed column
        data from Redis.

        Keyword-only arguments:

            * *include_pk=False* - whether to include the primary key in the
                returned data (we need to get this in some cases, so we fetch
                it anyway; if you want it, we can return it to you - just be
                careful with the namedtuple option - see the warning below)
            * *decode=True* - whether to take a pass through normal data
                decoding in the model (will not return an entity/model)
            * *ff=_dict_data_factory* - the type of data to return from the
                select after all filters/limits/order_by are applied

        .. warning:: If ``include_pk = True`` and if you don't provide
          the primary key column, it will be appended to your list of columns.

        .. note:: if you want to provide a new factory function for the returned
          data, it must be of the form (below is the actual dict factory
          function)

        ::

            def _dict_data_factory(columns):
                _dict = dict
                _zip = zip
                def make(data):
                    # do whatever you need to turn your tuple of columns plus
                    # your list of data into whatever you want:
                    return _dict(_zip(columns, data))
                return make

        Available factory functions:

            * *``rom.query._dict_data_factory``* - default
            * *``rom.query._list_data_factory``* - lowest overhead, as the
              ``data`` passed in above is a list that you can do anything to
            * *``rom.query._tuple_data_factory``* - when you want tuples instead
            * *``rom.query._namedtuple_data_factory``* - get namedtuples, see
              see warning below

        .. warning:: If you use the ``_namedtuple_data_factory``, and your
          columns include underscore prefixes, they will be stripped. If this
          results in a name collision, you *will* get an exception. If you want
          differerent behavior, write your own 20 line factory function that
          does exactly what you want, and pass it; they are really easy!

        '''
        include_pk = kwargs.pop('include_pk', False)
        decode = kwargs.pop('decode', True)
        ff = kwargs.pop('ff', _dict_data_factory)

        if isinstance(column_names[0], (list, tuple)):
            column_names = column_names[0]

        if not column_names:
            raise QueryError("Must provide at least one column to query for raw data")

        if len(set(column_names)) != len(column_names):
            raise QueryError("Column names must be unique")

        missing = [c for c in column_names if c not in self._model._columns]
        if missing:
            raise QueryError("No such columns known: %r"%(missing,))

        remove_last = False
        if self._model._pkey not in column_names:
            column_names += (self._model._pkey,)
            remove_last = not include_pk

        return self.replace(select=(column_names, decode, remove_last, ff))

    def replace(self, **kwargs):
        '''
        Copy the Query object, optionally replacing the filters, order_by, or
        limit information on the copy. This is mostly an internal detail that
        you can ignore.
        '''
        data = {
            'model': self._model,
            'filters': self._filters,
            'order_by': self._order_by,
            'limit': self._limit,
            'select': self._select,
        }
        data.update(**kwargs)
        return Query(**data)

    def filter(self, **kwargs):
        '''
        Only columns/attributes that have been specified as having an index with
        the ``index=True`` option on the column definition can be filtered with
        this method. Prefix, suffix, and pattern match filters must be provided
        using the ``.startswith()``, ``.endswith()``, and the ``.like()``
        methods on the query object, respectively. Geo location queries should
        be performed using the ``.near()`` method.

        Filters should be of the form::

            # for numeric ranges, use None for open-ended ranges
            attribute=(min, max)

            # you can also query for equality by passing a single number
            attribute=value

            # for string searches, passing a plain string will require that
            # string to be in the index as a literal
            attribute=string

            # to perform an 'or' query on strings, you can pass a list of
            # strings
            attribute=[string1, string2]

        As an example, the following will return entities that have both
        ``hello`` and ``world`` in the ``String`` column ``scol`` and has a
        ``Numeric`` column ``ncol`` with value between 2 and 10 (including the
        endpoints)::

            results = MyModel.query \\
                .filter(scol='hello') \\
                .filter(scol='world') \\
                .filter(ncol=(2, 10)) \\
                .all()

        If you only want to match a single value as part of your range query,
        you can pass an integer, float, or Decimal object by itself, similar
        to the ``Model.get_by()`` method::

            results = MyModel.query \\
                .filter(ncol=5) \\
                .execute()

        .. note:: Trying to use a range query `attribute=(min, max)` on indexed
            string columns won't return any results.
        .. note:: This method only filters columns that have been defined with
            ``index=True``.

        '''
        cur_filters = list(self._filters)
        for attr, value in kwargs.items():
            value = self._check(attr, value, which='filter')

            if isinstance(value, NUMERIC_TYPES):
                # for simple numeric equiality filters
                value = (value, value)

            if isinstance(value, six.string_types):
                cur_filters.append('%s:%s'%(attr, value))

            elif six.PY3 and isinstance(value, bytes):
                cur_filters.append('%s:%s'%(attr, value.decode('latin-1')))

            elif isinstance(value, tuple):
                if value is NOT_NULL:
                    from .columns import OneToOne, ManyToOne
                    ctype = type(self._model._columns[attr])
                    if not issubclass(ctype, (OneToOne, ManyToOne)):
                        raise QueryError("Can only query for non-null column values " \
                            "on OneToOne or ManyToOne columns, %r is of type %r"%(attr, ctype))

                if len(value) != 2:
                    raise QueryError("Numeric ranges require 2 endpoints, you provided %s with %r"%(len(value), value))

                tt = []
                for v in value:
                    if isinstance(v, date):
                        v = dt2ts(v)

                    if isinstance(v, dtime):
                        v = t2ts(v)
                    tt.append(v)

                value = tt

                cur_filters.append((attr, value[0], value[1]))

            elif isinstance(value, list) and value:
                cur_filters.append(['%s:%s'%(attr, _ts(v)) for v in value])

            else:
                raise QueryError("Sorry, we don't know how to filter %r by %r"%(attr, value))
        return self.replace(filters=tuple(cur_filters))

    def startswith(self, **kwargs):
        '''
        When provided with keyword arguments of the form ``col=prefix``, this
        will limit the entities returned to those that have a word with the
        provided prefix in the specified column(s). This requires that the
        ``prefix=True`` option was provided during column definition.

        Usage::

            User.query.startswith(email='user@').execute()

        '''
        new = []
        for k, v in kwargs.items():
            v = self._check(k, v, 'startswith')
            new.append(Prefix(k, v))
        return self.replace(filters=self._filters+tuple(new))

    def endswith(self, **kwargs):
        '''
        When provided with keyword arguments of the form ``col=suffix``, this
        will limit the entities returned to those that have a word with the
        provided suffix in the specified column(s). This requires that the
        ``suffix=True`` option was provided during column definition.

        Usage::

            User.query.endswith(email='@gmail.com').execute()

        '''
        new = []
        for k, v in kwargs.items():
            v = self._check(k, v, 'endswith')
            new.append(Suffix(k, v[::-1]))
        return self.replace(filters=self._filters+tuple(new))

    def like(self, **kwargs):
        '''
        When provided with keyword arguments of the form ``col=pattern``, this
        will limit the entities returned to those that include the provided
        pattern. Note that 'like' queries require that the ``prefix=True``
        option must have been provided as part of the column definition.

        Patterns allow for 4 wildcard characters, whose semantics are as
        follows:

            * *?* - will match 0 or 1 of any character
            * *\** - will match 0 or more of any character
            * *+* - will match 1 or more of any character
            * *!* - will match exactly 1 of any character

        As an example, imagine that you have enabled the required prefix
        matching on your ``User.email`` column. And lets say that you want to
        find everyone with an email address that contains the name 'frank'
        before the ``@`` sign. You can use either of the following patterns
        to discover those users.

            * *\*frank\*@*
            * *\*frank\*@*

        .. note:: Like queries implicitly start at the beginning of strings
          checked, so if you want to match a pattern that doesn't start at
          the beginning of a string, you should prefix it with one of the
          wildcard characters (like ``*`` as we did with the 'frank' pattern).
        '''
        new = []
        for k, v in kwargs.items():
            v = self._check(k, v, 'like')
            new.append(Pattern(k, v))
        return self.replace(filters=self._filters+tuple(new))

    def near(self, name, lon, lat, distance, measure, count=None):
        if name not in self._model._geo:
            raise ValueError("provided index name must be defined as a geo index")

        measure = measure.lower()
        if measure not in ALLOWED_DIST:
            raise ValueError("distance measure must be one of %r"%(ALLOWED_DIST,))

        return self.replace(filters=self._filters + (Geofilter(name, lon, lat, distance, measure, count),))

    def order_by(self, column):
        '''
        When provided with a column name, will sort the results of your query::

            # returns all users, ordered by the created_at column in
            # descending order
            User.query.order_by('-created_at').execute()
        '''
        cname = column.lstrip('-')
        col = self._check(cname)
        if type(col).__name__ in ('String', 'Text', 'Json') and col._keygen.__name__ not in _STRING_SORT_KEYGENS:
            warnings.warn("You are trying to order by a non-numeric column %r. "
                          "Unless you have provided your own keygen or are using "
                          "one of the sortable keygens: (%s), this probably won't "
                          "work the way you expect it to."%(cname, STRING_SORT_KEYGENS_STR),
                          stacklevel=2)

        return self.replace(order_by=column)

    def limit(self, offset, count):
        '''
        Will limit the number of results returned from a query::

            # returns the most recent 25 users
            User.query.order_by('-created_at').limit(0, 25).execute()
        '''
        return self.replace(limit=(offset, count))

    def count(self):
        '''
        Will return the total count of the objects that match the specified
        filters.::

            # counts the number of users created in the last 24 hours
            User.query.filter(created_at=(time.time()-86400, time.time())).count()
        '''
        filters = self._filters
        if self._order_by:
            filters += (self._order_by.lstrip('-'),)
        if not filters:
            # We can actually count entities here...
            size = _connect(self._model).hlen(self._model._namespace + '::')
            limit = self._limit or (0, 2**64)
            size = max(size - max(limit[0], 0), 0)
            return min(size, limit[1])

        return self._model._gindex.count(_connect(self._model), filters)

    def _search(self):
        if not (self._filters or self._order_by):
            raise QueryError("You are missing filter or order criteria")
        limit = () if not self._limit else self._limit
        return self._model._gindex.search(
            _connect(self._model), self._filters, self._order_by, *limit)

    def iter_result(self, timeout=30, pagesize=100, no_hscan=False):
        '''
        Iterate over the results of your query instead of getting them all with
        `.all()`. Will only perform a single query. If you expect that your
        processing will take more than 30 seconds to process 100 items, you
        should pass `timeout` and `pagesize` to reflect an appropriate timeout
        and page size to fetch at once.

        Usage::

            for user in User.query.endswith(email='@gmail.com').iter_result():
                # do something with user
                ...

        '''
        if not self._filters and not self._order_by:
            if self._model._columns[self._model._pkey]._index:
                return self._iter_all_pkey()
            conn = _connect(self._model)
            version = list(map(int, conn.info()['redis_version'].split('.')[:2]))
            if version >= [2,8] and not no_hscan:
                return self._iter_all_hscan()
            return self._iter_all()
        return self._iter_results(timeout, pagesize)

    def _iter_results(self, timeout=30, pagesize=100):
        conn = _connect(self._model)
        limit = self._limit or (0, 2**64)
        start = max(limit[0], 0)
        ns = '%s:'%self._model._namespace
        key = self.cached_result(timeout)

        remaining = limit[1]
        ids = [None]
        i = start
        cols = None
        if self._select:
            cols = self._select[0]
            dcols = json.dumps(cols)
            data_gen = iter(_select_generator(None, self._model, *self._select))
            next(data_gen) # prime the generator

        while ids and remaining > 0:
            # refresh the key
            conn.expire(key, timeout)
            ids = list(map(int, conn.zrange(key, i, i+min(remaining, pagesize)-1)))
            if not ids:
                break

            i += len(ids)
            if cols:
                _ids = json.dumps(ids)
                for data in _json_loads(_get_column_data(conn, [ns], [_ids, dcols])):
                    yield data_gen.send(data)
                    remaining -= 1

            else:
                # No need to fill up memory with paginated items hanging around the
                # session. Remove entities from the session as they come in, if they
                # weren't already in the session.
                isk = set(session.known.keys())
                for ent in self._model.get(ids):
                    if ent._pk not in isk:
                        session.forget(ent)
                    yield ent
                    remaining -= 1

    def _iter_all(self):
        conn = _connect(self._model)
        limit = self._limit or (0, 2**64)
        start = max(limit[0], 0)
        prefix = '%s:'%self._model._namespace
        max_id = int(conn.get('%s%s:'%(prefix, self._model._pkey)) or '0')

        cols = None
        if self._select:
            cols = self._select[0]
            dcols = json.dumps(cols)
            data_gen = iter(_select_generator(None, self._model, *self._select))
            next(data_gen) # prime the generator

        # We could use HSCAN here, except that we may get duplicates
        # as we are iterating. That's not good or expected behavior :/
        remaining = max(limit[1], 0)
        ids = [None]
        i = 1
        while ids and i <= max_id and remaining > 0:
            ids = list(range(i, i + 100))
            i += 100
            if cols:
                _ids = json.dumps(ids)
                for data in _json_loads(_get_column_data(conn, [prefix], [_ids, dcols])):
                    if start:
                        start -= 1
                    elif remaining > 0:
                        remaining -= 1
                        yield data_gen.send(data)

            else:
                isk = set(session.known.keys())
                for ent in self._model.get(ids):
                    if ent._pk not in isk:
                        # Same session comment as from _iter_results()
                        session.forget(ent)
                    if start:
                        start -= 1
                    elif remaining > 0:
                        remaining -= 1
                        yield ent

    def _iter_all_hscan(self):
        conn = _connect(self._model)
        limit = self._limit or (0, 2**64)
        start = max(limit[0], 0)
        ns = self._model._namespace + ':'
        tkey = ns + str(uuid.uuid4())
        pkey = self._model._pkey

        remaining = max(limit[1], 0)
        cursor = 0
        ids = ''
        cols = None
        ids = []
        cols = ''
        if self._select:
            cols = self._select[0]
            data_gen = iter(_select_generator(ids, self._model, *self._select))
            next(data_gen) # prime the generator
        dcols = json.dumps(cols)

        while cursor != '0' and remaining > 0:
            result = _scan_fetch_index_hash(conn, [ns, tkey], [cursor, json.dumps(ids or ''), dcols])
            cursor, data = _json_loads(result)

            del ids[:]
            for mdata in data:
                if cols:
                    # we are doing a special select, so ... do what we were going to do :)
                    d = data_gen.send(mdata)
                    if start:
                        start -= 1
                    elif remaining > 0:
                        remaining -= 1
                        yield d
                    else:
                        break

                else:
                    # Turn the flattened data into a dict so we can instantiate the
                    # result and/or fetch the object from the session.
                    mdata = iter(mdata)
                    mdata = dict(zip(mdata, mdata))
                    id = mdata[pkey]
                    ids.append(int(id))
                    # Try to get the shared entity from the session, even though
                    # we just fetched the data from Redis.
                    ent = session.get(ns + id)
                    if not ent:
                        ent = self._model(_loading=True, **mdata)
                        # Same session comment as from _iter_results()
                        session.forget(ent)
                    else:
                        # do a refresh...
                        ent.__init__(_loading=True, **mdata)

                    if start:
                        start -= 1
                    elif remaining > 0:
                        remaining -= 1
                        yield ent
                    else:
                        break

    def _iter_all_pkey(self):
        conn = _connect(self._model)
        limit = self._limit or (0, 2**64)
        start = max(limit[0], 0)
        remaining = limit[1]
        ns = '%s:'%self._model._namespace
        index = '%s:%s:idx'%(self._model._namespace, self._model._pkey)
        max_id = int((conn.zrevrange(index, 0, 0) or (0,))[0])

        cols = None
        if self._select:
            cols = self._select[0]
            dcols = json.dumps(cols)
            data_gen = iter(_select_generator(None, self._model, *self._select))
            next(data_gen) # prime the generator

        i = 1
        if start:
            # skip over the offset at the beginning
            i = int((conn.zrange(index, start, start) or (0,))[0])

        while i <= max_id and remaining > 0:
            ids = conn.zrangebyscore(index, i, i+99)
            i += 100
            if cols:
                _ids = json.dumps(list(map(int, ids)))
                for data in _json_loads(_get_column_data(conn, [ns], [_ids, dcols])):
                    if remaining > 0:
                        remaining -= 1
                        yield data_gen.send(data)
                    else:
                        break

            else:
                isk = set(session.known.keys())
                for ent in self._model.get(ids):
                    # Same session comment as from _iter_results()
                    if ent._pk not in isk:
                        session.forget(ent)
                    if remaining > 0:
                        remaining -= 1
                        yield ent

    def __iter__(self):
        return self.iter_result()

    def cached_result(self, timeout):
        '''
        This will execute the query, returning the key where a ZSET of your
        results will be stored for pagination, further operations, etc.

        The timeout must be a positive integer number of seconds for which to
        set the expiration time on the key (this is to ensure that any cached
        query results are eventually deleted, unless you make the explicit
        step to use the PERSIST command).

        .. note:: Limit clauses are ignored and not passed.

        Usage::

            ukey = User.query.endswith(email='@gmail.com').cached_result(30)
            for i in xrange(0, conn.zcard(ukey), 100):
                # refresh the expiration
                conn.expire(ukey, 30)
                users = User.get(conn.zrange(ukey, i, i+99))
                ...
        '''
        if not (self._filters or self._order_by):
            raise QueryError("You are missing filter or order criteria")
        timeout = int(timeout)
        if timeout < 1:
            raise QueryError("You must specify a timeout >= 1, you gave %r"%timeout)
        return self._model._gindex.search(
            _connect(self._model), self._filters, self._order_by, timeout=timeout)

    def execute(self):
        '''
        Actually executes the query, returning any entities that match the
        filters, ordered by the specified ordering (if any), limited by any
        earlier limit calls.
        '''
        return list(self)

    def all(self):
        '''
        Alias for ``execute()``.
        '''
        return self.execute()

    def first(self):
        '''
        Returns only the first result from the query, if any.
        '''
        lim = [0, 1]
        if self._limit:
            lim[0] = self._limit[0]
        if not self._filters and not self._order_by:
            for ent in self:
                return ent
            return None
        ids = self.limit(*lim)._search()
        if ids:
            return self._model.get(ids[0])
        return None

    def delete(self, blocksize=100):
        '''
        Will delete the entities that match at the time the query is executed.

        Used like::

            MyModel.query.filter(email=...).delete()
            MyModel.query.endswith(email='@host.com').delete()

        .. warning:: can't be used on models on either side of a ``OneToMany``,
            ``ManyToOne``, or ``OneToOne`` relationship.
        '''

        from .columns import MODELS_REFERENCED
        if not self._model._no_fk or self._model._namespace in MODELS_REFERENCED:
            raise QueryError("Can't delete entities of models with foreign key relationships")

        de = []
        i = 0
        for result in self.iter_result(pagesize=blocksize):
            de.append(result)
            i += 1
            if i >= blocksize:
                session.delete(de) # one round-trip to delete "chunk" items
                del de[:]
                i = 0

        if de:
            session.delete(de)

def _select_generator(lst, model, cols, decode, remove_last, factory):
    from itertools import islice
    final = factory(cols[:-1]) if remove_last else factory(cols)
    if decode:
        inter = final
        if remove_last or factory is not _dict_data_factory:
            inter = _dict_data_factory(cols)
    # Get the primary key, and the number of columns we are returning from the
    # query.
    pki = cols.index(model._pkey)
    wanted = len(cols) - remove_last

    data = yield
    if decode:
        while 1:
            # We know which column is the primary key, so can just access it
            # directly.
            if lst is not None:
                lst.append(int(data[pki]))
            # we need to decode, so use both the intermediate and final factories
            inst = model(_loading=True, _bypass_session_entirely=True, **inter(data))
            data = yield final([getattr(inst, c) for c in islice(cols, wanted)])

    else:
        while 1:
            # We know which column is the primary key, so can just access it
            # directly.
            if lst is not None:
                lst.append(int(data[pki]))
            # Fix for Redis' weird JSON null handling.
            data = yield final([None if c is False else c for c in islice(data, wanted)])

def _json_loads(data):
    if six.PY3 and isinstance(data, six.binary_type):
        data = data.decode('utf-8')
    return json.loads(data)

_get_column_data = _script_load('''
local namespace = KEYS[1]
local ids = cjson.decode(ARGV[1])
local cols = cjson.decode(ARGV[2])

local results = {}
local result
for _, id in ipairs(ids) do
    id = namespace .. id
    if redis.call('EXISTS', id) == 1 then
        result = redis.call('HMGET', id, unpack(cols))
        if #result > 0 then
            table.insert(results, result)
        end
    end
end
return cjson.encode(results)
''')

_scan_fetch_index_hash = _script_load('''
local namespace = KEYS[1]
local hkey = namespace .. ':'
local tkey = KEYS[2]
local has_cols = false
local cols, result

-- Keep track of ids we've already returned. We need to write here because
-- we can't write after the HSCAN call below.
if #ARGV[2] > 2 then
    redis.call('SADD', tkey, unpack(cjson.decode(ARGV[2])))
end
if #ARGV[3] > 2 then
    has_cols = true
    cols = cjson.decode(ARGV[3])
end

-- Make sure the temporary set goes away.
redis.call('EXPIRE', tkey, 30)

local results = {}
local pair = redis.call('HSCAN', hkey, ARGV[1])
local cursor = pair[1]
local contents = pair[2]
while #contents > 0 do
    -- ignore the index contents
    table.remove(contents)
    -- get the id
    local id = table.remove(contents)

    if redis.call('SISMEMBER', tkey, id) == 0 then
        if has_cols then
            result = redis.call('HMGET', namespace .. id, unpack(cols))
        else
            result = redis.call('HGETALL', namespace .. id)
        end
        if #result > 0 then
            table.insert(results, result)
        end
    end
end
return cjson.encode({cursor, results})
''')

__all__ = [k for k, v in globals().items() if getattr(v, '__doc__', None) and k not in _skip]
__all__.sort()