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

from queries import Grok

def print_monthly_views(site, pages, year, month):
    grok = Grok(site)
    for page in pages:
        result = grok.get_views_for_month(page, year, month)
        print result['daily_views']

def main():
    """ main script. """
    from argparse import ArgumentParser
    description = 'Extract traffic statistics of Wikipedia articles.'
    parser = ArgumentParser(description=description)
    parser.add_argument("-l", "--lang",
                        type=str,
                        dest="lang",
                        default="en",
                        required=True,
                        help="Language code for Wikipedia")
    parser.add_argument("-y", "--year",
                        type=int,
                        dest="year",
                        default="en",
                        required=True,
                        help="Year")
    parser.add_argument("-m", "--month",
                        type=int,
                        dest="month",
                        default="en",
                        required=True,
                        help="Month")
    parser.add_argument("page", nargs='*',
                        metavar="PAGE",
                        help='A list of pages')
    args = parser.parse_args()
    print_monthly_views(args.lang, args.page, args.year, args.month)

if __name__ == '__main__':
    main()
<commit_msg>Fix default values of Argument Parser<commit_after># -*- coding: utf-8 -*-

from queries import Grok

def print_monthly_views(site, pages, year, month):
    grok = Grok(site)
    for page in pages:
        result = grok.get_views_for_month(page, year, month)
        print result['daily_views']

def main():
    """ main script. """
    from argparse import ArgumentParser
    description = 'Extract traffic statistics of Wikipedia articles.'
    parser = ArgumentParser(description=description)
    parser.add_argument("-l", "--lang",
                        type=str,
                        dest="lang",
                        default="en",
                        required=False,
                        help="Language code for Wikipedia")
    parser.add_argument("-y", "--year",
                        type=int,
                        dest="year",
                        required=False,
                        help="Year")
    parser.add_argument("-m", "--month",
                        type=int,
                        dest="month",
                        required=False,
                        help="Month")
    parser.add_argument("page", nargs='*',
                        metavar="PAGE",
                        help='A list of pages')
    args = parser.parse_args()
    print_monthly_views(args.lang, args.page, args.year, args.month)

if __name__ == '__main__':
    main()
<|endoftext|>


<commit_before># -*- coding: utf-8 -*-
"""
    pygments.cmdline
    ~~~~~~~~~~~~~~~~

    Command line interface.

    :copyright: Copyright 2006-2014 by the Pygments team, see AUTHORS.
    :license: BSD, see LICENSE for details.
"""

from __future__ import print_function

import sys
import getopt
from textwrap import dedent

from pygments import __version__, highlight
from pygments.util import ClassNotFound, OptionError, docstring_headline, \
    text_type, guess_decode
from pygments.lexers import get_all_lexers, get_lexer_by_name, guess_lexer, \
    get_lexer_for_filename, find_lexer_class, TextLexer
from pygments.formatters.latex import LatexEmbeddedLexer, LatexFormatter
from pygments.formatters import get_all_formatters, get_formatter_by_name, \
    get_formatter_for_filename, find_formatter_class, \
    TerminalFormatter  # pylint:disable-msg=E0611
from pygments.filters import get_all_filters, find_filter_class
from pygments.styles import get_all_styles, get_style_by_name

USAGE = """\
Usage: %s [-l <lexer> | -g] [-F <filter>[:<options>]] [-f <formatter>]
          [-O <options>] [-P <option=value>] [-o <outfile>] [<infile>]

       %s -S <style> -f <formatter> [-a <arg>] [-O <options>] [-P <option=value>]
       %s -L [<which> ...]
       %s -N <filename>
       %s -H <type> <name>
       %s -h | -V

Highlight the input file and write the result to <outfile>.

If no input file is given, use stdin, if -o is not given, use stdout.

<lexer> is a lexer name (query all lexer names with -L). If -l is not
given, the lexer is guessed from the extension of the input file name
(this obviously doesn't work if the input is stdin).  If -g is passed,
attempt to guess the lexer from the file contents, or pass through as
plain text if this fails (this can work for stdin).

Likewise, <formatter> is a formatter name, and will be guessed from
the extension of the output file name. If no output file is given,
the terminal formatter will be used by default.

With the -O option, you can give the lexer and formatter a comma-
separated list of options, e.g. ``-O bg=light,python=cool``.

The -P option adds lexer and formatter options like the -O option, but
you can only give one option per -P. That way, the option value may
contain commas and equals signs, which it can't with -O, e.g.
``-P "heading=Pygments, the Python highlighter".

With the -F option, you can add filters to the token stream, you can
give options in the same way as for -O after a colon (note: there must
not be spaces around the colon).

The -O, -P and -F options can be given multiple times.

With the -S option, print out style definitions for style <style>
for formatter <formatter>. The argument given by -a is formatter
dependent.

The -L option lists lexers, formatters, styles or filters -- set
`which` to the thing you want to list (e.g. "styles"), or omit it to
list everything.

The -N option guesses and prints out a lexer name based solely on
the given filename. It does not take input or highlight anything.
If no specific lexer can be determined "text" is returned.

The -H option prints detailed help for the object <name> of type <type>,
where <type> is one of "lexer", "formatter" or "filter".

The -h option prints this help.
The -V option prints the package version.
"""

def _parse_options(o_strs):
    opts = {}
    if not o_strs:
        return opts
    for o_str in o_strs:
        if not o_str:
            continue
        o_args = o_str.split(',')
        for o_arg in o_args:
            o_arg = o_arg.strip()
            try:
                o_key, o_val = o_arg.split('=', 1)
                o_key = o_key.strip()
                o_val = o_val.strip()
            except ValueError:
                opts[o_arg] = True
            else:
                opts[o_key] = o_val
    return opts

def _parse_filters(f_strs):
    filters = []
    if not f_strs:
        return filters
    for f_str in f_strs:
        if ':' in f_str:
            fname, fopts = f_str.split(':', 1)
            filters.append((fname, _parse_options([fopts])))
        else:
            filters.append((f_str, {}))
    return filters

def _print_help(what, name):
    try:
        if what == 'lexer':
            cls = find_lexer_class(name)
            print("Help on the %s lexer:" % cls.name)
            print(dedent(cls.__doc__))
        elif what == 'formatter':
            cls = find_formatter_class(name)
            print("Help on the %s formatter:" % cls.name)
            print(dedent(cls.__doc__))
        elif what == 'filter':
            cls = find_filter_class(name)
            print("Help on the %s filter:" % name)
            print(dedent(cls.__doc__))
    except AttributeError:
        print("%s not found!" % what, file=sys.stderr)

def _print_list(what):
    if what == 'lexer':
        print()
        print("Lexers:")
        print("~~~~~~~")

        info = []
        for fullname, names, exts, _ in get_all_lexers():
            tup = (', '.join(names)+':', fullname,
                   exts and '(filenames ' + ', '.join(exts) + ')' or '')
            info.append(tup)
        info.sort()
        for i in info:
            print(('* %s\n    %s %s') % i)

    elif what == 'formatter':
        print()
        print("Formatters:")
        print("~~~~~~~~~~~")

        info = []
        for cls in get_all_formatters():
            doc = docstring_headline(cls)
            tup = (', '.join(cls.aliases) + ':', doc, cls.filenames and
                   '(filenames ' + ', '.join(cls.filenames) + ')' or '')
            info.append(tup)
        info.sort()
        for i in info:
            print(('* %s\n    %s %s') % i)

    elif what == 'filter':
        print()
        print("Filters:")
        print("~~~~~~~~")

        for name in get_all_filters():
            cls = find_filter_class(name)
            print("* " + name + ':')
            print("    %s" % docstring_headline(cls))

    elif what == 'style':
        print()
        print("Styles:")
        print("~~~~~~~")

        for name in get_all_styles():
            cls = get_style_by_name(name)
            print("* " + name + ':')
            print("    %s" % docstring_headline(cls))

def _get_termencoding():
    """Return terminal encoding for stdin/stdout.

    Defaults to preferred locale encoding.
    """
    import locale
    defencoding = locale.getpreferredencoding()
    inencoding = getattr(sys.stdin, 'encoding', None) or defencoding
    outencoding = getattr(sys.stdout, 'encoding', None) or defencoding
    return inencoding, outencoding

def main(args=sys.argv):
    """
    Main command line entry point.
    """
    # pylint: disable-msg=R0911,R0912,R0915

    usage = USAGE % ((args[0],) * 6)

    if sys.platform in ['win32', 'cygwin']:
        try:
            # Provide coloring under Windows, if possible
            import colorama
            colorama.init()
        except ImportError:
            pass

    try:
        popts, args = getopt.getopt(args[1:], "l:f:F:o:O:P:LS:a:N:hVHg")
    except getopt.GetoptError:
        print(usage, file=sys.stderr)
        return 2
    opts = {}
    O_opts = []
    P_opts = []
    F_opts = []
    for opt, arg in popts:
        if opt == '-O':
            O_opts.append(arg)
        elif opt == '-P':
            P_opts.append(arg)
        elif opt == '-F':
            F_opts.append(arg)
        opts[opt] = arg

    if not opts and not args:
        print(usage)
        return 0

    if opts.pop('-h', None) is not None:
        print(usage)
        return 0

    if opts.pop('-V', None) is not None:
        print('Pygments version %s, (c) 2006-2014 by Georg Brandl.' % __version__)
        return 0

    # handle ``pygmentize -L``
    L_opt = opts.pop('-L', None)
    if L_opt is not None:
        if opts:
            print(usage, file=sys.stderr)
            return 2

        # print version
        main(['', '-V'])
        if not args:
            args = ['lexer', 'formatter', 'filter', 'style']
        for arg in args:
            _print_list(arg.rstrip('s'))
        return 0

    # handle ``pygmentize -H``
    H_opt = opts.pop('-H', None)
    if H_opt is not None:
        if opts or len(args) != 2:
            print(usage, file=sys.stderr)
            return 2

        what, name = args
        if what not in ('lexer', 'formatter', 'filter'):
            print(usage, file=sys.stderr)
            return 2

        _print_help(what, name)
        return 0

    # parse -O options
    parsed_opts = _parse_options(O_opts)
    opts.pop('-O', None)

    # parse -P options
    for p_opt in P_opts:
        try:
            name, value = p_opt.split('=', 1)
        except ValueError:
            parsed_opts[p_opt] = True
        else:
            parsed_opts[name] = value
    opts.pop('-P', None)

    # handle ``pygmentize -N``
    infn = opts.pop('-N', None)
    if infn is not None:
        try:
            lexer = get_lexer_for_filename(infn, **parsed_opts)
        except ClassNotFound as err:
            lexer = TextLexer()
        except OptionError as err:
            print('Error:', err, file=sys.stderr)
            return 1

        print(lexer.aliases[0])
        return 0

    # handle ``pygmentize -S``
    S_opt = opts.pop('-S', None)
    a_opt = opts.pop('-a', None)
    if S_opt is not None:
        f_opt = opts.pop('-f', None)
        if not f_opt:
            print(usage, file=sys.stderr)
            return 2
        if opts or args:
            print(usage, file=sys.stderr)
            return 2

        try:
            parsed_opts['style'] = S_opt
            fmter = get_formatter_by_name(f_opt, **parsed_opts)
        except ClassNotFound as err:
            print(err, file=sys.stderr)
            return 1

        arg = a_opt or ''
        try:
            print(fmter.get_style_defs(arg))
        except Exception as err:
            print('Error:', err, file=sys.stderr)
            return 1
        return 0

    # if no -S is given, -a is not allowed
    if a_opt is not None:
        print(usage, file=sys.stderr)
        return 2

    # parse -F options
    F_opts = _parse_filters(F_opts)
    opts.pop('-F', None)

    # select formatter
    outfn = opts.pop('-o', None)
    fmter = opts.pop('-f', None)
    if fmter:
        try:
            fmter = get_formatter_by_name(fmter, **parsed_opts)
        except (OptionError, ClassNotFound) as err:
            print('Error:', err, file=sys.stderr)
            return 1

    if outfn:
        if not fmter:
            try:
                fmter = get_formatter_for_filename(outfn, **parsed_opts)
            except (OptionError, ClassNotFound) as err:
                print('Error:', err, file=sys.stderr)
                return 1
        try:
            outfile = open(outfn, 'wb')
        except Exception as err:
            print('Error: cannot open outfile:', err, file=sys.stderr)
            return 1
    else:
        if not fmter:
            fmter = TerminalFormatter(**parsed_opts)
        outfile = sys.stdout

    # select lexer
    lexer = opts.pop('-l', None)
    if lexer:
        try:
            lexer = get_lexer_by_name(lexer, **parsed_opts)
        except (OptionError, ClassNotFound) as err:
            print('Error:', err, file=sys.stderr)
            return 1

    if args:
        if len(args) > 1:
            print(usage, file=sys.stderr)
            return 2

        infn = args[0]
        try:
            with open(infn, 'rb') as infp:
                code = infp.read()
        except Exception as err:
            print('Error: cannot read infile:', err, file=sys.stderr)
            return 1
        if 'encoding' not in parsed_opts:
            code = guess_decode(code)

        if not lexer:
            try:
                lexer = get_lexer_for_filename(infn, code, **parsed_opts)
            except ClassNotFound as err:
                if '-g' in opts:
                    try:
                        lexer = guess_lexer(code, **parsed_opts)
                    except ClassNotFound:
                        lexer = TextLexer(**parsed_opts)
                else:
                    print('Error:', err, file=sys.stderr)
                    return 1
            except OptionError as err:
                print('Error:', err, file=sys.stderr)
                return 1

    else:
        code = sys.stdin.read()
        if not isinstance(code, text_type):
            # Python 2; Python 3's terminal is already fine
            code = code.decode(_get_termencoding()[0])
        if not lexer:
            try:
                lexer = guess_lexer(code, **parsed_opts)
            except ClassNotFound:
                lexer = TextLexer(**parsed_opts)

    # When using the LaTeX formatter and the option `escapeinside` is
    # specified, we need a special lexer which collects escaped text
    # before running the chosen language lexer.
    escapeinside = parsed_opts.get('escapeinside', '')
    if len(escapeinside) == 2 and isinstance(fmter, LatexFormatter):
        left = escapeinside[0]
        right = escapeinside[1]
        lexer = LatexEmbeddedLexer(left, right, lexer)

    # No encoding given? Use latin1 if output file given,
    # stdin/stdout encoding otherwise.
    # (This is a compromise, I'm not too happy with it...)
    if 'encoding' not in parsed_opts and 'outencoding' not in parsed_opts:
        if outfn:
            # encoding pass-through
            fmter.encoding = 'latin1'
        else:
            if sys.version_info < (3,):
                # use terminal encoding; Python 3's terminals already do that
                lexer.encoding, fmter.encoding = _get_termencoding()
    elif not outfn and sys.version_info > (3,):
        # output to terminal with encoding -> use .buffer
        outfile = sys.stdout.buffer

    # ... and do it!
    try:
        # process filters
        for fname, fopts in F_opts:
            lexer.add_filter(fname, **fopts)
        highlight(code, lexer, fmter, outfile)
    except Exception:
        import traceback
        info = traceback.format_exception(*sys.exc_info())
        msg = info[-1].strip()
        if len(info) >= 3:
            # extract relevant file and position info
            msg += '\n   (f%s)' % info[-2].split('\n')[0].strip()[1:]
        print(file=sys.stderr)
        print('*** Error while highlighting:', file=sys.stderr)
        print(msg, file=sys.stderr)
        return 1

    return 0
<commit_msg>Closes #799: read from binary buffer stdin on Python 3 if explicit encoding given<commit_after># -*- coding: utf-8 -*-
"""
    pygments.cmdline
    ~~~~~~~~~~~~~~~~

    Command line interface.

    :copyright: Copyright 2006-2014 by the Pygments team, see AUTHORS.
    :license: BSD, see LICENSE for details.
"""

from __future__ import print_function

import sys
import getopt
from textwrap import dedent

from pygments import __version__, highlight
from pygments.util import ClassNotFound, OptionError, docstring_headline, \
    text_type, guess_decode
from pygments.lexers import get_all_lexers, get_lexer_by_name, guess_lexer, \
    get_lexer_for_filename, find_lexer_class, TextLexer
from pygments.formatters.latex import LatexEmbeddedLexer, LatexFormatter
from pygments.formatters import get_all_formatters, get_formatter_by_name, \
    get_formatter_for_filename, find_formatter_class, \
    TerminalFormatter  # pylint:disable-msg=E0611
from pygments.filters import get_all_filters, find_filter_class
from pygments.styles import get_all_styles, get_style_by_name

USAGE = """\
Usage: %s [-l <lexer> | -g] [-F <filter>[:<options>]] [-f <formatter>]
          [-O <options>] [-P <option=value>] [-o <outfile>] [<infile>]

       %s -S <style> -f <formatter> [-a <arg>] [-O <options>] [-P <option=value>]
       %s -L [<which> ...]
       %s -N <filename>
       %s -H <type> <name>
       %s -h | -V

Highlight the input file and write the result to <outfile>.

If no input file is given, use stdin, if -o is not given, use stdout.

<lexer> is a lexer name (query all lexer names with -L). If -l is not
given, the lexer is guessed from the extension of the input file name
(this obviously doesn't work if the input is stdin).  If -g is passed,
attempt to guess the lexer from the file contents, or pass through as
plain text if this fails (this can work for stdin).

Likewise, <formatter> is a formatter name, and will be guessed from
the extension of the output file name. If no output file is given,
the terminal formatter will be used by default.

With the -O option, you can give the lexer and formatter a comma-
separated list of options, e.g. ``-O bg=light,python=cool``.

The -P option adds lexer and formatter options like the -O option, but
you can only give one option per -P. That way, the option value may
contain commas and equals signs, which it can't with -O, e.g.
``-P "heading=Pygments, the Python highlighter".

With the -F option, you can add filters to the token stream, you can
give options in the same way as for -O after a colon (note: there must
not be spaces around the colon).

The -O, -P and -F options can be given multiple times.

With the -S option, print out style definitions for style <style>
for formatter <formatter>. The argument given by -a is formatter
dependent.

The -L option lists lexers, formatters, styles or filters -- set
`which` to the thing you want to list (e.g. "styles"), or omit it to
list everything.

The -N option guesses and prints out a lexer name based solely on
the given filename. It does not take input or highlight anything.
If no specific lexer can be determined "text" is returned.

The -H option prints detailed help for the object <name> of type <type>,
where <type> is one of "lexer", "formatter" or "filter".

The -h option prints this help.
The -V option prints the package version.
"""

def _parse_options(o_strs):
    opts = {}
    if not o_strs:
        return opts
    for o_str in o_strs:
        if not o_str:
            continue
        o_args = o_str.split(',')
        for o_arg in o_args:
            o_arg = o_arg.strip()
            try:
                o_key, o_val = o_arg.split('=', 1)
                o_key = o_key.strip()
                o_val = o_val.strip()
            except ValueError:
                opts[o_arg] = True
            else:
                opts[o_key] = o_val
    return opts

def _parse_filters(f_strs):
    filters = []
    if not f_strs:
        return filters
    for f_str in f_strs:
        if ':' in f_str:
            fname, fopts = f_str.split(':', 1)
            filters.append((fname, _parse_options([fopts])))
        else:
            filters.append((f_str, {}))
    return filters

def _print_help(what, name):
    try:
        if what == 'lexer':
            cls = find_lexer_class(name)
            print("Help on the %s lexer:" % cls.name)
            print(dedent(cls.__doc__))
        elif what == 'formatter':
            cls = find_formatter_class(name)
            print("Help on the %s formatter:" % cls.name)
            print(dedent(cls.__doc__))
        elif what == 'filter':
            cls = find_filter_class(name)
            print("Help on the %s filter:" % name)
            print(dedent(cls.__doc__))
    except AttributeError:
        print("%s not found!" % what, file=sys.stderr)

def _print_list(what):
    if what == 'lexer':
        print()
        print("Lexers:")
        print("~~~~~~~")

        info = []
        for fullname, names, exts, _ in get_all_lexers():
            tup = (', '.join(names)+':', fullname,
                   exts and '(filenames ' + ', '.join(exts) + ')' or '')
            info.append(tup)
        info.sort()
        for i in info:
            print(('* %s\n    %s %s') % i)

    elif what == 'formatter':
        print()
        print("Formatters:")
        print("~~~~~~~~~~~")

        info = []
        for cls in get_all_formatters():
            doc = docstring_headline(cls)
            tup = (', '.join(cls.aliases) + ':', doc, cls.filenames and
                   '(filenames ' + ', '.join(cls.filenames) + ')' or '')
            info.append(tup)
        info.sort()
        for i in info:
            print(('* %s\n    %s %s') % i)

    elif what == 'filter':
        print()
        print("Filters:")
        print("~~~~~~~~")

        for name in get_all_filters():
            cls = find_filter_class(name)
            print("* " + name + ':')
            print("    %s" % docstring_headline(cls))

    elif what == 'style':
        print()
        print("Styles:")
        print("~~~~~~~")

        for name in get_all_styles():
            cls = get_style_by_name(name)
            print("* " + name + ':')
            print("    %s" % docstring_headline(cls))

def _get_termencoding():
    """Return terminal encoding for stdin/stdout.

    Defaults to preferred locale encoding.
    """
    import locale
    defencoding = locale.getpreferredencoding()
    inencoding = getattr(sys.stdin, 'encoding', None) or defencoding
    outencoding = getattr(sys.stdout, 'encoding', None) or defencoding
    return inencoding, outencoding

def main(args=sys.argv):
    """
    Main command line entry point.
    """
    # pylint: disable-msg=R0911,R0912,R0915

    usage = USAGE % ((args[0],) * 6)

    if sys.platform in ['win32', 'cygwin']:
        try:
            # Provide coloring under Windows, if possible
            import colorama
            colorama.init()
        except ImportError:
            pass

    try:
        popts, args = getopt.getopt(args[1:], "l:f:F:o:O:P:LS:a:N:hVHg")
    except getopt.GetoptError:
        print(usage, file=sys.stderr)
        return 2
    opts = {}
    O_opts = []
    P_opts = []
    F_opts = []
    for opt, arg in popts:
        if opt == '-O':
            O_opts.append(arg)
        elif opt == '-P':
            P_opts.append(arg)
        elif opt == '-F':
            F_opts.append(arg)
        opts[opt] = arg

    if not opts and not args:
        print(usage)
        return 0

    if opts.pop('-h', None) is not None:
        print(usage)
        return 0

    if opts.pop('-V', None) is not None:
        print('Pygments version %s, (c) 2006-2014 by Georg Brandl.' % __version__)
        return 0

    # handle ``pygmentize -L``
    L_opt = opts.pop('-L', None)
    if L_opt is not None:
        if opts:
            print(usage, file=sys.stderr)
            return 2

        # print version
        main(['', '-V'])
        if not args:
            args = ['lexer', 'formatter', 'filter', 'style']
        for arg in args:
            _print_list(arg.rstrip('s'))
        return 0

    # handle ``pygmentize -H``
    H_opt = opts.pop('-H', None)
    if H_opt is not None:
        if opts or len(args) != 2:
            print(usage, file=sys.stderr)
            return 2

        what, name = args
        if what not in ('lexer', 'formatter', 'filter'):
            print(usage, file=sys.stderr)
            return 2

        _print_help(what, name)
        return 0

    # parse -O options
    parsed_opts = _parse_options(O_opts)
    opts.pop('-O', None)

    # parse -P options
    for p_opt in P_opts:
        try:
            name, value = p_opt.split('=', 1)
        except ValueError:
            parsed_opts[p_opt] = True
        else:
            parsed_opts[name] = value
    opts.pop('-P', None)

    # handle ``pygmentize -N``
    infn = opts.pop('-N', None)
    if infn is not None:
        try:
            lexer = get_lexer_for_filename(infn, **parsed_opts)
        except ClassNotFound as err:
            lexer = TextLexer()
        except OptionError as err:
            print('Error:', err, file=sys.stderr)
            return 1

        print(lexer.aliases[0])
        return 0

    # handle ``pygmentize -S``
    S_opt = opts.pop('-S', None)
    a_opt = opts.pop('-a', None)
    if S_opt is not None:
        f_opt = opts.pop('-f', None)
        if not f_opt:
            print(usage, file=sys.stderr)
            return 2
        if opts or args:
            print(usage, file=sys.stderr)
            return 2

        try:
            parsed_opts['style'] = S_opt
            fmter = get_formatter_by_name(f_opt, **parsed_opts)
        except ClassNotFound as err:
            print(err, file=sys.stderr)
            return 1

        arg = a_opt or ''
        try:
            print(fmter.get_style_defs(arg))
        except Exception as err:
            print('Error:', err, file=sys.stderr)
            return 1
        return 0

    # if no -S is given, -a is not allowed
    if a_opt is not None:
        print(usage, file=sys.stderr)
        return 2

    # parse -F options
    F_opts = _parse_filters(F_opts)
    opts.pop('-F', None)

    # select formatter
    outfn = opts.pop('-o', None)
    fmter = opts.pop('-f', None)
    if fmter:
        try:
            fmter = get_formatter_by_name(fmter, **parsed_opts)
        except (OptionError, ClassNotFound) as err:
            print('Error:', err, file=sys.stderr)
            return 1

    if outfn:
        if not fmter:
            try:
                fmter = get_formatter_for_filename(outfn, **parsed_opts)
            except (OptionError, ClassNotFound) as err:
                print('Error:', err, file=sys.stderr)
                return 1
        try:
            outfile = open(outfn, 'wb')
        except Exception as err:
            print('Error: cannot open outfile:', err, file=sys.stderr)
            return 1
    else:
        if not fmter:
            fmter = TerminalFormatter(**parsed_opts)
        outfile = sys.stdout

    # select lexer
    lexer = opts.pop('-l', None)
    if lexer:
        try:
            lexer = get_lexer_by_name(lexer, **parsed_opts)
        except (OptionError, ClassNotFound) as err:
            print('Error:', err, file=sys.stderr)
            return 1

    if args:
        if len(args) > 1:
            print(usage, file=sys.stderr)
            return 2

        infn = args[0]
        try:
            with open(infn, 'rb') as infp:
                code = infp.read()
        except Exception as err:
            print('Error: cannot read infile:', err, file=sys.stderr)
            return 1
        if 'encoding' not in parsed_opts:
            code = guess_decode(code)

        if not lexer:
            try:
                lexer = get_lexer_for_filename(infn, code, **parsed_opts)
            except ClassNotFound as err:
                if '-g' in opts:
                    try:
                        lexer = guess_lexer(code, **parsed_opts)
                    except ClassNotFound:
                        lexer = TextLexer(**parsed_opts)
                else:
                    print('Error:', err, file=sys.stderr)
                    return 1
            except OptionError as err:
                print('Error:', err, file=sys.stderr)
                return 1

    else:
        if 'encoding' in parsed_opts:
            if sys.version_info > (3,):
                # Python 3: we have to use .buffer
                code = sys.stdin.buffer.read()
            else:
                code = sys.stdin.read()
            # the lexer will do the decoding
        else:
            code = sys.stdin.read()
            if not isinstance(code, text_type):
                # Python 2; Python 3's terminal is already fine
                code = code.decode(_get_termencoding()[0])
        if not lexer:
            try:
                lexer = guess_lexer(code, **parsed_opts)
            except ClassNotFound:
                lexer = TextLexer(**parsed_opts)

    # When using the LaTeX formatter and the option `escapeinside` is
    # specified, we need a special lexer which collects escaped text
    # before running the chosen language lexer.
    escapeinside = parsed_opts.get('escapeinside', '')
    if len(escapeinside) == 2 and isinstance(fmter, LatexFormatter):
        left = escapeinside[0]
        right = escapeinside[1]
        lexer = LatexEmbeddedLexer(left, right, lexer)

    # No encoding given? Use latin1 if output file given,
    # stdin/stdout encoding otherwise.
    # (This is a compromise, I'm not too happy with it...)
    if 'encoding' not in parsed_opts and 'outencoding' not in parsed_opts:
        if outfn:
            # encoding pass-through
            fmter.encoding = 'latin1'
        else:
            if sys.version_info < (3,):
                # use terminal encoding; Python 3's terminals already do that
                lexer.encoding, fmter.encoding = _get_termencoding()
    elif not outfn and sys.version_info > (3,):
        # output to terminal with encoding -> use .buffer
        outfile = sys.stdout.buffer

    # ... and do it!
    try:
        # process filters
        for fname, fopts in F_opts:
            lexer.add_filter(fname, **fopts)
        highlight(code, lexer, fmter, outfile)
    except Exception:
        import traceback
        info = traceback.format_exception(*sys.exc_info())
        msg = info[-1].strip()
        if len(info) >= 3:
            # extract relevant file and position info
            msg += '\n   (f%s)' % info[-2].split('\n')[0].strip()[1:]
        print(file=sys.stderr)
        print('*** Error while highlighting:', file=sys.stderr)
        print(msg, file=sys.stderr)
        return 1

    return 0
<|endoftext|>


<commit_before>"""Pyflakes support."""

from pyflakes import checker

from pylama.context import RunContext
from pylama.lint import LinterV2 as Abstract

m = checker.messages
CODES = {
    m.UnusedImport.message: "W0611",
    m.RedefinedWhileUnused.message: "W0404",
    m.ImportShadowedByLoopVar.message: "W0621",
    m.ImportStarUsed.message: "W0401",
    m.ImportStarUsage.message: "W0401",
    m.UndefinedName.message: "E0602",
    m.DoctestSyntaxError.message: "W0511",
    m.UndefinedExport.message: "E0603",
    m.UndefinedLocal.message: "E0602",
    m.DuplicateArgument.message: "E1122",
    m.LateFutureImport.message: "W0410",
    m.UnusedVariable.message: "W0612",
    m.ReturnOutsideFunction.message: "E0104",
}

class Linter(Abstract):
    """Pyflakes runner."""

    name = "pyflakes"

    def run_check(self, context: RunContext):  # noqa
        """Check code with pyflakes."""
        params = context.get_params("pyflakes")
        builtins = params.get("builtins", "")
        if builtins:
            builtins = builtins.split(",")

        check = checker.Checker(context.ast, context.filename, builtins=builtins)
        for msg in check.messages:
            context.push(
                lnum=msg.lineno,
                col=msg.col + 1,
                text=msg.message % msg.message_args,
                number=CODES.get(msg.message, ""),
                source="pyflakes",
            )

#  pylama:ignore=E501,C0301
<commit_msg>Modify to keep backward compatibility<commit_after>"""Pyflakes support."""

from pyflakes import checker

from pylama.context import RunContext
from pylama.lint import LinterV2 as Abstract

m = checker.messages
CODES = {
    m.UnusedImport.message: "W0611",
    m.RedefinedWhileUnused.message: "W0404",
    m.ImportShadowedByLoopVar.message: "W0621",
    m.ImportStarUsed.message: "W0401",
    m.ImportStarUsage.message: "W0401",
    m.UndefinedName.message: "E0602",
    m.DoctestSyntaxError.message: "W0511",
    m.UndefinedExport.message: "E0603",
    m.UndefinedLocal.message: "E0602",
    m.DuplicateArgument.message: "E1122",
    m.LateFutureImport.message: "W0410",
    m.UnusedVariable.message: "W0612",
    m.ReturnOutsideFunction.message: "E0104",
}

# RedefinedInListComp and ReturnWithArgsInsideGenerator were removed at pyflakes 2.5.0:
#   https://github.com/PyCQA/pyflakes/commit/2246217295dc8cb30ef4a7b9d8dc449ce32e603a
if hasattr(m, "RedefinedInListComp"):
    CODES[m.RedefinedInListComp.message] = "W0621"
if hasattr(m, "ReturnWithArgsInsideGenerator"):
    CODES[m.ReturnWithArgsInsideGenerator.message] = "E0106"

class Linter(Abstract):
    """Pyflakes runner."""

    name = "pyflakes"

    def run_check(self, context: RunContext):  # noqa
        """Check code with pyflakes."""
        params = context.get_params("pyflakes")
        builtins = params.get("builtins", "")
        if builtins:
            builtins = builtins.split(",")

        check = checker.Checker(context.ast, context.filename, builtins=builtins)
        for msg in check.messages:
            context.push(
                lnum=msg.lineno,
                col=msg.col + 1,
                text=msg.message % msg.message_args,
                number=CODES.get(msg.message, ""),
                source="pyflakes",
            )

#  pylama:ignore=E501,C0301
<|endoftext|>


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

"""
.. See the NOTICE file distributed with this work for additional information
   regarding copyright ownership.

   Licensed under the Apache License, Version 2.0 (the "License");
   you may not use this file except in compliance with the License.
   You may obtain a copy of the License at

       http://www.apache.org/licenses/LICENSE-2.0

   Unless required by applicable law or agreed to in writing, software
   distributed under the License is distributed on an "AS IS" BASIS,
   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
   See the License for the specific language governing permissions and
   limitations under the License.
"""

import os
import re
from collections import deque

class fastqreader(object):
    """
    Module for reading single end and paired end FASTQ files
    """

    def __init__(self):
        """
        Initialise the module
        """
        self.fastq1 = ''
        self.fastq2 = ''

        self.f1_file = None
        self.f2_file = None

        self.f1_eof = False
        self.f2_eof = False

        self.f1_output_file = None
        self.f2_output_file = None

        self.output_tag = ''
        self.output_file_count = 0

        self.paired = False

    def openFastQ(self, file1, file2=None):
        """
        Create file handles for reading the FastQ files

        Parameters
        ----------
        file1 : str
            Location of the first FASTQ file
        file2 : str
            Location of a paired end FASTQ file.
        """
        self.fastq1 = file1
        self.f1_file = open(self.fastq1, "r")
        self.f1_eof = False

        if file2 is not None:
            self.fastq2 = file2
            self.f2_file = open(self.fastq2, "r")
            self.f2_eof = False
            self.paired = True

    def closeFastQ(self):
        """
        Close file handles for the FastQ files.
        """
        self.f1_file.close()

        if self.paired is True:
            self.f2_file.close()

    def eof(self, side=1):
        """
        Indicate if the end of the file has been reached

        Parameters
        ----------
        side : int
            1 or 2
        """
        if side == 1:
            return self.f1_eof
        elif side == 2:
            return self.f2_eof

        return "ERROR"

    def next(self, side=1):
        """
        Get the next read element for the specific FastQ file pair

        Parameters
        ----------
        side : int
            1 or 2 to get the element from the relevant end (DEFAULT: 1)

        Returns
        -------
        dict
            id : str
                Sequence ID
            seq : str
                Called sequence
            add : str
                Plus sign
            score : str
                Base call score
        """
        read_id = ''
        read_seq = ''
        read_addition = ''
        read_score = ''

        if side == 1:
            try:
                read_id = self.f1_file.readline()
                read_seq = self.f1_file.readline()
                read_addition = self.f1_file.readline()
                read_score = self.f1_file.readline()
            except IndexError:
                self.f1_eof = True
                return False

        elif side == 2:
            try:
                read_id = self.f2_file.readline()
                read_seq = self.f2_file.readline()
                read_addition = self.f2_file.readline()
                read_score = self.f2_file.readline()
            except IndexError:
                self.f2_eof = True
                return False
        else:
            return 'ERROR'

        return {
            'id': read_id.rstrip(),
            'seq': read_seq,
            'add': read_addition,
            'score': read_score
        }

    def createOutputFiles(self, tag=''):
        """
        Create and open the file handles for the output files

        Parameters
        ----------
        tag : str
            Tag to identify the output files (DEFAULT: '')
        """
        if tag != '' and self.output_tag != tag:
            self.output_tag = tag

        fq1 = self.fastq1.split("/")
        new_suffix = "." + str(self.output_tag) + "_" + str(self.output_file_count) + ".fastq"
        fq1[-1] = re.sub('.fastq$', new_suffix, fq1[-1])
        fq1.insert(-1, "tmp")

        if os.path.isdir("/".join(fq1[0:-1])) is False:
            os.mkdir("/".join(fq1[0:-1]))

        self.f1_output_file = open("/".join(fq1), "w")

        if self.paired is True:
            fq2 = self.fastq2.split("/")
            new_suffix = "." + str(self.output_tag) + "_" + str(self.output_file_count) + ".fastq"
            fq2[-1] = re.sub('.fastq$', new_suffix, fq2[-1])
            fq2.insert(-1, "tmp")
            self.f2_output_file = open("/".join(fq2), "w")

    def writeOutput(self, read, side=1):
        """
        Writer to print the extracted lines

        Parameters
        ----------
        read : dict
            Read is the dictionary object returned from self.next()
        side : int
            The side that the read has coe from (DEFAULT: 1)

        Returns
        -------
        bool
            False if a value other than 1 or 2 is entered for the side.
        """
        line = read["id"] + "\n" + read["seq"] + read["add"] + read["score"]
        if side == 1:
            self.f1_output_file.write(line)
        elif side == 2:
            self.f2_output_file.write(line)
        else:
            return False
        return True

    def closeOutputFiles(self):
        """
        Close the output file handles
        """
        self.f1_output_file.close()

        if self.paired is True:
            self.f2_output_file.close()

    def incrementOutputFiles(self):
        """
        Increment the counter and create new files for splitting the original
        FastQ paired end files.
        """
        self.closeOutputFiles()

        self.output_file_count += 1

        self.createOutputFiles(self.output_tag)
<commit_msg>Added handling for the end of file<commit_after>#!/usr/bin/env python

"""
.. See the NOTICE file distributed with this work for additional information
   regarding copyright ownership.

   Licensed under the Apache License, Version 2.0 (the "License");
   you may not use this file except in compliance with the License.
   You may obtain a copy of the License at

       http://www.apache.org/licenses/LICENSE-2.0

   Unless required by applicable law or agreed to in writing, software
   distributed under the License is distributed on an "AS IS" BASIS,
   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
   See the License for the specific language governing permissions and
   limitations under the License.
"""

import os
import re
from collections import deque

class fastqreader(object):
    """
    Module for reading single end and paired end FASTQ files
    """

    def __init__(self):
        """
        Initialise the module
        """
        self.fastq1 = ''
        self.fastq2 = ''

        self.f1_file = None
        self.f2_file = None

        self.f1_eof = False
        self.f2_eof = False

        self.f1_output_file = None
        self.f2_output_file = None

        self.output_tag = ''
        self.output_file_count = 0

        self.paired = False

    def openFastQ(self, file1, file2=None):
        """
        Create file handles for reading the FastQ files

        Parameters
        ----------
        file1 : str
            Location of the first FASTQ file
        file2 : str
            Location of a paired end FASTQ file.
        """
        self.fastq1 = file1
        self.f1_file = open(self.fastq1, "r")
        self.f1_eof = False

        if file2 is not None:
            self.fastq2 = file2
            self.f2_file = open(self.fastq2, "r")
            self.f2_eof = False
            self.paired = True

    def closeFastQ(self):
        """
        Close file handles for the FastQ files.
        """
        self.f1_file.close()

        if self.paired is True:
            self.f2_file.close()

    def eof(self, side=1):
        """
        Indicate if the end of the file has been reached

        Parameters
        ----------
        side : int
            1 or 2
        """
        if side == 1:
            return self.f1_eof
        elif side == 2:
            return self.f2_eof

        return "ERROR"

    def next(self, side=1):
        """
        Get the next read element for the specific FastQ file pair

        Parameters
        ----------
        side : int
            1 or 2 to get the element from the relevant end (DEFAULT: 1)

        Returns
        -------
        dict
            id : str
                Sequence ID
            seq : str
                Called sequence
            add : str
                Plus sign
            score : str
                Base call score
        """
        read_id = ''
        read_seq = ''
        read_addition = ''
        read_score = ''

        if side == 1:
            try:
                read_id = self.f1_file.readline()
                read_seq = self.f1_file.readline()
                read_addition = self.f1_file.readline()
                read_score = self.f1_file.readline()
                if read_id == "":
                    raise EOFError
            except EOFError:
                self.f1_eof = True
                return False

        elif side == 2:
            try:
                read_id = self.f2_file.readline()
                read_seq = self.f2_file.readline()
                read_addition = self.f2_file.readline()
                read_score = self.f2_file.readline()
                if read_id == "":
                    raise EOFError
            except EOFError:
                self.f2_eof = True
                return False
        else:
            return 'ERROR'

        return {
            'id': read_id.rstrip(),
            'seq': read_seq,
            'add': read_addition,
            'score': read_score
        }

    def createOutputFiles(self, tag=''):
        """
        Create and open the file handles for the output files

        Parameters
        ----------
        tag : str
            Tag to identify the output files (DEFAULT: '')
        """
        if tag != '' and self.output_tag != tag:
            self.output_tag = tag

        fq1 = self.fastq1.split("/")
        new_suffix = "." + str(self.output_tag) + "_" + str(self.output_file_count) + ".fastq"
        fq1[-1] = re.sub('.fastq$', new_suffix, fq1[-1])
        fq1.insert(-1, "tmp")

        if os.path.isdir("/".join(fq1[0:-1])) is False:
            os.mkdir("/".join(fq1[0:-1]))

        self.f1_output_file = open("/".join(fq1), "w")

        if self.paired is True:
            fq2 = self.fastq2.split("/")
            new_suffix = "." + str(self.output_tag) + "_" + str(self.output_file_count) + ".fastq"
            fq2[-1] = re.sub('.fastq$', new_suffix, fq2[-1])
            fq2.insert(-1, "tmp")
            self.f2_output_file = open("/".join(fq2), "w")

    def writeOutput(self, read, side=1):
        """
        Writer to print the extracted lines

        Parameters
        ----------
        read : dict
            Read is the dictionary object returned from self.next()
        side : int
            The side that the read has coe from (DEFAULT: 1)

        Returns
        -------
        bool
            False if a value other than 1 or 2 is entered for the side.
        """
        line = read["id"] + "\n" + read["seq"] + read["add"] + read["score"]
        if side == 1:
            self.f1_output_file.write(line)
        elif side == 2:
            self.f2_output_file.write(line)
        else:
            return False
        return True

    def closeOutputFiles(self):
        """
        Close the output file handles
        """
        self.f1_output_file.close()

        if self.paired is True:
            self.f2_output_file.close()

    def incrementOutputFiles(self):
        """
        Increment the counter and create new files for splitting the original
        FastQ paired end files.
        """
        self.closeOutputFiles()

        self.output_file_count += 1

        self.createOutputFiles(self.output_tag)
<|endoftext|>


<commit_before>import os
import pytest
import re
from subprocess import check_call
from tempfile import mkdtemp
from shutil import rmtree, copyfile
from time import sleep
from platform import system

from genomepy.plugin import init_plugins, activate
from genomepy.utils import cmd_ok
from genomepy.utils import mkdir_p
from genomepy.functions import Genome
from genomepy.plugins.bwa import BwaPlugin
from genomepy.plugins.gmap import GmapPlugin
from genomepy.plugins.minimap2 import Minimap2Plugin
from genomepy.plugins.bowtie2 import Bowtie2Plugin
from genomepy.plugins.hisat2 import Hisat2Plugin
from genomepy.plugins.star import StarPlugin
from genomepy.plugins.blacklist import BlacklistPlugin

@pytest.fixture(scope="module")
def tempdir():
    """Temporary directory."""
    tmpdir = mkdtemp()
    yield tmpdir
    rmtree(tmpdir)

@pytest.fixture(scope="module", params=["unzipped", "bgzipped"])
def genome(request, tempdir):
    """Create a test genome."""
    name = "dm3"  # Use fake name for blacklist test
    fafile = "tests/data/small_genome.fa"
    bgzipped = True if request.param == "bgzipped" else False

    # Input needs to be bgzipped, depending on param
    if os.path.exists(fafile + ".gz"):
        if not bgzipped:
            check_call(["gunzip", fafile + ".gz"])
    elif bgzipped:
        check_call(["bgzip", fafile])

    tmpdir = os.path.join(tempdir, request.param, name)
    mkdir_p(tmpdir)

    if bgzipped:
        fafile += ".gz"

    copyfile(fafile, os.path.join(tmpdir, os.path.basename(fafile)))
    for p in init_plugins():
        activate(p)
    # provide the fixture value
    yield Genome(name, genome_dir=os.path.join(tempdir, request.param))
    if os.path.exists(fafile) and not bgzipped:
        check_call(["bgzip", fafile])

@pytest.fixture(scope="module", params=["no-overwrite", "overwrite"])
def force(request):
    return request.param

def force_test(p, fname, genome, force):
    """check if a plugin file was properly overwritten (or not) depending on force flag"""
    t0 = os.path.getmtime(fname)
    # OSX rounds down getmtime to the second
    if system() != "Linux":
        sleep(1)
    p.after_genome_download(genome, force=force)
    t1 = os.path.getmtime(fname)
    assert t0 != t1 if force else t0 == t1

def test_blacklist(genome, force):
    """Create blacklist."""
    assert os.path.exists(genome.filename)

    force = True if force == "overwrite" else False
    p = BlacklistPlugin()

    p.after_genome_download(genome, force=force)
    fname = re.sub(".fa(.gz)?$", ".blacklist.bed.gz", genome.filename)
    assert os.path.exists(fname)

    force_test(p, fname, genome, force)

def test_bwa(genome, force):
    """Create bwa index."""
    assert os.path.exists(genome.filename)

    force = True if force == "overwrite" else False
    if cmd_ok("bwa"):
        p = BwaPlugin()
        p.after_genome_download(genome, force=force)
        dirname = os.path.dirname(genome.filename)
        index_dir = os.path.join(dirname, "index", "bwa")
        fname = os.path.join(index_dir, "{}.fa.sa".format(genome.name))
        assert os.path.exists(index_dir)
        assert os.path.exists(fname)

        force_test(p, fname, genome, force)

def test_minimap2(genome, force):
    """Create minimap2 index."""
    assert os.path.exists(genome.filename)

    force = True if force == "overwrite" else False
    if cmd_ok("minimap2"):
        p = Minimap2Plugin()
        p.after_genome_download(genome, force=force)
        dirname = os.path.dirname(genome.filename)
        index_dir = os.path.join(dirname, "index", "minimap2")
        fname = os.path.join(index_dir, "{}.mmi".format(genome.name))
        assert os.path.exists(index_dir)
        assert os.path.exists(fname)

        force_test(p, fname, genome, force)

def test_bowtie2(genome, force):
    """Create bbowtie2 index."""
    assert os.path.exists(genome.filename)

    force = True if force == "overwrite" else False
    if cmd_ok("bowtie2"):
        p = Bowtie2Plugin()
        p.after_genome_download(genome, force=force)
        dirname = os.path.dirname(genome.filename)
        index_dir = os.path.join(dirname, "index", "bowtie2")
        fname = os.path.join(index_dir, "{}.1.bt2".format(genome.name))
        assert os.path.exists(index_dir)
        assert os.path.exists(fname)

        force_test(p, fname, genome, force)

def test_hisat2(genome, force):
    """Create hisat2 index."""
    assert os.path.exists(genome.filename)

    force = True if force == "overwrite" else False
    if cmd_ok("hisat2-build"):
        p = Hisat2Plugin()
        p.after_genome_download(genome)
        dirname = os.path.dirname(genome.filename)
        index_dir = os.path.join(dirname, "index", "hisat2")
        fname = os.path.join(index_dir, "{}.1.ht2".format(genome.name))
        assert os.path.exists(index_dir)
        assert os.path.exists(fname)

        force_test(p, fname, genome, force)

def test_star(genome, force):
    """Create star index."""
    assert os.path.exists(genome.filename)

    force = True if force == "overwrite" else False
    if cmd_ok("STAR"):
        p = StarPlugin()
        p.after_genome_download(genome)
        dirname = os.path.dirname(genome.filename)
        index_dir = os.path.join(dirname, "index", "star")
        fname = os.path.join(index_dir, "SA")
        assert os.path.exists(index_dir)
        assert os.path.exists(fname)

        force_test(p, fname, genome, force)

def test_gmap(genome, force):
    """Create gmap index."""
    assert os.path.exists(genome.filename)

    force = True if force == "overwrite" else False
    if cmd_ok("gmap"):
        p = GmapPlugin()
        p.after_genome_download(genome, force=force)
        dirname = os.path.dirname(genome.filename)
        index_dir = os.path.join(dirname, "index", "gmap")
        fname = os.path.join(index_dir, "{}.maps".format(genome.name))
        assert os.path.exists(index_dir)
        assert os.path.exists(fname)

        force_test(p, fname, genome, force)
<commit_msg>update test to search for unzipped bed<commit_after>import os
import pytest
import re
from subprocess import check_call
from tempfile import mkdtemp
from shutil import rmtree, copyfile
from time import sleep
from platform import system

from genomepy.plugin import init_plugins, activate
from genomepy.utils import cmd_ok
from genomepy.utils import mkdir_p
from genomepy.functions import Genome
from genomepy.plugins.bwa import BwaPlugin
from genomepy.plugins.gmap import GmapPlugin
from genomepy.plugins.minimap2 import Minimap2Plugin
from genomepy.plugins.bowtie2 import Bowtie2Plugin
from genomepy.plugins.hisat2 import Hisat2Plugin
from genomepy.plugins.star import StarPlugin
from genomepy.plugins.blacklist import BlacklistPlugin

@pytest.fixture(scope="module")
def tempdir():
    """Temporary directory."""
    tmpdir = mkdtemp()
    yield tmpdir
    rmtree(tmpdir)

@pytest.fixture(scope="module", params=["unzipped", "bgzipped"])
def genome(request, tempdir):
    """Create a test genome."""
    name = "dm3"  # Use fake name for blacklist test
    fafile = "tests/data/small_genome.fa"
    bgzipped = True if request.param == "bgzipped" else False

    # Input needs to be bgzipped, depending on param
    if os.path.exists(fafile + ".gz"):
        if not bgzipped:
            check_call(["gunzip", fafile + ".gz"])
    elif bgzipped:
        check_call(["bgzip", fafile])

    tmpdir = os.path.join(tempdir, request.param, name)
    mkdir_p(tmpdir)

    if bgzipped:
        fafile += ".gz"

    copyfile(fafile, os.path.join(tmpdir, os.path.basename(fafile)))
    for p in init_plugins():
        activate(p)
    # provide the fixture value
    yield Genome(name, genome_dir=os.path.join(tempdir, request.param))
    if os.path.exists(fafile) and not bgzipped:
        check_call(["bgzip", fafile])

@pytest.fixture(scope="module", params=["no-overwrite", "overwrite"])
def force(request):
    return request.param

def force_test(p, fname, genome, force):
    """check if a plugin file was properly overwritten (or not) depending on force flag"""
    t0 = os.path.getmtime(fname)
    # OSX rounds down getmtime to the second
    if system() != "Linux":
        sleep(1)
    p.after_genome_download(genome, force=force)
    t1 = os.path.getmtime(fname)
    assert t0 != t1 if force else t0 == t1

def test_blacklist(genome, force):
    """Create blacklist."""
    assert os.path.exists(genome.filename)

    force = True if force == "overwrite" else False
    p = BlacklistPlugin()

    p.after_genome_download(genome, force=force)
    fname = re.sub(".fa(.gz)?$", ".blacklist.bed", genome.filename)
    assert os.path.exists(fname)

    force_test(p, fname, genome, force)

def test_bwa(genome, force):
    """Create bwa index."""
    assert os.path.exists(genome.filename)

    force = True if force == "overwrite" else False
    if cmd_ok("bwa"):
        p = BwaPlugin()
        p.after_genome_download(genome, force=force)
        dirname = os.path.dirname(genome.filename)
        index_dir = os.path.join(dirname, "index", "bwa")
        fname = os.path.join(index_dir, "{}.fa.sa".format(genome.name))
        assert os.path.exists(index_dir)
        assert os.path.exists(fname)

        force_test(p, fname, genome, force)

def test_minimap2(genome, force):
    """Create minimap2 index."""
    assert os.path.exists(genome.filename)

    force = True if force == "overwrite" else False
    if cmd_ok("minimap2"):
        p = Minimap2Plugin()
        p.after_genome_download(genome, force=force)
        dirname = os.path.dirname(genome.filename)
        index_dir = os.path.join(dirname, "index", "minimap2")
        fname = os.path.join(index_dir, "{}.mmi".format(genome.name))
        assert os.path.exists(index_dir)
        assert os.path.exists(fname)

        force_test(p, fname, genome, force)

def test_bowtie2(genome, force):
    """Create bbowtie2 index."""
    assert os.path.exists(genome.filename)

    force = True if force == "overwrite" else False
    if cmd_ok("bowtie2"):
        p = Bowtie2Plugin()
        p.after_genome_download(genome, force=force)
        dirname = os.path.dirname(genome.filename)
        index_dir = os.path.join(dirname, "index", "bowtie2")
        fname = os.path.join(index_dir, "{}.1.bt2".format(genome.name))
        assert os.path.exists(index_dir)
        assert os.path.exists(fname)

        force_test(p, fname, genome, force)

def test_hisat2(genome, force):
    """Create hisat2 index."""
    assert os.path.exists(genome.filename)

    force = True if force == "overwrite" else False
    if cmd_ok("hisat2-build"):
        p = Hisat2Plugin()
        p.after_genome_download(genome)
        dirname = os.path.dirname(genome.filename)
        index_dir = os.path.join(dirname, "index", "hisat2")
        fname = os.path.join(index_dir, "{}.1.ht2".format(genome.name))
        assert os.path.exists(index_dir)
        assert os.path.exists(fname)

        force_test(p, fname, genome, force)

def test_star(genome, force):
    """Create star index."""
    assert os.path.exists(genome.filename)

    force = True if force == "overwrite" else False
    if cmd_ok("STAR"):
        p = StarPlugin()
        p.after_genome_download(genome)
        dirname = os.path.dirname(genome.filename)
        index_dir = os.path.join(dirname, "index", "star")
        fname = os.path.join(index_dir, "SA")
        assert os.path.exists(index_dir)
        assert os.path.exists(fname)

        force_test(p, fname, genome, force)

def test_gmap(genome, force):
    """Create gmap index."""
    assert os.path.exists(genome.filename)

    force = True if force == "overwrite" else False
    if cmd_ok("gmap"):
        p = GmapPlugin()
        p.after_genome_download(genome, force=force)
        dirname = os.path.dirname(genome.filename)
        index_dir = os.path.join(dirname, "index", "gmap")
        fname = os.path.join(index_dir, "{}.maps".format(genome.name))
        assert os.path.exists(index_dir)
        assert os.path.exists(fname)

        force_test(p, fname, genome, force)
<|endoftext|>


<commit_before>import codecs
import os
import sys
import time
import re

wpull_hook = globals().get('wpull_hook')  # silence code checkers
tries = 0
total_tries = 0
max_gallery_page = None
max_scraps_page = None
max_favorites_page = None
item_dir = os.environ['item_dir']

def print_(*args, **kwargs):
    print(*args, **kwargs)
    sys.stdout.flush()

def accept_url(url_info, record_info, verdict, reasons):
    global max_gallery_page
    global max_scraps_page

    url = url_info['url']

    if verdict:
        match = re.search(r'furaffinity\.net/(\w+)/([^/]+)/(\d+)/', url)

        if match:
            what_type = match.group(1).lower()

            if what_type in ('gallery', 'scraps', 'favorites'):
                num = int(match.group(3))

                if what_type == 'gallery':
                    if max_gallery_page is not None and num > max_gallery_page:
                        print_('Pagination complete for gallery')
                        return False
                elif what_type == 'scraps':
                    if max_scraps_page is not None and num > max_scraps_page:
                        print_('Pagination complete for scraps')
                        return False
                elif what_type == 'favorites':
                    if max_favorites_page is not None and num > max_favorites_page:
                        print_('Pagination complete for favorites')
                        return False
                else:
                    raise Exception('Unknown what type!')

        if '/themes/classic/' in url_info['url']:
            return False
        if 'google-analytics.com' in url_info['url']:
            return False
        if 'google.com/analytics' in url_info['url']:
            return False
        if 'gstatic.com/analytics' in url_info['url']:
            return False
        if re.search(r'/commissions/.*/(add|manage)/$', url_info['url']):
            return False

    if not verdict and 'facdn.net' in url and '/full/' in record_info['referrer']:
        return True

    return verdict

def handle_pre_response(url_info, record_info, response_info):
    return wpull_hook.actions.NORMAL

def handle_response(url_info, record_info, response_info):
    global tries
    global total_tries
    total_tries += 1

    if total_tries > 1000:
        raise Exception('Too many tries in this session!')

    status_code = response_info['status_code']

    if status_code != 404 and status_code >= 400:
        print_('Uh oh!. Sleeping...')
        time.sleep(60)
        tries += 1

        if tries > 5:
            raise Exception('Giving up')
        else:
            return wpull_hook.actions.RETRY

    tries = 0

    return wpull_hook.actions.NORMAL

def handle_error(url_info, record_info, error_info):
    global tries
    global total_tries
    total_tries += 1

    if total_tries > 1000:
        raise Exception('Too many tries in this session!')

    tries += 1
    if tries > 5:
        raise Exception('Giving up')

    return wpull_hook.actions.NORMAL

def get_urls(filename, url_info, document_info):
    urls = []

    if 'furaffinity.net' in url_info['hostname']:
        with open(filename, 'r', encoding='utf8', errors='replace') as file:
            text = file.read(1048576)

        check_ok_content(text)

        if not is_text_404(text):
            with open(os.path.join(item_dir, 'usernames.txt'), 'a', encoding='utf8', errors='replace') as file:
                for username in scrape_usernames(text):
                    file.write(username)
                    file.write('\n')

            url = url_info['url']
            check_pagination(text, url)

        print_('Looking good so far..')

        match = re.search(r'Download', text)

        if match:
            urls.append({
                'url': match.group(1)
            })

    return urls

def check_ok_content(text):
    ok_text_found = (
        'Commission Info' in text or
        'Submission information:' in text or
        'This user cannot be found.' in text or
        'The journal you are trying to find is not in our database.' in text or
        'The submission you are trying to find is not in our database.' in text
    )

    is_404_error_page = is_text_404(text)

    if ok_text_found and not is_404_error_page and '/logout/' not in text:
        print_('Problem detected. Not logged in! Sleeping.')
        time.sleep(60)
        raise Exception('Not logged in!')
    elif ok_text_found and not is_404_error_page and 'Toggle to hide Mature and Adult submissions.' not in text:
        print_('Problem detected. Cannot view adult material! Sleeping.')
        time.sleep(60)
        raise Exception('Cannot view adult material!')

def is_text_404(text):
    is_404_error_page = (
        'This user cannot be found.' in text or
        'This user has voluntarily disabled access to their userpage.' in text or
        'The journal you are trying to find is not in our database.' in text or
        'The submission you are trying to find is not in our database.' in text
    )
    return is_404_error_page

def scrape_usernames(text):
    for match in re.finditer(r'href="/user/([^"]+)"', text):
        username = match.group(1)
        username = username.strip('/')
        yield username

def check_pagination(text, url):
    global max_gallery_page
    global max_scraps_page
    global max_favorites_page

    match = re.search(r'furaffinity\.net/(\w+)/([^/]+)/(\d+)/', url)

    if match:
        # if codecs.encode('kzyuggc_znxrf_gebtqbe_fnq', 'rot_13') not in text:
        #     raise Exception('Could not find pagination form!')

        what_type = match.group(1).lower()

        if what_type in ('gallery', 'scraps', 'favorites') and 'There are no submissions to list' in text:
            num = int(match.group(3))

            if what_type == 'gallery':
                if max_gallery_page is None:
                    max_gallery_page = num
            elif what_type == 'scraps':
                if max_scraps_page is None:
                    max_scraps_page = num
            elif what_type == 'favorites':
                if max_favorites_page is None:
                    max_favorites_page = num
            else:
                raise Exception('Unknown what type!')

def wait_time(seconds, url_info, url_record, response, error):
    if 'facdn.net' in url_info['hostname'] and not error:
        return 0
    else:
        return seconds

# wpull_hook.callbacks.engine_run = engine_run
# wpull_hook.callbacks.resolve_dns = resolve_dns
wpull_hook.callbacks.accept_url = accept_url
# wpull_hook.callbacks.queued_url = queued_url
# wpull_hook.callbacks.dequeued_url = dequeued_url
wpull_hook.callbacks.handle_pre_response = handle_pre_response
wpull_hook.callbacks.handle_response = handle_response
wpull_hook.callbacks.handle_error = handle_error
wpull_hook.callbacks.get_urls = get_urls
wpull_hook.callbacks.wait_time = wait_time
# wpull_hook.callbacks.finish_statistics = finish_statistics
# wpull_hook.callbacks.exit_status = exit_status
wpull_hook.callbacks.version = 3
<commit_msg>Don't download thumbnails from favorites<commit_after>import codecs
import os
import sys
import time
import re

wpull_hook = globals().get('wpull_hook')  # silence code checkers
tries = 0
total_tries = 0
max_gallery_page = None
max_scraps_page = None
max_favorites_page = None
item_dir = os.environ['item_dir']

def print_(*args, **kwargs):
    print(*args, **kwargs)
    sys.stdout.flush()

def accept_url(url_info, record_info, verdict, reasons):
    global max_gallery_page
    global max_scraps_page

    url = url_info['url']

    if verdict:
        match = re.search(r'furaffinity\.net/(\w+)/([^/]+)/(\d+)/', url)

        if match:
            what_type = match.group(1).lower()

            if what_type in ('gallery', 'scraps', 'favorites'):
                num = int(match.group(3))

                if what_type == 'gallery':
                    if max_gallery_page is not None and num > max_gallery_page:
                        print_('Pagination complete for gallery')
                        return False
                elif what_type == 'scraps':
                    if max_scraps_page is not None and num > max_scraps_page:
                        print_('Pagination complete for scraps')
                        return False
                elif what_type == 'favorites':
                    if max_favorites_page is not None and num > max_favorites_page:
                        print_('Pagination complete for favorites')
                        return False
                else:
                    raise Exception('Unknown what type!')

        if '/themes/classic/' in url_info['url']:
            return False
        if 'google-analytics.com' in url_info['url']:
            return False
        if 'google.com/analytics' in url_info['url']:
            return False
        if 'gstatic.com/analytics' in url_info['url']:
            return False
        if re.search(r'/commissions/.*/(add|manage)/$', url_info['url']):
            return False

        if 'facdn.net' in url and 'furaffinity.net/favorites/' in record_info['referrer']:
            return False

    if not verdict and 'facdn.net' in url and '/full/' in record_info['referrer']:
        return True

    return verdict

def handle_pre_response(url_info, record_info, response_info):
    return wpull_hook.actions.NORMAL

def handle_response(url_info, record_info, response_info):
    global tries
    global total_tries
    total_tries += 1

    if total_tries > 1000:
        raise Exception('Too many tries in this session!')

    status_code = response_info['status_code']

    if status_code != 404 and status_code >= 400:
        print_('Uh oh!. Sleeping...')
        time.sleep(60)
        tries += 1

        if tries > 5:
            raise Exception('Giving up')
        else:
            return wpull_hook.actions.RETRY

    tries = 0

    return wpull_hook.actions.NORMAL

def handle_error(url_info, record_info, error_info):
    global tries
    global total_tries
    total_tries += 1

    if total_tries > 1000:
        raise Exception('Too many tries in this session!')

    tries += 1
    if tries > 5:
        raise Exception('Giving up')

    return wpull_hook.actions.NORMAL

def get_urls(filename, url_info, document_info):
    urls = []

    if 'furaffinity.net' in url_info['hostname']:
        with open(filename, 'r', encoding='utf8', errors='replace') as file:
            text = file.read(1048576)

        check_ok_content(text)

        if not is_text_404(text):
            with open(os.path.join(item_dir, 'usernames.txt'), 'a', encoding='utf8', errors='replace') as file:
                for username in scrape_usernames(text):
                    file.write(username)
                    file.write('\n')

            url = url_info['url']
            check_pagination(text, url)

        print_('Looking good so far..')

        match = re.search(r'Download', text)

        if match:
            urls.append({
                'url': match.group(1)
            })

    return urls

def check_ok_content(text):
    ok_text_found = (
        'Commission Info' in text or
        'Submission information:' in text or
        'This user cannot be found.' in text or
        'The journal you are trying to find is not in our database.' in text or
        'The submission you are trying to find is not in our database.' in text
    )

    is_404_error_page = is_text_404(text)

    if ok_text_found and not is_404_error_page and '/logout/' not in text:
        print_('Problem detected. Not logged in! Sleeping.')
        time.sleep(60)
        raise Exception('Not logged in!')
    elif ok_text_found and not is_404_error_page and 'Toggle to hide Mature and Adult submissions.' not in text:
        print_('Problem detected. Cannot view adult material! Sleeping.')
        time.sleep(60)
        raise Exception('Cannot view adult material!')

def is_text_404(text):
    is_404_error_page = (
        'This user cannot be found.' in text or
        'This user has voluntarily disabled access to their userpage.' in text or
        'The journal you are trying to find is not in our database.' in text or
        'The submission you are trying to find is not in our database.' in text
    )
    return is_404_error_page

def scrape_usernames(text):
    for match in re.finditer(r'href="/user/([^"]+)"', text):
        username = match.group(1)
        username = username.strip('/')
        yield username

def check_pagination(text, url):
    global max_gallery_page
    global max_scraps_page
    global max_favorites_page

    match = re.search(r'furaffinity\.net/(\w+)/([^/]+)/(\d+)/', url)

    if match:
        # if codecs.encode('kzyuggc_znxrf_gebtqbe_fnq', 'rot_13') not in text:
        #     raise Exception('Could not find pagination form!')

        what_type = match.group(1).lower()

        if what_type in ('gallery', 'scraps', 'favorites') and 'There are no submissions to list' in text:
            num = int(match.group(3))

            if what_type == 'gallery':
                if max_gallery_page is None:
                    max_gallery_page = num
            elif what_type == 'scraps':
                if max_scraps_page is None:
                    max_scraps_page = num
            elif what_type == 'favorites':
                if max_favorites_page is None:
                    max_favorites_page = num
            else:
                raise Exception('Unknown what type!')

def wait_time(seconds, url_info, url_record, response, error):
    if 'facdn.net' in url_info['hostname'] and not error:
        return 0
    else:
        return seconds

# wpull_hook.callbacks.engine_run = engine_run
# wpull_hook.callbacks.resolve_dns = resolve_dns
wpull_hook.callbacks.accept_url = accept_url
# wpull_hook.callbacks.queued_url = queued_url
# wpull_hook.callbacks.dequeued_url = dequeued_url
wpull_hook.callbacks.handle_pre_response = handle_pre_response
wpull_hook.callbacks.handle_response = handle_response
wpull_hook.callbacks.handle_error = handle_error
wpull_hook.callbacks.get_urls = get_urls
wpull_hook.callbacks.wait_time = wait_time
# wpull_hook.callbacks.finish_statistics = finish_statistics
# wpull_hook.callbacks.exit_status = exit_status
wpull_hook.callbacks.version = 3
<|endoftext|>