<commit_before>'''
This module contains the main class of the pyPaSWAS suite.

An effort is made to comply as much as possible with the options and arguments
as used by NCBI blastall version 2.2.21.
'''

from pyPaSWAS import parse_cli, set_logger, normalize_file_path
from pyPaSWAS.Core import resource_filename
from pyPaSWAS.Core.Exceptions import InvalidOptionException, ReaderException
from pyPaSWAS.Core.Formatters import DefaultFormatter, SamFormatter,TrimmerFormatter, PlotterFormatter, FASTA
from pyPaSWAS.Core.Programs import Aligner,Trimmer, ComBaRMapper,  ComBaRIndexer, GenomePlotter, Palindrome
from pyPaSWAS.Core.Readers import BioPythonReader
from pyPaSWAS.Core.Scores import BasicScore, CustomScore, DnaRnaScore, Blosum62Score, Blosum80Score, IrysScore, PalindromeScore
from pyPaSWAS.Core.HitList import HitList
import logging
import os.path

class Pypaswas(object):
    '''
    This class represents the main program. It parses the command line and runs
    one of the programs from the pyPaSWAS suite.
    Settings and arguments are stored in two variables created in parse_cli.

    An effort is made to comply as much as possible with the options and arguments
    as used by NCBI blastall version 2.2.21.
    '''
    def __init__(self, config=None):
        self.outputfile = None
        self.score = None
        self.output_format = None
        # set initial logger. Will be modified after settings have been read.
        self.logger = None
        self.program = None
        # Default settings are stored in:
        if config == None:
            self.config_file = resource_filename(__name__, '/Core/cfg/defaults.cfg')
        self._get_default_logger()
        self.settings = None
        self.arguments = None

    def _get_default_logger(self):
        '''Sets a default logger which will be changed after the settings have been parsed.'''
        self.logger = logging.getLogger()
        self.logger.setLevel(logging.WARNING)

        # configure logging to console
        console_format = logging.Formatter('%(levelname)s - %(message)s')
        console_handler = logging.StreamHandler()
        console_handler.setLevel(logging.WARNING)
        console_handler.setFormatter(console_format)
        self.logger.addHandler(console_handler)

    def _get_formatter(self, results):
        '''Sets the format of the output file.'''
        formatter = ''
        if self.output_format == 'SAM':
            formatter = SamFormatter(self.logger, results, self.outputfile)
        elif self.output_format == 'GRAPH':
            from pyPaSWAS.Core.GraphFormatter import GraphFormatter
            formatter = GraphFormatter(self.logger, results, self.outputfile, self.settings)
        elif self.output_format == "trimmedFasta":
            formatter = TrimmerFormatter(self.logger, results, self.outputfile)
        elif self.output_format == "FASTA":
            formatter = FASTA(self.logger, results, self.outputfile)
        elif self.output_format == "plot":
            formatter = PlotterFormatter(self.logger, results, self.outputfile)
        else:
            formatter = DefaultFormatter(self.logger, results, self.outputfile)
        return formatter

    def _set_outfile(self):
        '''
        Checks whether or not the required output file can be created.
        An InvalidOptionException is issued if the file exists already.
        '''
        path = self.settings.out_file
        if path in (None, '', '.'):
            path = 'output'
        path = normalize_file_path(path)
        # check write access
        if self.settings.override_output != "T" and os.path.exists(path):
            raise IOError('File {0} already exists.'.format(path))
        else:
            destination_dir = os.path.split(path)[0]
            if os.path.isdir(destination_dir):
                if not os.access(destination_dir, os.W_OK):
                    raise IOError('Cannot write to {0}.'.format(destination_dir))
            else:
                try:
                    os.makedirs(destination_dir)
                except:
                    raise IOError('Cannot create direcory {0}.'.format(destination_dir))
        self.outputfile = path

    def _get_query_sequences(self, queryfile, start =0, end=None):
        queryfile = normalize_file_path(queryfile)
        if not os.path.exists(queryfile):
            raise InvalidOptionException('File {0} does not exist.'.format(queryfile))
        else:
            reader = BioPythonReader(self.logger, queryfile, self.settings.filetype1, self.settings.limit_length)
            reader.read_records(start,end)
            reader.sort_records()
            return reader.get_records()

    def _get_target_sequences(self, databasefile, start =0, end=None):
        '''Gets the target (database) sequences from the database file. '''
        databasefile = normalize_file_path(databasefile)
        if not os.path.exists(databasefile):
            raise InvalidOptionException('File {0} does not exist.'.format(databasefile))
        else:
            reader = BioPythonReader(self.logger, databasefile, self.settings.filetype2, self.settings.limit_length)
            reader.read_records(start, end)

            if self.score.score_type == 'DNA_RNA':
                reader.complement_records()
            elif self.score.score_type == 'PALINDROME':
                reader.complement_records_only()
            elif self.score.score_type == "IRYS":
                reader.reverse_records()
            reader.sort_records()
            return reader.get_records()

    def _set_scoring_matrix(self):
        ''' Instantiate the scoring matrix. For more information refer to the
            documentation of the different score types. '''
        matrix_name = self.settings.matrix_name.upper()
        score = None
        if matrix_name == 'DNA-RNA':
            score = DnaRnaScore(self.logger, self.settings)
        elif matrix_name == 'PALINDROME':
            score = PalindromeScore(self.logger, self.settings)
        elif matrix_name == 'BASIC':
            score = BasicScore(self.logger, self.settings)
        elif matrix_name == 'BLOSUM62':
            score = Blosum62Score(self.logger, self.settings)
        elif matrix_name == 'BLOSUM80':
            score = Blosum80Score(self.logger, self.settings)
        elif matrix_name == 'CUSTOM':
            score = CustomScore(self.logger, self.settings)
        elif matrix_name == "IRYS":
            score = IrysScore(self.logger, self.settings)
        else:
            raise InvalidOptionException(matrix_name + ' is not a valid substitution matrix')
        self.score = score

    def _set_output_format(self):
        '''Determines the output format.
            Currently only TXT for text and SAM for SAM are supported.
        '''
        if self.settings.out_format.upper() == 'SAM':
            self.output_format = 'SAM'
        elif self.settings.out_format.upper() == 'TXT':
            self.output_format = 'TXT'
        elif self.settings.out_format.upper() == 'TRIMMEDFASTA':
            self.output_format = 'trimmedFasta'
        elif self.settings.out_format.upper() == 'PLOT':
            self.output_format = 'plot'
        elif self.settings.out_format.upper() == 'FASTA':
            self.output_format = 'FASTA'
        elif self.settings.out_format.upper() == "GRAPH":
            self.output_format = 'GRAPH'
        else:
            raise InvalidOptionException('Invalid output format {0}.'.format(self.settings.out_format))

    def _set_program(self):
        '''Determines what program from the suite should be used and instantiates it'''
        if self.settings.program == 'aligner':
            self.program = Aligner(self.logger, self.score, self.settings)
        elif self.settings.program == 'trimmer':
            self.program = Trimmer(self.logger, self.score, self.settings)
        elif self.settings.program == 'mapper':
            self.program = ComBaRMapper(self.logger, self.score, self.settings, self.arguments)
            self.logger.warning("Removing limits on length of sequences for ComBaR mapping!")
            self.settings.limit_length = 10**20
            #self.logger.warning("Setting query_step to 1 for ComBaR mapping")
            #self.settings.query_step = 1
        elif self.settings.program == "indexer":
            self.program = ComBaRIndexer(self.logger, self.score, self.settings, self.arguments)
            self.logger.warning("Removing limits on length of sequences for ComBaR mapping!")
            self.settings.limit_length = 10**20
        elif self.settings.program == "plotter":
            self.program = GenomePlotter(self.logger, self.score, self.settings, self.arguments)
            self.logger.warning("Removing limits on length of sequences for genome plotter!")
            self.settings.limit_length = 10**20
        elif self.settings.program == "palindrome":
            self.program = Palindrome(self.logger, self.score, self.settings)
            self.logger.warning("Forcing output to FASTA")
            self.output_format = "FASTA"
            self.logger.warning("Forcing query step to 1")
            self.settings.query_step = "1"
            self.logger.warning("Forcing sequence step to 1")
            self.settings.sequence_step = "1"
            self.logger.warning("Forcing Matrix to PALINDROME")
            self.settings.matrix_name = "PALINDROME"
            self.score = PalindromeScore(self.logger, self.settings)
        else:
            raise InvalidOptionException('Invalid program selected {0}'.format(self.settings.program))

    def run(self):
        '''The main program of pyPaSWAS.'''
        # Read command-line arguments
        self.settings, self.arguments = parse_cli(self.config_file)
        self.logger = set_logger(self.settings)
        self.logger.info("Initializing application...")
        self._set_outfile()
        self._set_scoring_matrix()
        self.logger.info('Application initialized.')
        self.logger.info('Setting program...')
        self._set_output_format()
        self._set_program()
        self.logger.info('Program set.')

        queriesToProcess = True

        query_start = int(self.settings.start_query)
        query_end = int(self.settings.start_query) + int(self.settings.query_step)
        if query_end > int(self.settings.end_query) and int(self.settings.start_query) != int(self.settings.end_query):
            query_end = int(self.settings.end_query)

        start_index = int(self.settings.start_target)

        end_index = int(self.settings.start_target) + int(self.settings.sequence_step)
        if end_index > int(self.settings.end_target) and int(self.settings.start_target) != int(self.settings.end_target):
            end_index = int(self.settings.end_target)

        results = HitList(self.logger)

        while queriesToProcess:
            self.logger.info('Reading query sequences {} {}...'.format(query_start, query_end))
            try:
                query_sequences = self._get_query_sequences(self.arguments[0], start=query_start, end=query_end)
                self.logger.info('Query sequences OK.')
            except ReaderException:
                queriesToProcess = False

            sequencesToProcess = True
            if not self.settings.program == "palindrome":
                start_index = int(self.settings.start_target)
                end_index = int(self.settings.start_target) + int(self.settings.sequence_step)
                if end_index > int(self.settings.end_target) and int(self.settings.start_target) != int(self.settings.end_target):
                    end_index = int(self.settings.end_target)

            while queriesToProcess and sequencesToProcess:

                self.logger.info('Reading target sequences {}, {}...'.format(start_index,end_index))
                try:
                    target_sequences = self._get_target_sequences(self.arguments[1], start=start_index, end=end_index)
                    self.logger.info('Target sequences OK.')
                except ReaderException:
                    sequencesToProcess = False

                if not sequencesToProcess or not queriesToProcess or len(query_sequences) == 0 or len(target_sequences) == 0:
                    sequencesToProcess = False
                    self.logger.info('Processing done')
                else:
                    self.logger.info('Processing {0}- vs {1}-sequences'.format(len(query_sequences),
                                                                            len(target_sequences)))
                    results.extend(self.program.process(query_sequences, target_sequences, self))

                if sequencesToProcess and len(target_sequences) <= end_index:
                    # for palindrome program, skip directly to next
                    sequencesToProcess = False
                    self.logger.info('Processing done')

                start_index = start_index + int(self.settings.sequence_step)
                end_index = end_index + int(self.settings.sequence_step)
                if  self.settings.program == "palindrome" or (int(self.settings.end_target) > 0 and int(self.settings.end_target) < end_index):
                    sequencesToProcess = False

            if int(self.settings.end_query) > 0 and int(self.settings.end_query) < query_end:
                queriesToProcess = False
            query_start = query_start + int(self.settings.query_step)
            query_end = query_end + int(self.settings.query_step)

        nhits = len(results.hits)
        # retrieve and print results!
        self.logger.info('Processing OK ({} hits found).'.format(nhits))
        if nhits > 0:
            self.logger.info('Formatting output...')
            formatter = self._get_formatter(results)
            self.logger.info('Formatting OK.')

            self.logger.info('Writing output...')
            formatter.print_results()
            self.logger.info('Writing OK.')
            self.logger.info('Finished')
        else:
            self.logger.warning('No suitable hits produced, exiting...')

if __name__ == '__main__':
    try:
        ppw = Pypaswas()
        ppw.run()
    except Exception as exception:
        # Show complete exception when running in DEBUG
        if (hasattr(ppw.settings, 'loglevel') and
            getattr(logging, 'DEBUG') == ppw.logger.getEffectiveLevel()):
            ppw.logger.exception(str(exception))
        else:
            print 'Something unforeseen caused the program to terminate.\n' \
                  'The internal error message was: ', ','.join(exception.args)
<commit_msg>config file bug fix<commit_after>'''
This module contains the main class of the pyPaSWAS suite.

An effort is made to comply as much as possible with the options and arguments
as used by NCBI blastall version 2.2.21.
'''

from pyPaSWAS import parse_cli, set_logger, normalize_file_path
from pyPaSWAS.Core import resource_filename
from pyPaSWAS.Core.Exceptions import InvalidOptionException, ReaderException
from pyPaSWAS.Core.Formatters import DefaultFormatter, SamFormatter,TrimmerFormatter, PlotterFormatter, FASTA
from pyPaSWAS.Core.Programs import Aligner,Trimmer, ComBaRMapper,  ComBaRIndexer, GenomePlotter, Palindrome
from pyPaSWAS.Core.Readers import BioPythonReader
from pyPaSWAS.Core.Scores import BasicScore, CustomScore, DnaRnaScore, Blosum62Score, Blosum80Score, IrysScore, PalindromeScore
from pyPaSWAS.Core.HitList import HitList
import logging
import os.path

class Pypaswas(object):
    '''
    This class represents the main program. It parses the command line and runs
    one of the programs from the pyPaSWAS suite.
    Settings and arguments are stored in two variables created in parse_cli.

    An effort is made to comply as much as possible with the options and arguments
    as used by NCBI blastall version 2.2.21.
    '''
    def __init__(self, config=None):
        self.outputfile = None
        self.score = None
        self.output_format = None
        # set initial logger. Will be modified after settings have been read.
        self.logger = None
        self.program = None
        # Default settings are stored in:
        if config == None:
            self.config_file = resource_filename(__name__, '/Core/cfg/defaults.cfg')
        else:
            self.config_file = config

        self._get_default_logger()
        self.settings = None
        self.arguments = None

    def _get_default_logger(self):
        '''Sets a default logger which will be changed after the settings have been parsed.'''
        self.logger = logging.getLogger()
        self.logger.setLevel(logging.WARNING)

        # configure logging to console
        console_format = logging.Formatter('%(levelname)s - %(message)s')
        console_handler = logging.StreamHandler()
        console_handler.setLevel(logging.WARNING)
        console_handler.setFormatter(console_format)
        self.logger.addHandler(console_handler)

    def _get_formatter(self, results):
        '''Sets the format of the output file.'''
        formatter = ''
        if self.output_format == 'SAM':
            formatter = SamFormatter(self.logger, results, self.outputfile)
        elif self.output_format == 'GRAPH':
            from pyPaSWAS.Core.GraphFormatter import GraphFormatter
            formatter = GraphFormatter(self.logger, results, self.outputfile, self.settings)
        elif self.output_format == "trimmedFasta":
            formatter = TrimmerFormatter(self.logger, results, self.outputfile)
        elif self.output_format == "FASTA":
            formatter = FASTA(self.logger, results, self.outputfile)
        elif self.output_format == "plot":
            formatter = PlotterFormatter(self.logger, results, self.outputfile)
        else:
            formatter = DefaultFormatter(self.logger, results, self.outputfile)
        return formatter

    def _set_outfile(self):
        '''
        Checks whether or not the required output file can be created.
        An InvalidOptionException is issued if the file exists already.
        '''
        path = self.settings.out_file
        if path in (None, '', '.'):
            path = 'output'
        path = normalize_file_path(path)
        # check write access
        if self.settings.override_output != "T" and os.path.exists(path):
            raise IOError('File {0} already exists.'.format(path))
        else:
            destination_dir = os.path.split(path)[0]
            if os.path.isdir(destination_dir):
                if not os.access(destination_dir, os.W_OK):
                    raise IOError('Cannot write to {0}.'.format(destination_dir))
            else:
                try:
                    os.makedirs(destination_dir)
                except:
                    raise IOError('Cannot create direcory {0}.'.format(destination_dir))
        self.outputfile = path

    def _get_query_sequences(self, queryfile, start =0, end=None):
        queryfile = normalize_file_path(queryfile)
        if not os.path.exists(queryfile):
            raise InvalidOptionException('File {0} does not exist.'.format(queryfile))
        else:
            reader = BioPythonReader(self.logger, queryfile, self.settings.filetype1, self.settings.limit_length)
            reader.read_records(start,end)
            reader.sort_records()
            return reader.get_records()

    def _get_target_sequences(self, databasefile, start =0, end=None):
        '''Gets the target (database) sequences from the database file. '''
        databasefile = normalize_file_path(databasefile)
        if not os.path.exists(databasefile):
            raise InvalidOptionException('File {0} does not exist.'.format(databasefile))
        else:
            reader = BioPythonReader(self.logger, databasefile, self.settings.filetype2, self.settings.limit_length)
            reader.read_records(start, end)

            if self.score.score_type == 'DNA_RNA':
                reader.complement_records()
            elif self.score.score_type == 'PALINDROME':
                reader.complement_records_only()
            elif self.score.score_type == "IRYS":
                reader.reverse_records()
            reader.sort_records()
            return reader.get_records()

    def _set_scoring_matrix(self):
        ''' Instantiate the scoring matrix. For more information refer to the
            documentation of the different score types. '''
        matrix_name = self.settings.matrix_name.upper()
        score = None
        if matrix_name == 'DNA-RNA':
            score = DnaRnaScore(self.logger, self.settings)
        elif matrix_name == 'PALINDROME':
            score = PalindromeScore(self.logger, self.settings)
        elif matrix_name == 'BASIC':
            score = BasicScore(self.logger, self.settings)
        elif matrix_name == 'BLOSUM62':
            score = Blosum62Score(self.logger, self.settings)
        elif matrix_name == 'BLOSUM80':
            score = Blosum80Score(self.logger, self.settings)
        elif matrix_name == 'CUSTOM':
            score = CustomScore(self.logger, self.settings)
        elif matrix_name == "IRYS":
            score = IrysScore(self.logger, self.settings)
        else:
            raise InvalidOptionException(matrix_name + ' is not a valid substitution matrix')
        self.score = score

    def _set_output_format(self):
        '''Determines the output format.
            Currently only TXT for text and SAM for SAM are supported.
        '''
        if self.settings.out_format.upper() == 'SAM':
            self.output_format = 'SAM'
        elif self.settings.out_format.upper() == 'TXT':
            self.output_format = 'TXT'
        elif self.settings.out_format.upper() == 'TRIMMEDFASTA':
            self.output_format = 'trimmedFasta'
        elif self.settings.out_format.upper() == 'PLOT':
            self.output_format = 'plot'
        elif self.settings.out_format.upper() == 'FASTA':
            self.output_format = 'FASTA'
        elif self.settings.out_format.upper() == "GRAPH":
            self.output_format = 'GRAPH'
        else:
            raise InvalidOptionException('Invalid output format {0}.'.format(self.settings.out_format))

    def _set_program(self):
        '''Determines what program from the suite should be used and instantiates it'''
        if self.settings.program == 'aligner':
            self.program = Aligner(self.logger, self.score, self.settings)
        elif self.settings.program == 'trimmer':
            self.program = Trimmer(self.logger, self.score, self.settings)
        elif self.settings.program == 'mapper':
            self.program = ComBaRMapper(self.logger, self.score, self.settings, self.arguments)
            self.logger.warning("Removing limits on length of sequences for ComBaR mapping!")
            self.settings.limit_length = 10**20
            #self.logger.warning("Setting query_step to 1 for ComBaR mapping")
            #self.settings.query_step = 1
        elif self.settings.program == "indexer":
            self.program = ComBaRIndexer(self.logger, self.score, self.settings, self.arguments)
            self.logger.warning("Removing limits on length of sequences for ComBaR mapping!")
            self.settings.limit_length = 10**20
        elif self.settings.program == "plotter":
            self.program = GenomePlotter(self.logger, self.score, self.settings, self.arguments)
            self.logger.warning("Removing limits on length of sequences for genome plotter!")
            self.settings.limit_length = 10**20
        elif self.settings.program == "palindrome":
            self.program = Palindrome(self.logger, self.score, self.settings)
            self.logger.warning("Forcing output to FASTA")
            self.output_format = "FASTA"
            self.logger.warning("Forcing query step to 1")
            self.settings.query_step = "1"
            self.logger.warning("Forcing sequence step to 1")
            self.settings.sequence_step = "1"
            self.logger.warning("Forcing Matrix to PALINDROME")
            self.settings.matrix_name = "PALINDROME"
            self.score = PalindromeScore(self.logger, self.settings)
        else:
            raise InvalidOptionException('Invalid program selected {0}'.format(self.settings.program))

    def run(self):
        '''The main program of pyPaSWAS.'''
        # Read command-line arguments
        self.settings, self.arguments = parse_cli(self.config_file)
        self.logger = set_logger(self.settings)
        self.logger.info("Initializing application...")
        self._set_outfile()
        self._set_scoring_matrix()
        self.logger.info('Application initialized.')
        self.logger.info('Setting program...')
        self._set_output_format()
        self._set_program()
        self.logger.info('Program set.')

        queriesToProcess = True

        query_start = int(self.settings.start_query)
        query_end = int(self.settings.start_query) + int(self.settings.query_step)
        if query_end > int(self.settings.end_query) and int(self.settings.start_query) != int(self.settings.end_query):
            query_end = int(self.settings.end_query)

        start_index = int(self.settings.start_target)

        end_index = int(self.settings.start_target) + int(self.settings.sequence_step)
        if end_index > int(self.settings.end_target) and int(self.settings.start_target) != int(self.settings.end_target):
            end_index = int(self.settings.end_target)

        results = HitList(self.logger)

        while queriesToProcess:
            self.logger.info('Reading query sequences {} {}...'.format(query_start, query_end))
            try:
                query_sequences = self._get_query_sequences(self.arguments[0], start=query_start, end=query_end)
                self.logger.info('Query sequences OK.')
            except ReaderException:
                queriesToProcess = False

            sequencesToProcess = True
            if not self.settings.program == "palindrome":
                start_index = int(self.settings.start_target)
                end_index = int(self.settings.start_target) + int(self.settings.sequence_step)
                if end_index > int(self.settings.end_target) and int(self.settings.start_target) != int(self.settings.end_target):
                    end_index = int(self.settings.end_target)

            while queriesToProcess and sequencesToProcess:

                self.logger.info('Reading target sequences {}, {}...'.format(start_index,end_index))
                try:
                    target_sequences = self._get_target_sequences(self.arguments[1], start=start_index, end=end_index)
                    self.logger.info('Target sequences OK.')
                except ReaderException:
                    sequencesToProcess = False

                if not sequencesToProcess or not queriesToProcess or len(query_sequences) == 0 or len(target_sequences) == 0:
                    sequencesToProcess = False
                    self.logger.info('Processing done')
                else:
                    self.logger.info('Processing {0}- vs {1}-sequences'.format(len(query_sequences),
                                                                            len(target_sequences)))
                    results.extend(self.program.process(query_sequences, target_sequences, self))

                if sequencesToProcess and len(target_sequences) <= end_index:
                    # for palindrome program, skip directly to next
                    sequencesToProcess = False
                    self.logger.info('Processing done')

                start_index = start_index + int(self.settings.sequence_step)
                end_index = end_index + int(self.settings.sequence_step)
                if  self.settings.program == "palindrome" or (int(self.settings.end_target) > 0 and int(self.settings.end_target) < end_index):
                    sequencesToProcess = False

            if int(self.settings.end_query) > 0 and int(self.settings.end_query) < query_end:
                queriesToProcess = False
            query_start = query_start + int(self.settings.query_step)
            query_end = query_end + int(self.settings.query_step)

        nhits = len(results.hits)
        # retrieve and print results!
        self.logger.info('Processing OK ({} hits found).'.format(nhits))
        if nhits > 0:
            self.logger.info('Formatting output...')
            formatter = self._get_formatter(results)
            self.logger.info('Formatting OK.')

            self.logger.info('Writing output...')
            formatter.print_results()
            self.logger.info('Writing OK.')
            self.logger.info('Finished')
        else:
            self.logger.warning('No suitable hits produced, exiting...')

if __name__ == '__main__':
    try:
        ppw = Pypaswas()
        ppw.run()
    except Exception as exception:
        # Show complete exception when running in DEBUG
        if (hasattr(ppw.settings, 'loglevel') and
            getattr(logging, 'DEBUG') == ppw.logger.getEffectiveLevel()):
            ppw.logger.exception(str(exception))
        else:
            print 'Something unforeseen caused the program to terminate.\n' \
                  'The internal error message was: ', ','.join(exception.args)
<|endoftext|>


<commit_before># encoding: utf-8
# pylint: disable=missing-docstring
"""
This file contains initialization data for development usage only.

You can execute this code via ``invoke app.db.init_development_data``
"""
from app.extensions import db, api

from app.modules.users.models import User
from app.modules.auth.models import OAuth2Client

def init_users():
    root_user = User(username='root', email='root@localhost', password='q', is_admin=True)
    docs_user = User(username='docs', email='docs@localhost', password='w', is_active=False)
    regular_user = User(username='user', email='user@localhost', password='w')
    db.session.add(root_user)
    db.session.add(docs_user)
    db.session.add(regular_user)
    db.session.commit()
    return root_user, docs_user, regular_user

def init_auth(docs_user):
    # TODO: OpenAPI documentation has to have OAuth2 Implicit Flow instead
    # of Resource Owner Password Credentials Flow
    oauth2_client = OAuth2Client(
        client_id='docs',
        client_secret='KQ()SWK)SQK)QWSKQW(SKQ)S(QWSQW(SJ*HQ&HQW*SQ*^SSQWSGQSG',
        user_id=docs_user.id,
        _default_scopes=' '.join(api.api_v1.authorizations['oauth2_password']['scopes'])
    )
    db.session.add(oauth2_client)
    db.session.commit()
    return oauth2_client

def init():
    # pylint: disable=unused-variable
    assert User.query.count() == 0, \
        "Database is not empty. You should not re-apply fixtures! Aborted."

    root_user, docs_user, regular_user = init_users()
    init_auth(docs_user)
<commit_msg>Fixed initial development data<commit_after># encoding: utf-8
# pylint: disable=missing-docstring
"""
This file contains initialization data for development usage only.

You can execute this code via ``invoke app.db.init_development_data``
"""
from app.extensions import db, api

from app.modules.users.models import User
from app.modules.auth.models import OAuth2Client

def init_users():
    root_user = User(
        username='root',
        email='root@localhost',
        password='q',
        is_active=True,
        is_admin=True
    )
    db.session.add(root_user)
    docs_user = User(
        username='docs',
        email='docs@localhost',
        password='w',
        is_active=False
    )
    db.session.add(docs_user)
    regular_user = User(
        username='user',
        email='user@localhost',
        password='w',
        is_active=True
    )
    db.session.add(regular_user)
    db.session.commit()
    return root_user, docs_user, regular_user

def init_auth(docs_user):
    # TODO: OpenAPI documentation has to have OAuth2 Implicit Flow instead
    # of Resource Owner Password Credentials Flow
    oauth2_client = OAuth2Client(
        client_id='docs',
        client_secret='KQ()SWK)SQK)QWSKQW(SKQ)S(QWSQW(SJ*HQ&HQW*SQ*^SSQWSGQSG',
        user_id=docs_user.id,
        _default_scopes=' '.join(api.api_v1.authorizations['oauth2_password']['scopes'])
    )
    db.session.add(oauth2_client)
    db.session.commit()
    return oauth2_client

def init():
    # pylint: disable=unused-variable
    assert User.query.count() == 0, \
        "Database is not empty. You should not re-apply fixtures! Aborted."

    root_user, docs_user, regular_user = init_users()
    init_auth(docs_user)
<|endoftext|>


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

import re
import time

from base64 import urlsafe_b64encode

from module.plugins.internal.SimpleCrypter import SimpleCrypter, create_getInfo

class DlProtectCom(SimpleCrypter):
    __name__    = "DlProtectCom"
    __type__    = "crypter"
    __version__ = "0.07"
    __status__  = "testing"

    __pattern__ = r'https?://(?:www\.)?dl-protect\.com/((en|fr)/)?\w+'
    __config__  = [("activated"            , "bool", "Activated"                                        , True),
                   ("use_premium"          , "bool", "Use premium account if available"                 , True),
                   ("use_subfolder"        , "bool", "Save package to subfolder"                        , True),
                   ("subfolder_per_package", "bool", "Create a subfolder for each package"              , True),
                   ("max_wait"             , "int" , "Reconnect if waiting time is greater than minutes", 10  )]

    __description__ = """Dl-protect.com decrypter plugin"""
    __license__     = "GPLv3"
    __authors__     = [("Walter Purcaro", "vuolter@gmail.com")]

    COOKIES = [("dl-protect.com", "l", "en")]

    OFFLINE_PATTERN = r'Unfortunately, the link you are looking for is not found'

    def get_links(self):
        #: Direct link with redirect
        if not re.match(r"https?://(?:www\.)?dl-protect\.com/.+", self.req.http.lastEffectiveURL):
            return [self.req.http.lastEffectiveURL]

        post_req = {'key'       : re.search(r'name="key" value="(.+?)"', self.data).group(1),
                    'submitform': ""}

        if "Please click on continue to see the links" in self.data:
            post_req['submitform'] = "Continue"
            self.wait(2)

        else:
            mstime  = int(round(time.time() * 1000))
            b64time = "_" + urlsafe_b64encode(str(mstime)).replace("=", "%3D")

            post_req.update({'i'         : b64time,
                             'submitform': "Decrypt+link"})

            if "Password :" in self.data:
                post_req['pwd'] = self.get_password()

            if "Security Code" in self.data:
                m = re.search(r'/captcha\.php\?key=(.+?)"', self.data)
                if m is not None:
                    captcha_code = self.captcha.decrypt("http://www.dl-protect.com/captcha.php?key=" + m.group(1), input_type="gif")
                    post_req['secure'] = captcha_code

        self.data = self.load(self.pyfile.url, post=post_req)

        for errmsg in ("The password is incorrect", "The security code is incorrect"):
            if errmsg in self.data:
                self.fail(_(errmsg[1:]))

        return re.findall(r'', self.data)

getInfo = create_getInfo(DlProtectCom)
<commit_msg>Addition of captcha code in debug messages in DlProtect.py<commit_after># -*- coding: utf-8 -*-

import re
import time

from base64 import urlsafe_b64encode

from module.plugins.internal.SimpleCrypter import SimpleCrypter, create_getInfo

class DlProtectCom(SimpleCrypter):
    __name__    = "DlProtectCom"
    __type__    = "crypter"
    __version__ = "0.07"
    __status__  = "testing"

    __pattern__ = r'https?://(?:www\.)?dl-protect\.com/((en|fr)/)?\w+'
    __config__  = [("activated"            , "bool", "Activated"                                        , True),
                   ("use_premium"          , "bool", "Use premium account if available"                 , True),
                   ("use_subfolder"        , "bool", "Save package to subfolder"                        , True),
                   ("subfolder_per_package", "bool", "Create a subfolder for each package"              , True),
                   ("max_wait"             , "int" , "Reconnect if waiting time is greater than minutes", 10  )]

    __description__ = """Dl-protect.com decrypter plugin"""
    __license__     = "GPLv3"
    __authors__     = [("Walter Purcaro", "vuolter@gmail.com")]

    COOKIES = [("dl-protect.com", "l", "en")]

    OFFLINE_PATTERN = r'Unfortunately, the link you are looking for is not found'

    def get_links(self):
        #: Direct link with redirect
        if not re.match(r"https?://(?:www\.)?dl-protect\.com/.+", self.req.http.lastEffectiveURL):
            return [self.req.http.lastEffectiveURL]

        post_req = {'key'       : re.search(r'name="key" value="(.+?)"', self.data).group(1),
                    'submitform': ""}

        if "Please click on continue to see the links" in self.data:
            post_req['submitform'] = "Continue"
            self.wait(2)

        else:
            mstime  = int(round(time.time() * 1000))
            b64time = "_" + urlsafe_b64encode(str(mstime)).replace("=", "%3D")

            post_req.update({'i'         : b64time,
                             'submitform': "Decrypt+link"})

            if "Password :" in self.data:
                post_req['pwd'] = self.get_password()

            if "Security Code" in self.data:
                m = re.search(r'/captcha\.php\?key=(.+?)"', self.data)
                if m is not None:
                    captcha_code = self.captcha.decrypt("http://www.dl-protect.com/captcha.php?key=" + m.group(1), input_type="gif")
                    self.log_debug("Captcha code: %s" % captcha_code)
                    post_req['secure'] = captcha_code

        self.data = self.load(self.pyfile.url, post=post_req)

        for errmsg in ("The password is incorrect", "The security code is incorrect"):
            if errmsg in self.data:
                self.fail(_(errmsg[1:]))

        return re.findall(r'', self.data)

getInfo = create_getInfo(DlProtectCom)
<|endoftext|>


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

import re
import urllib

import BeautifulSoup

from ..captcha.ReCaptcha import ReCaptcha
from ..internal.SimpleHoster import SimpleHoster

class ZippyshareCom(SimpleHoster):
    __name__ = "ZippyshareCom"
    __type__ = "hoster"
    __version__ = "0.92"
    __status__ = "testing"

    __pattern__ = r'http://www\d{0,3}\.zippyshare\.com/v(/|iew\.jsp.*key=)(?P<KEY>[\w^_]+)'
    __config__ = [("activated", "bool", "Activated", True),
                  ("use_premium", "bool", "Use premium account if available", True),
                  ("fallback", "bool",
                   "Fallback to free download if premium fails", True),
                  ("chk_filesize", "bool", "Check file size", True),
                  ("max_wait", "int", "Reconnect if waiting time is greater than minutes", 10)]

    __description__ = """Zippyshare.com hoster plugin"""
    __license__ = "GPLv3"
    __authors__ = [("Walter Purcaro", "vuolter@gmail.com"),
                   ("sebdelsol", "seb.morin@gmail.com"),
                   ("GammaC0de", "nitzo2001[AT]yahoo[DOT]com")]

    COOKIES = [("zippyshare.com", "ziplocale", "en")]

    NAME_PATTERN = r'(<title>Zippyshare.com - |"/)(?P<N>[^/]+)(</title>|";)'
    SIZE_PATTERN = r'>Size:.+?">(?P<S>[\d.,]+) (?P<U>[\w^_]+)'
    OFFLINE_PATTERN = r'does not exist (anymore )?on this server<'
    TEMP_OFFLINE_PATTERN = None

    LINK_PREMIUM_PATTERN = r"document.location = '(.+?)'"

    def setup(self):
        self.chunk_limit = -1
        self.multiDL = True
        self.resume_download = True

    def handle_free(self, pyfile):
        self.captcha = ReCaptcha(pyfile)
        captcha_key = self.captcha.detect_key()

        if captcha_key:
            try:
                self.link = re.search(self.LINK_PREMIUM_PATTERN, self.data)
                self.captcha.challenge()

            except Exception, e:
                self.error(e)

        else:
            self.link = self.fixurl(self.get_link())
            if ".com/pd/" in self.link:
                self.load(self.link)
                self.link = self.link.replace(".com/pd/", ".com/d/")

        if self.link and pyfile.name == "file.html":
            pyfile.name = urllib.unquote(self.link.split('/')[-1])

    def get_link(self):
        #: Get all the scripts inside the html body
        soup = BeautifulSoup.BeautifulSoup(self.data)
        scripts = [
            s.getText() for s in soup.body.findAll(
                'script',
                type='text/javascript') if "('dlbutton').href =" in s.getText()]

        #: Emulate a document in JS
        inits = ['''
                var document = {}
                document.getElementById = function(x) {
                    if (!this.hasOwnProperty(x)) {
                        this[x] = {getAttribute : function(x) { return this[x] } }
                    }
                    return this[x]
                }
                ''']

        #: inits is meant to be populated with the initialization of all the DOM elements found in the scripts
        eltRE = r'getElementById\([\'"](.+?)[\'"]\)(\.)?(getAttribute\([\'"])?(\w+)?([\'"]\))?'
        for m in re.findall(eltRE, ' '.join(scripts)):
            JSid, JSattr = m[0], m[3]
            values = filter(None, (elt.get(JSattr, None)
                                   for elt in soup.findAll(id=JSid)))
            if values:
                inits.append('document.getElementById("%s")["%s"] = "%s"' % (
                    JSid, JSattr, values[-1]))

        #: Add try/catch in JS to handle deliberate errors
        scripts = ['\n'.join(('try{', script, '} catch(err){}'))
                   for script in scripts]

        #: Get the file's url by evaluating all the scripts
        scripts = inits + scripts + ['document.dlbutton.href']

        return self.js.eval('\n'.join(scripts))

<commit_msg>[ZippyshareCom] Update<commit_after># -*- coding: utf-8 -*-

import re
import urllib

import BeautifulSoup

from ..captcha.ReCaptcha import ReCaptcha
from ..internal.SimpleHoster import SimpleHoster

class ZippyshareCom(SimpleHoster):
    __name__ = "ZippyshareCom"
    __type__ = "hoster"
    __version__ = "0.93"
    __status__ = "testing"

    __pattern__ = r'http://(?P<HOST>www\d{0,3}\.zippyshare\.com)/(?:[vd]/|view\.jsp.*key=)(?P<KEY>[\w^_]+)'
    __config__ = [("activated", "bool", "Activated", True),
                  ("use_premium", "bool", "Use premium account if available", True),
                  ("fallback", "bool", "Fallback to free download if premium fails", True),
                  ("chk_filesize", "bool", "Check file size", True),
                  ("max_wait", "int", "Reconnect if waiting time is greater than minutes", 10)]

    __description__ = """Zippyshare.com hoster plugin"""
    __license__ = "GPLv3"
    __authors__ = [("Walter Purcaro", "vuolter@gmail.com"),
                   ("sebdelsol", "seb.morin@gmail.com"),
                   ("GammaC0de", "nitzo2001[AT]yahoo[DOT]com")]

    COOKIES = [("zippyshare.com", "ziplocale", "en")]

    URL_REPLACEMENTS = [(__pattern__ + ".*", r'http://\g<HOST>/v/\g<KEY>/file.html')]

    NAME_PATTERN = r'>Name:.+?">(?P<N>.+?)<'
    SIZE_PATTERN = r'>Size:.+?">(?P<S>[\d.,]+) (?P<U>[\w^_]+)'
    OFFLINE_PATTERN = r'does not exist (anymore )?on this server<'
    TEMP_OFFLINE_PATTERN = r'^unmatchable$'

    LINK_PATTERN = r"document.location = '(.+?)'"

    def setup(self):
        self.chunk_limit = -1
        self.multiDL = True
        self.resume_download = True

    def handle_free(self, pyfile):
        self.captcha = ReCaptcha(pyfile)
        captcha_key = self.captcha.detect_key()

        if captcha_key:
            try:
                self.link = re.search(self.LINK_PATTERN, self.data)
                self.captcha.challenge()

            except Exception, e:
                self.error(e)

        else:
            self.link = self.fixurl(self.get_link())
            if ".com/pd/" in self.link:
                self.load(self.link)
                self.link = self.link.replace(".com/pd/", ".com/d/")

        if self.link and pyfile.name == "file.html":
            pyfile.name = urllib.unquote(self.link.split('/')[-1])

    def get_link(self):
        #: Get all the scripts inside the html body
        soup = BeautifulSoup.BeautifulSoup(self.data)
        scripts = [
            s.getText() for s in soup.body.findAll(
                'script',
                type='text/javascript') if "('dlbutton').href =" in s.getText()]

        #: Emulate a document in JS
        inits = ['''
                var document = {}
                document.getElementById = function(x) {
                    if (!this.hasOwnProperty(x)) {
                        this[x] = {getAttribute : function(x) { return this[x] } }
                    }
                    return this[x]
                }
                ''']

        #: inits is meant to be populated with the initialization of all the DOM elements found in the scripts
        eltRE = r'getElementById\([\'"](.+?)[\'"]\)(\.)?(getAttribute\([\'"])?(\w+)?([\'"]\))?'
        for m in re.findall(eltRE, ' '.join(scripts)):
            JSid, JSattr = m[0], m[3]
            values = filter(None, (elt.get(JSattr, None)
                                   for elt in soup.findAll(id=JSid)))
            if values:
                inits.append('document.getElementById("%s")["%s"] = "%s"' % (
                    JSid, JSattr, values[-1]))

        #: Add try/catch in JS to handle deliberate errors
        scripts = ['\n'.join(('try{', script, '} catch(err){}'))
                   for script in scripts]

        #: Get the file's url by evaluating all the scripts
        scripts = inits + scripts + ['document.dlbutton.href']

        return self.js.eval('\n'.join(scripts))

<|endoftext|>


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

from django import test

from .. import models

class TestCommentForm(test.TestCase):

    def test_author_is_required(self):
        self.fail('test_author_is_required FAULT')

    def test_email_is_required(self):
        self.fail('test_email_is_required FAULT')

    def test_email_pattern(self):
        self.fail('test_email_pattern FAULT')

    def test_content_is_required(self):
        self.fail('test_content_is_required FAULT')
<commit_msg>Test comment form #15<commit_after># -*- coding: utf-8 -*-
import unittest

from django import test

from .. import models
from .. import forms

class TestFormBase(test.TestCase):
    @classmethod
    def setUpTestData(cls):
        # Validation errors messages
        cls.REQUIRED_ERROR_MSG = 'This field is required.'
        cls.EMAIL_ERROR_MSG = 'Enter a valid email address.'

class TestCommentForm(TestFormBase):
    def test_author_is_required(self):
        form = forms.CommentForm(models.Comment().__dict__)
        error_msg = form['author'].errors[0]

        self.assertEqual(error_msg, self.REQUIRED_ERROR_MSG)

    def test_email_is_required(self):
        form = forms.CommentForm(models.Comment().__dict__)
        error_msg = form['email'].errors[0]

        self.assertEqual(error_msg, self.REQUIRED_ERROR_MSG)

    def test_malformed_email(self):
        form = forms.CommentForm(
            models.Comment(email='malformed@email').__dict__)
        error_msg = form['email'].errors[0]

        self.assertEqual(error_msg, self.EMAIL_ERROR_MSG)

    def test_content_is_required(self):
        form = forms.CommentForm(models.Comment().__dict__)
        error_msg = form['content'].errors[0]

        self.assertEqual(error_msg, self.REQUIRED_ERROR_MSG)
<|endoftext|>


<commit_before>try:
    import pkg_resources
    pkg_resources.require('Thrift')
except ImportError:
    pass
from cassandra.ttypes import Column, ColumnOrSuperColumn, ColumnParent, \
    ColumnPath, ConsistencyLevel, NotFoundException, SlicePredicate, \
    SliceRange, SuperColumn, Mutation, Deletion, Clock, KeyRange

import time, sys

__all__ = ['gm_timestamp', 'ColumnFamily']

def gm_timestamp():
    """
    Returns
    -------
    int : UNIX epoch time in GMT
    """
    return int(time.time() * 1e6)

def create_SlicePredicate(columns, column_start, column_finish, column_reversed, column_count):
    if columns is not None:
        return SlicePredicate(column_names=columns)
    sr = SliceRange(start=column_start, finish=column_finish,
                    reversed=column_reversed, count=column_count)
    return SlicePredicate(slice_range=sr)

class ColumnFamily(object):
    def __init__(self, client, column_family, buffer_size=1024,
                 read_consistency_level=ConsistencyLevel.ONE,
                 write_consistency_level=ConsistencyLevel.ONE,
                 timestamp=gm_timestamp, super=False,
                 dict_class=dict):
        """
        Construct a ColumnFamily

        Parameters
        ----------
        client   : cassandra.Cassandra.Client
            Cassandra client with thrift API
        column_family : str
            The name of this ColumnFamily
        buffer_size : int
            When calling get_range(), the intermediate results need to be
            buffered if we are fetching many rows, otherwise the Cassandra
            server will overallocate memory and fail.  This is the size of
            that buffer.
        read_consistency_level : ConsistencyLevel
            Affects the guaranteed replication factor before returning from
            any read operation
        write_consistency_level : ConsistencyLevel
            Affects the guaranteed replication factor before returning from
            any write operation
        timestamp : function
            The default timestamp function returns:
            int(time.mktime(time.gmtime()))
            Or the number of seconds since Unix epoch in GMT.
            Set timestamp to replace the default timestamp function with your
            own.
        super : bool
            Whether this ColumnFamily has SuperColumns
        dict_class : class (must act like the dict type)
            The default dict_class is dict.
            If the order of columns matter to you, pass your own dictionary
            class, or python 2.7's new collections.OrderedDict. All returned
            rows and subcolumns are instances of this.
        """
        self.client = client
        self.column_family = column_family
        self.buffer_size = buffer_size
        self.read_consistency_level = read_consistency_level
        self.write_consistency_level = write_consistency_level
        self.timestamp = timestamp
        self.super = super
        self.dict_class = dict_class

    def _convert_Column_to_base(self, column, include_timestamp):
        if include_timestamp:
            return (column.value, column.clock.timestamp)
        return column.value

    def _convert_SuperColumn_to_base(self, super_column, include_timestamp):
        ret = self.dict_class()
        for column in super_column.columns:
            ret[column.name] = self._convert_Column_to_base(column, include_timestamp)
        return ret

    def _convert_ColumnOrSuperColumns_to_dict_class(self, list_col_or_super, include_timestamp):
        ret = self.dict_class()
        for col_or_super in list_col_or_super:
            if col_or_super.super_column is not None:
                col = col_or_super.super_column
                ret[col.name] = self._convert_SuperColumn_to_base(col, include_timestamp)
            else:
                col = col_or_super.column
                ret[col.name] = self._convert_Column_to_base(col, include_timestamp)
        return ret

    def _rcl(self, alternative):
        """Helper function that returns self.read_consistency_level if
        alternative is None, otherwise returns alternative"""
        if alternative is None:
            return self.read_consistency_level
        return alternative

    def _wcl(self, alternative):
        """Helper function that returns self.write_consistency_level
        if alternative is None, otherwise returns alternative"""
        if alternative is None:
            return self.write_consistency_level
        return alternative

    def get(self, key, columns=None, column_start="", column_finish="",
            column_reversed=False, column_count=100, include_timestamp=False,
            super_column=None, read_consistency_level = None):
        """
        Fetch a key from a Cassandra server

        Parameters
        ----------
        key : str
            The key to fetch
        columns : [str]
            Limit the columns or super_columns fetched to the specified list
        column_start : str
            Only fetch when a column or super_column is >= column_start
        column_finish : str
            Only fetch when a column or super_column is <= column_finish
        column_reversed : bool
            Fetch the columns or super_columns in reverse order. This will do
            nothing unless you passed a dict_class to the constructor.
        column_count : int
            Limit the number of columns or super_columns fetched per key
        include_timestamp : bool
            If true, return a (value, timestamp) tuple for each column
        super_column : str
            Return columns only in this super_column
        read_consistency_level : ConsistencyLevel
            Affects the guaranteed replication factor before returning from
            any read operation

        Returns
        -------
        if include_timestamp == True: {'column': ('value', timestamp)}
        else: {'column': 'value'}
        """
        cp = ColumnParent(column_family=self.column_family, super_column=super_column)
        sp = create_SlicePredicate(columns, column_start, column_finish,
                                   column_reversed, column_count)

        list_col_or_super = self.client.get_slice(key, cp, sp,
                                                  self._rcl(read_consistency_level))

        if len(list_col_or_super) == 0:
            raise NotFoundException()
        return self._convert_ColumnOrSuperColumns_to_dict_class(list_col_or_super, include_timestamp)

    def multiget(self, keys, columns=None, column_start="", column_finish="",
                 column_reversed=False, column_count=100, include_timestamp=False,
                 super_column=None, read_consistency_level = None):
        """
        Fetch multiple key from a Cassandra server

        Parameters
        ----------
        keys : [str]
            A list of keys to fetch
        columns : [str]
            Limit the columns or super_columns fetched to the specified list
        column_start : str
            Only fetch when a column or super_column is >= column_start
        column_finish : str
            Only fetch when a column or super_column is <= column_finish
        column_reversed : bool
            Fetch the columns or super_columns in reverse order. This will do
            nothing unless you passed a dict_class to the constructor.
        column_count : int
            Limit the number of columns or super_columns fetched per key
        include_timestamp : bool
            If true, return a (value, timestamp) tuple for each column
        super_column : str
            Return columns only in this super_column
        read_consistency_level : ConsistencyLevel
            Affects the guaranteed replication factor before returning from
            any read operation

        Returns
        -------
        if include_timestamp == True: {'key': {'column': ('value', timestamp)}}
        else: {'key': {'column': 'value'}}
        """
        cp = ColumnParent(column_family=self.column_family, super_column=super_column)
        sp = create_SlicePredicate(columns, column_start, column_finish,
                                   column_reversed, column_count)

        keymap = self.client.multiget_slice(keys, cp, sp,
                                            self._rcl(read_consistency_level))

        ret = dict()
        for key, columns in keymap.iteritems():
            if len(columns) > 0:
                ret[key] = self._convert_ColumnOrSuperColumns_to_dict_class(columns, include_timestamp)
        return ret

    def get_count(self, key, super_column=None, read_consistency_level = None):
        """
        Count the number of columns for a key

        Parameters
        ----------
        key : str
            The key with which to count columns
        super_column : str
            Count the columns only in this super_column
        read_consistency_level : ConsistencyLevel
            Affects the guaranteed replication factor before returning from
            any read operation

        Returns
        -------
        int Count of columns
        """
        cp = ColumnParent(column_family=self.column_family, super_column=super_column)
        sp = SlicePredicate(slice_range=SliceRange(start='',
                                                   finish='',
                                                   count=sys.maxint))
        return self.client.get_count(key, cp, sp,
                                     self._rcl(read_consistency_level))

    def get_range(self, start="", finish="", columns=None, column_start="",
                  column_finish="", column_reversed=False, column_count=100,
                  row_count=None, include_timestamp=False,
                  super_column=None, read_consistency_level = None):
        """
        Get an iterator over keys in a specified range

        Parameters
        ----------
        start : str
            Start from this key (inclusive)
        finish : str
            End at this key (inclusive)
        columns : [str]
            Limit the columns or super_columns fetched to the specified list
        column_start : str
            Only fetch when a column or super_column is >= column_start
        column_finish : str
            Only fetch when a column or super_column is <= column_finish
        column_reversed : bool
            Fetch the columns or super_columns in reverse order. This will do
            nothing unless you passed a dict_class to the constructor.
        column_count : int
            Limit the number of columns or super_columns fetched per key
        row_count : int
            Limit the number of rows fetched
        include_timestamp : bool
            If true, return a (value, timestamp) tuple for each column
        super_column : string
            Return columns only in this super_column
        read_consistency_level : ConsistencyLevel
            Affects the guaranteed replication factor before returning from
            any read operation

        Returns
        -------
        iterator over ('key', {'column': 'value'})
        """
        cp = ColumnParent(column_family=self.column_family, super_column=super_column)
        sp = create_SlicePredicate(columns, column_start, column_finish,
                                   column_reversed, column_count)

        count = 0
        i = 0
        last_key = start

        buffer_size = self.buffer_size
        if row_count is not None:
            buffer_size = min(row_count, self.buffer_size)
        while True:
            key_range = KeyRange(start_key=last_key, end_key=finish, count=buffer_size)
            key_slices = self.client.get_range_slices(cp, sp, key_range,
                                                     self._rcl(read_consistency_level))
            # This may happen if nothing was ever inserted
            if key_slices is None:
                return
            for j, key_slice in enumerate(key_slices):
                # Ignore the first element after the first iteration
                # because it will be a duplicate.
                if j == 0 and i != 0:
                    continue
                yield (key_slice.key,
                       self._convert_ColumnOrSuperColumns_to_dict_class(key_slice.columns, include_timestamp))
                count += 1
                if row_count is not None and count >= row_count:
                    return

            if len(key_slices) != self.buffer_size:
                return
            last_key = key_slices[-1].key
            i += 1

    def insert(self, key, columns, write_consistency_level = None):
        """
        Insert or update columns for a key

        Parameters
        ----------
        key : str
            The key to insert or update the columns at
        columns : dict
            Column: {'column': 'value'}
            SuperColumn: {'column': {'subcolumn': 'value'}}
            The columns or supercolumns to insert or update
        write_consistency_level : ConsistencyLevel
            Affects the guaranteed replication factor before returning from
            any write operation

        Returns
        -------
        int timestamp
        """
        clock = Clock(timestamp=self.timestamp())

        cols = []
        for c, v in columns.iteritems():
            if self.super:
                subc = [Column(name=subname, value=subvalue, clock=clock) \
                        for subname, subvalue in v.iteritems()]
                column = SuperColumn(name=c, columns=subc)
                cols.append(Mutation(column_or_supercolumn=ColumnOrSuperColumn(super_column=column)))
            else:
                column = Column(name=c, value=v, clock=clock)
                cols.append(Mutation(column_or_supercolumn=ColumnOrSuperColumn(column=column)))
        self.client.batch_mutate({key: {self.column_family: cols}},
                                 self._wcl(write_consistency_level))
        return clock.timestamp

    def remove(self, key, columns=None, super_column=None, write_consistency_level = None):
        """
        Remove a specified key or columns

        Parameters
        ----------
        key : str
            The key to remove. If columns is not set, remove all columns
        columns : list
            Delete the columns or super_columns in this list
        super_column : str
            Delete the columns from this super_column
        write_consistency_level : ConsistencyLevel
            Affects the guaranteed replication factor before returning from
            any write operation

        Returns
        -------
        int timestamp
        """
        clock = Clock(timestamp=self.timestamp())
        if columns is not None:
            # Deletion doesn't support SliceRange predicates as of Cassandra 0.6.0,
            # so we can't add column_start, column_finish, etc... yet
            sp = SlicePredicate(column_names=columns)
            deletion = Deletion(clock=clock, super_column=super_column, predicate=sp)
            mutation = Mutation(deletion=deletion)
            self.client.batch_mutate({key: {self.column_family: [mutation]}},
                                     self._wcl(write_consistency_level))
        else:
            cp = ColumnPath(column_family=self.column_family, super_column=super_column)
            self.client.remove(key, cp, clock,
                               self._wcl(write_consistency_level))
        return clock.timestamp
<commit_msg>fix OverflowError on get_count()<commit_after>try:
    import pkg_resources
    pkg_resources.require('Thrift')
except ImportError:
    pass
from cassandra.ttypes import Column, ColumnOrSuperColumn, ColumnParent, \
    ColumnPath, ConsistencyLevel, NotFoundException, SlicePredicate, \
    SliceRange, SuperColumn, Mutation, Deletion, Clock, KeyRange

import time, sys

__all__ = ['gm_timestamp', 'ColumnFamily']

def gm_timestamp():
    """
    Returns
    -------
    int : UNIX epoch time in GMT
    """
    return int(time.time() * 1e6)

def create_SlicePredicate(columns, column_start, column_finish, column_reversed, column_count):
    if columns is not None:
        return SlicePredicate(column_names=columns)
    sr = SliceRange(start=column_start, finish=column_finish,
                    reversed=column_reversed, count=column_count)
    return SlicePredicate(slice_range=sr)

class ColumnFamily(object):
    def __init__(self, client, column_family, buffer_size=1024,
                 read_consistency_level=ConsistencyLevel.ONE,
                 write_consistency_level=ConsistencyLevel.ONE,
                 timestamp=gm_timestamp, super=False,
                 dict_class=dict):
        """
        Construct a ColumnFamily

        Parameters
        ----------
        client   : cassandra.Cassandra.Client
            Cassandra client with thrift API
        column_family : str
            The name of this ColumnFamily
        buffer_size : int
            When calling get_range(), the intermediate results need to be
            buffered if we are fetching many rows, otherwise the Cassandra
            server will overallocate memory and fail.  This is the size of
            that buffer.
        read_consistency_level : ConsistencyLevel
            Affects the guaranteed replication factor before returning from
            any read operation
        write_consistency_level : ConsistencyLevel
            Affects the guaranteed replication factor before returning from
            any write operation
        timestamp : function
            The default timestamp function returns:
            int(time.mktime(time.gmtime()))
            Or the number of seconds since Unix epoch in GMT.
            Set timestamp to replace the default timestamp function with your
            own.
        super : bool
            Whether this ColumnFamily has SuperColumns
        dict_class : class (must act like the dict type)
            The default dict_class is dict.
            If the order of columns matter to you, pass your own dictionary
            class, or python 2.7's new collections.OrderedDict. All returned
            rows and subcolumns are instances of this.
        """
        self.client = client
        self.column_family = column_family
        self.buffer_size = buffer_size
        self.read_consistency_level = read_consistency_level
        self.write_consistency_level = write_consistency_level
        self.timestamp = timestamp
        self.super = super
        self.dict_class = dict_class

    def _convert_Column_to_base(self, column, include_timestamp):
        if include_timestamp:
            return (column.value, column.clock.timestamp)
        return column.value

    def _convert_SuperColumn_to_base(self, super_column, include_timestamp):
        ret = self.dict_class()
        for column in super_column.columns:
            ret[column.name] = self._convert_Column_to_base(column, include_timestamp)
        return ret

    def _convert_ColumnOrSuperColumns_to_dict_class(self, list_col_or_super, include_timestamp):
        ret = self.dict_class()
        for col_or_super in list_col_or_super:
            if col_or_super.super_column is not None:
                col = col_or_super.super_column
                ret[col.name] = self._convert_SuperColumn_to_base(col, include_timestamp)
            else:
                col = col_or_super.column
                ret[col.name] = self._convert_Column_to_base(col, include_timestamp)
        return ret

    def _rcl(self, alternative):
        """Helper function that returns self.read_consistency_level if
        alternative is None, otherwise returns alternative"""
        if alternative is None:
            return self.read_consistency_level
        return alternative

    def _wcl(self, alternative):
        """Helper function that returns self.write_consistency_level
        if alternative is None, otherwise returns alternative"""
        if alternative is None:
            return self.write_consistency_level
        return alternative

    def get(self, key, columns=None, column_start="", column_finish="",
            column_reversed=False, column_count=100, include_timestamp=False,
            super_column=None, read_consistency_level = None):
        """
        Fetch a key from a Cassandra server

        Parameters
        ----------
        key : str
            The key to fetch
        columns : [str]
            Limit the columns or super_columns fetched to the specified list
        column_start : str
            Only fetch when a column or super_column is >= column_start
        column_finish : str
            Only fetch when a column or super_column is <= column_finish
        column_reversed : bool
            Fetch the columns or super_columns in reverse order. This will do
            nothing unless you passed a dict_class to the constructor.
        column_count : int
            Limit the number of columns or super_columns fetched per key
        include_timestamp : bool
            If true, return a (value, timestamp) tuple for each column
        super_column : str
            Return columns only in this super_column
        read_consistency_level : ConsistencyLevel
            Affects the guaranteed replication factor before returning from
            any read operation

        Returns
        -------
        if include_timestamp == True: {'column': ('value', timestamp)}
        else: {'column': 'value'}
        """
        cp = ColumnParent(column_family=self.column_family, super_column=super_column)
        sp = create_SlicePredicate(columns, column_start, column_finish,
                                   column_reversed, column_count)

        list_col_or_super = self.client.get_slice(key, cp, sp,
                                                  self._rcl(read_consistency_level))

        if len(list_col_or_super) == 0:
            raise NotFoundException()
        return self._convert_ColumnOrSuperColumns_to_dict_class(list_col_or_super, include_timestamp)

    def multiget(self, keys, columns=None, column_start="", column_finish="",
                 column_reversed=False, column_count=100, include_timestamp=False,
                 super_column=None, read_consistency_level = None):
        """
        Fetch multiple key from a Cassandra server

        Parameters
        ----------
        keys : [str]
            A list of keys to fetch
        columns : [str]
            Limit the columns or super_columns fetched to the specified list
        column_start : str
            Only fetch when a column or super_column is >= column_start
        column_finish : str
            Only fetch when a column or super_column is <= column_finish
        column_reversed : bool
            Fetch the columns or super_columns in reverse order. This will do
            nothing unless you passed a dict_class to the constructor.
        column_count : int
            Limit the number of columns or super_columns fetched per key
        include_timestamp : bool
            If true, return a (value, timestamp) tuple for each column
        super_column : str
            Return columns only in this super_column
        read_consistency_level : ConsistencyLevel
            Affects the guaranteed replication factor before returning from
            any read operation

        Returns
        -------
        if include_timestamp == True: {'key': {'column': ('value', timestamp)}}
        else: {'key': {'column': 'value'}}
        """
        cp = ColumnParent(column_family=self.column_family, super_column=super_column)
        sp = create_SlicePredicate(columns, column_start, column_finish,
                                   column_reversed, column_count)

        keymap = self.client.multiget_slice(keys, cp, sp,
                                            self._rcl(read_consistency_level))

        ret = dict()
        for key, columns in keymap.iteritems():
            if len(columns) > 0:
                ret[key] = self._convert_ColumnOrSuperColumns_to_dict_class(columns, include_timestamp)
        return ret

    MAX_COUNT = 2**31-1
    def get_count(self, key, super_column=None, read_consistency_level = None):
        """
        Count the number of columns for a key

        Parameters
        ----------
        key : str
            The key with which to count columns
        super_column : str
            Count the columns only in this super_column
        read_consistency_level : ConsistencyLevel
            Affects the guaranteed replication factor before returning from
            any read operation

        Returns
        -------
        int Count of columns
        """
        cp = ColumnParent(column_family=self.column_family, super_column=super_column)
        sp = SlicePredicate(slice_range=SliceRange(start='',
                                                   finish='',
                                                   count=self.MAX_COUNT))
        return self.client.get_count(key, cp, sp,
                                     self._rcl(read_consistency_level))

    def get_range(self, start="", finish="", columns=None, column_start="",
                  column_finish="", column_reversed=False, column_count=100,
                  row_count=None, include_timestamp=False,
                  super_column=None, read_consistency_level = None):
        """
        Get an iterator over keys in a specified range

        Parameters
        ----------
        start : str
            Start from this key (inclusive)
        finish : str
            End at this key (inclusive)
        columns : [str]
            Limit the columns or super_columns fetched to the specified list
        column_start : str
            Only fetch when a column or super_column is >= column_start
        column_finish : str
            Only fetch when a column or super_column is <= column_finish
        column_reversed : bool
            Fetch the columns or super_columns in reverse order. This will do
            nothing unless you passed a dict_class to the constructor.
        column_count : int
            Limit the number of columns or super_columns fetched per key
        row_count : int
            Limit the number of rows fetched
        include_timestamp : bool
            If true, return a (value, timestamp) tuple for each column
        super_column : string
            Return columns only in this super_column
        read_consistency_level : ConsistencyLevel
            Affects the guaranteed replication factor before returning from
            any read operation

        Returns
        -------
        iterator over ('key', {'column': 'value'})
        """
        cp = ColumnParent(column_family=self.column_family, super_column=super_column)
        sp = create_SlicePredicate(columns, column_start, column_finish,
                                   column_reversed, column_count)

        count = 0
        i = 0
        last_key = start

        buffer_size = self.buffer_size
        if row_count is not None:
            buffer_size = min(row_count, self.buffer_size)
        while True:
            key_range = KeyRange(start_key=last_key, end_key=finish, count=buffer_size)
            key_slices = self.client.get_range_slices(cp, sp, key_range,
                                                     self._rcl(read_consistency_level))
            # This may happen if nothing was ever inserted
            if key_slices is None:
                return
            for j, key_slice in enumerate(key_slices):
                # Ignore the first element after the first iteration
                # because it will be a duplicate.
                if j == 0 and i != 0:
                    continue
                yield (key_slice.key,
                       self._convert_ColumnOrSuperColumns_to_dict_class(key_slice.columns, include_timestamp))
                count += 1
                if row_count is not None and count >= row_count:
                    return

            if len(key_slices) != self.buffer_size:
                return
            last_key = key_slices[-1].key
            i += 1

    def insert(self, key, columns, write_consistency_level = None):
        """
        Insert or update columns for a key

        Parameters
        ----------
        key : str
            The key to insert or update the columns at
        columns : dict
            Column: {'column': 'value'}
            SuperColumn: {'column': {'subcolumn': 'value'}}
            The columns or supercolumns to insert or update
        write_consistency_level : ConsistencyLevel
            Affects the guaranteed replication factor before returning from
            any write operation

        Returns
        -------
        int timestamp
        """
        clock = Clock(timestamp=self.timestamp())

        cols = []
        for c, v in columns.iteritems():
            if self.super:
                subc = [Column(name=subname, value=subvalue, clock=clock) \
                        for subname, subvalue in v.iteritems()]
                column = SuperColumn(name=c, columns=subc)
                cols.append(Mutation(column_or_supercolumn=ColumnOrSuperColumn(super_column=column)))
            else:
                column = Column(name=c, value=v, clock=clock)
                cols.append(Mutation(column_or_supercolumn=ColumnOrSuperColumn(column=column)))
        self.client.batch_mutate({key: {self.column_family: cols}},
                                 self._wcl(write_consistency_level))
        return clock.timestamp

    def remove(self, key, columns=None, super_column=None, write_consistency_level = None):
        """
        Remove a specified key or columns

        Parameters
        ----------
        key : str
            The key to remove. If columns is not set, remove all columns
        columns : list
            Delete the columns or super_columns in this list
        super_column : str
            Delete the columns from this super_column
        write_consistency_level : ConsistencyLevel
            Affects the guaranteed replication factor before returning from
            any write operation

        Returns
        -------
        int timestamp
        """
        clock = Clock(timestamp=self.timestamp())
        if columns is not None:
            # Deletion doesn't support SliceRange predicates as of Cassandra 0.6.0,
            # so we can't add column_start, column_finish, etc... yet
            sp = SlicePredicate(column_names=columns)
            deletion = Deletion(clock=clock, super_column=super_column, predicate=sp)
            mutation = Mutation(deletion=deletion)
            self.client.batch_mutate({key: {self.column_family: [mutation]}},
                                     self._wcl(write_consistency_level))
        else:
            cp = ColumnPath(column_family=self.column_family, super_column=super_column)
            self.client.remove(key, cp, clock,
                               self._wcl(write_consistency_level))
        return clock.timestamp
<|endoftext|>


<commit_before>#! /usr/bin/env python
# -*- coding: utf-8 -*-
#
from rpython.rlib import jit

def print_help(argv):
    print """Welcome to Pycket.
%s [<option> ...] <argument> ...
 File and expression options:
  -e <exprs>, --eval <exprs> : Evaluate <exprs>, prints results
  -f <file>, --load <file> : Like -e '(load "<file>")' without printing
  -t <file>, --require <file> : Like -e '(require (file "<file>"))'
  -l <path>, --lib <path> : Like -e '(require (lib "<path>"))'
  -r <file>, --script <file> : Same as -f <file> -N <file> --
  -u <file>, --require-script <file> : Same as -t <file> -N <file> --
  -g <json>, --eval-json <json> : Load and instantiate a linklet json
  -gg <json> : Like -g but stops after instantiating

 Interaction options:
  -i, --repl : Run interactive read-eval-print loop; implies -v
  -n, --no-lib : Skip `(require (lib "<init-lib>"))' for -i/-e/-f/-r
  -v, --version : Show version

 Configuration options:
  --kernel : ignore everything, only load the #%%kernel (for development purposes)
  -I <path> : Set <init-lib> to <path> (sets language)
  -X <dir>, --collects <dir> : Main collects at <dir> (or "" disables all)
  -G <dir>, --config <dir> : Main configuration directory at <dir>
  -A <dir>, --addon <dir> : Addon directory at <dir>
  -N <file>, --name <file> : Sets `(find-system-path 'run-file)' to <file>
  --save-callgraph : save the jit output

 Meta options:
  --verbose <level> : Print the debug logs. <level> : [0,1] (defaults to 0)
  --jit <jitargs> : Set RPython JIT options may be 'default', 'off',
                    or 'param=value,param=value' list
  -- : No argument following this switch is used as a switch
  -h, --help : Show this information and exits, ignoring other options
Default options:
 If only an argument is specified, it is loaded and evaluated
""" % (argv[0])

file_expr_opts = ["-e", "--eval",
                  "-f", "--load",
                  "-t", "--require",
                  "-l", "--lib",
                  "-r", "--script",
                  "-u", "--require-script",
                  "-g", "--eval-json",
                  "-gg"]
inter_opts = ["-i", "--repl",
              "-n", "--no-lib",
              "-v", "--version"]
conf_opts = ["--kernel",
             "-I",
             "-X", "--collects",
             "-G", "--config",
             "-A", "--addon",
             "-N", "--name",
             "--save-callgraph"]
meta_opts = ["--verbose", "--jit", "-h"]

all_opts = file_expr_opts + inter_opts + conf_opts + meta_opts

verbosity_levels = ['0','1']

def add_name(names, name, val, replace=False):
    if name not in names or replace:
        names[name] = [val]
    else:
        names[name].append(val)

def is_rkt_file(name_str):
    return ".rkt" in name_str

def parse_args(argv):
    INIT = -1
    RETURN_OK = 0
    MISSING_ARG = 5
    JUST_EXIT = 3

    config = {
        'repl' : False,
        'no-lib' : False,
        'version' : False,
        'stop' : False,
        'just_kernel' : False,
        'verbose' : False
    }
    names = {
        # 'file': "",
        # 'exprs': "",
    }
    args = []
    retval = INIT
    i = 1
    to = len(argv)
    while i < to:
        if False:
            pass

        elif argv[i] in ["-e", "--eval"]:
            if to <= i + 1 or argv[i+1] in all_opts:
                print "missing argument after %s" % argv[i]
                retval = MISSING_ARG
                break

            i += 1
            add_name(names, 'loads', "eval")
            add_name(names, 'load_arguments', argv[i])
            retval = RETURN_OK

        elif argv[i] in ["-f", "--load", "-r", "--script"]:
            if to <= i + 1 or argv[i+1] in all_opts:
                print "missing argument after %s" % argv[i]
                retval = MISSING_ARG
                break

            if argv[i] in ["r", "--script"]:
                names['run-file'] = [argv[i]]

            i += 1
            add_name(names, 'loads', "load")
            add_name(names, 'load_arguments', argv[i])
            retval = RETURN_OK

        elif argv[i] in ["-t", "--require", "-u", "--require-script"]:
            if to <= i + 1 or argv[i+1] in all_opts:
                print "missing argument after %s" % argv[i]
                retval = MISSING_ARG
                break

            if argv[i] in ["-u", "--require-script"]:
                names['run-file'] = [argv[i]]

            i += 1
            add_name(names, 'loads', "file")
            add_name(names, 'load_arguments', argv[i])

            config['no-lib'] = True
            retval = RETURN_OK

        elif argv[i] in ["-l", "--lib"]:
            if to <= i + 1 or argv[i+1] in all_opts:
                print "missing argument after %s" % argv[i]
                retval = MISSING_ARG
                break

            i += 1
            add_name(names, 'loads', "lib")
            add_name(names, 'load_arguments', argv[i])

            config['no-lib'] = True
            retval = RETURN_OK

        elif argv[i] in ["-g", "-gg", "--eval-json"]:
            if argv[i] == "-gg":
                config['stop'] = True
            if to <= i + 1 or argv[i+1] in all_opts:
                print "missing argument after %s" % argv[i]
                retval = MISSING_ARG
                break

            i += 1
            add_name(names, 'json-linklets', argv[i])

            config['no-lib'] = True
            retval = RETURN_OK

        #########################
        # Interaction Options
        #########################

        elif argv[i] in ["-i", "--repl"]:
            config['repl'] = True
            config['version'] = True
            retval = RETURN_OK

        elif argv[i] in ["-n", "--no-lib"]:
            config['no-lib'] = True
            if retval == INIT:
                retval = JUST_EXIT

        elif argv[i] in ["-v", "--version"]:
            config['version'] = True

        #########################
        # Configuration Options
        #########################

        elif argv[i] == "--kernel":
            config['just_kernel'] = True
            config['no-lib'] = True
            retval = RETURN_OK

        elif argv[i] == "-I":
            if to <= i + 1 or argv[i+1] in all_opts:
                print "missing argument after -I"
                retval = MISSING_ARG
                break
            i += 1
            add_name(names, 'init-lib', argv[i])

        elif argv[i] in ["-X", "--collects"]:
            if to <= i + 1 or argv[i+1] in all_opts:
                print "missing argument after %s" % argv[i]
                retval = MISSING_ARG
                break
            i += 1
            add_name(names, 'set-collects-dir', argv[i])

        elif argv[i] in ["-G", "--config"]:
            if to <= i + 1 or argv[i+1] in all_opts:
                print "missing argument after %s" % argv[i]
                retval = MISSING_ARG
                break
            i += 1
            add_name(names, 'set-config-dir', argv[i])

        elif argv[i] in ["-A", "--addon"]:
            if to <= i + 1 or argv[i+1] in all_opts:
                print "missing argument after %s" % argv[i]
                retval = MISSING_ARG
                break
            i += 1
            add_name(names, 'set-addon-dir', argv[i])

        elif argv[i] in ["-N", "--name"]:
            if to <= i + 1 or argv[i+1] in all_opts:
                print "missing argument after %s" % argv[i]
                retval = MISSING_ARG
                break
            i += 1
            add_name(names, 'set-run-file', argv[i])

        elif argv[i] == '--save-callgraph':
            config['save-callgraph'] = True

        #########################
        # Meta Options
        #########################

        elif argv[i] == "--verbose":
            config['verbose'] = True
            add_name(names, 'verbosity_level', '0')
            if to > i + 1 and not is_rkt_file(argv[i+1]) and argv[i+1] not in all_opts:
                vl = argv[i+1]
                if vl not in verbosity_levels:
                    print("--verbose <level> can only be one of : %s" % verbosity_levels)
                    retval = MISSING_ARG
                    break
                add_name(names, 'verbosity_level', vl, replace=True)
                i += 1

        elif argv[i] == "--jit":
            if to <= i + 1:
                print "missing argument after --jit"
                retval = 2
                break
            i += 1
            jitarg = argv[i]
            jit.set_user_param(None, jitarg)

        elif argv[i] == "--":
            i += 1
            break

        elif argv[i] in ["-h", "--help", "/?", "-?", "/h", "/help"]:
            print_help(argv)
            if retval == INIT:
                retval = JUST_EXIT

        else:
            if '.rkt' in argv[i]:
                add_name(names, 'req-file', argv[i])
                config['no-lib'] = True

                i += 1
                retval = RETURN_OK
            else:
                print "Bad switch : %s" % argv[i]
                retval = MISSING_ARG
            break
        i += 1

    if retval == INIT:
        config['repl'] = True
        config['version'] = True
        retval = RETURN_OK

    if i <= to: #pending args
        args = argv[i:to]

    return config, names, args, retval

# EOF
<commit_msg>any positive integer can be given as verbosity level, all the console_logs with lesser verbosity levels will appear<commit_after>#! /usr/bin/env python
# -*- coding: utf-8 -*-
#
from rpython.rlib import jit

def print_help(argv):
    print """Welcome to Pycket.
%s [<option> ...] <argument> ...
 File and expression options:
  -e <exprs>, --eval <exprs> : Evaluate <exprs>, prints results
  -f <file>, --load <file> : Like -e '(load "<file>")' without printing
  -t <file>, --require <file> : Like -e '(require (file "<file>"))'
  -l <path>, --lib <path> : Like -e '(require (lib "<path>"))'
  -r <file>, --script <file> : Same as -f <file> -N <file> --
  -u <file>, --require-script <file> : Same as -t <file> -N <file> --
  -g <json>, --eval-json <json> : Load and instantiate a linklet json
  -gg <json> : Like -g but stops after instantiating

 Interaction options:
  -i, --repl : Run interactive read-eval-print loop; implies -v
  -n, --no-lib : Skip `(require (lib "<init-lib>"))' for -i/-e/-f/-r
  -v, --version : Show version

 Configuration options:
  --kernel : ignore everything, only load the #%%kernel (for development purposes)
  -I <path> : Set <init-lib> to <path> (sets language)
  -X <dir>, --collects <dir> : Main collects at <dir> (or "" disables all)
  -G <dir>, --config <dir> : Main configuration directory at <dir>
  -A <dir>, --addon <dir> : Addon directory at <dir>
  -N <file>, --name <file> : Sets `(find-system-path 'run-file)' to <file>
  --save-callgraph : save the jit output

 Meta options:
  --verbose <level> : Print the debug logs. <level> : natural number (defaults to 0)
  --jit <jitargs> : Set RPython JIT options may be 'default', 'off',
                    or 'param=value,param=value' list
  -- : No argument following this switch is used as a switch
  -h, --help : Show this information and exits, ignoring other options
Default options:
 If only an argument is specified, it is loaded and evaluated
""" % (argv[0])

file_expr_opts = ["-e", "--eval",
                  "-f", "--load",
                  "-t", "--require",
                  "-l", "--lib",
                  "-r", "--script",
                  "-u", "--require-script",
                  "-g", "--eval-json",
                  "-gg"]
inter_opts = ["-i", "--repl",
              "-n", "--no-lib",
              "-v", "--version"]
conf_opts = ["--kernel",
             "-I",
             "-X", "--collects",
             "-G", "--config",
             "-A", "--addon",
             "-N", "--name",
             "--save-callgraph"]
meta_opts = ["--verbose", "--jit", "-h"]

all_opts = file_expr_opts + inter_opts + conf_opts + meta_opts

def add_name(names, name, val, replace=False):
    if name not in names or replace:
        names[name] = [val]
    else:
        names[name].append(val)

def is_rkt_file(name_str):
    return ".rkt" in name_str

def parse_args(argv):
    INIT = -1
    RETURN_OK = 0
    MISSING_ARG = 5
    JUST_EXIT = 3

    config = {
        'repl' : False,
        'no-lib' : False,
        'version' : False,
        'stop' : False,
        'just_kernel' : False,
        'verbose' : False
    }
    names = {
        # 'file': "",
        # 'exprs': "",
    }
    args = []
    retval = INIT
    i = 1
    to = len(argv)
    while i < to:
        if False:
            pass

        elif argv[i] in ["-e", "--eval"]:
            if to <= i + 1 or argv[i+1] in all_opts:
                print "missing argument after %s" % argv[i]
                retval = MISSING_ARG
                break

            i += 1
            add_name(names, 'loads', "eval")
            add_name(names, 'load_arguments', argv[i])
            retval = RETURN_OK

        elif argv[i] in ["-f", "--load", "-r", "--script"]:
            if to <= i + 1 or argv[i+1] in all_opts:
                print "missing argument after %s" % argv[i]
                retval = MISSING_ARG
                break

            if argv[i] in ["r", "--script"]:
                names['run-file'] = [argv[i]]

            i += 1
            add_name(names, 'loads', "load")
            add_name(names, 'load_arguments', argv[i])
            retval = RETURN_OK

        elif argv[i] in ["-t", "--require", "-u", "--require-script"]:
            if to <= i + 1 or argv[i+1] in all_opts:
                print "missing argument after %s" % argv[i]
                retval = MISSING_ARG
                break

            if argv[i] in ["-u", "--require-script"]:
                names['run-file'] = [argv[i]]

            i += 1
            add_name(names, 'loads', "file")
            add_name(names, 'load_arguments', argv[i])

            config['no-lib'] = True
            retval = RETURN_OK

        elif argv[i] in ["-l", "--lib"]:
            if to <= i + 1 or argv[i+1] in all_opts:
                print "missing argument after %s" % argv[i]
                retval = MISSING_ARG
                break

            i += 1
            add_name(names, 'loads', "lib")
            add_name(names, 'load_arguments', argv[i])

            config['no-lib'] = True
            retval = RETURN_OK

        elif argv[i] in ["-g", "-gg", "--eval-json"]:
            if argv[i] == "-gg":
                config['stop'] = True
            if to <= i + 1 or argv[i+1] in all_opts:
                print "missing argument after %s" % argv[i]
                retval = MISSING_ARG
                break

            i += 1
            add_name(names, 'json-linklets', argv[i])

            config['no-lib'] = True
            retval = RETURN_OK

        #########################
        # Interaction Options
        #########################

        elif argv[i] in ["-i", "--repl"]:
            config['repl'] = True
            config['version'] = True
            retval = RETURN_OK

        elif argv[i] in ["-n", "--no-lib"]:
            config['no-lib'] = True
            if retval == INIT:
                retval = JUST_EXIT

        elif argv[i] in ["-v", "--version"]:
            config['version'] = True

        #########################
        # Configuration Options
        #########################

        elif argv[i] == "--kernel":
            config['just_kernel'] = True
            config['no-lib'] = True
            retval = RETURN_OK

        elif argv[i] == "-I":
            if to <= i + 1 or argv[i+1] in all_opts:
                print "missing argument after -I"
                retval = MISSING_ARG
                break
            i += 1
            add_name(names, 'init-lib', argv[i])

        elif argv[i] in ["-X", "--collects"]:
            if to <= i + 1 or argv[i+1] in all_opts:
                print "missing argument after %s" % argv[i]
                retval = MISSING_ARG
                break
            i += 1
            add_name(names, 'set-collects-dir', argv[i])

        elif argv[i] in ["-G", "--config"]:
            if to <= i + 1 or argv[i+1] in all_opts:
                print "missing argument after %s" % argv[i]
                retval = MISSING_ARG
                break
            i += 1
            add_name(names, 'set-config-dir', argv[i])

        elif argv[i] in ["-A", "--addon"]:
            if to <= i + 1 or argv[i+1] in all_opts:
                print "missing argument after %s" % argv[i]
                retval = MISSING_ARG
                break
            i += 1
            add_name(names, 'set-addon-dir', argv[i])

        elif argv[i] in ["-N", "--name"]:
            if to <= i + 1 or argv[i+1] in all_opts:
                print "missing argument after %s" % argv[i]
                retval = MISSING_ARG
                break
            i += 1
            add_name(names, 'set-run-file', argv[i])

        elif argv[i] == '--save-callgraph':
            config['save-callgraph'] = True

        #########################
        # Meta Options
        #########################

        elif argv[i] == "--verbose":
            config['verbose'] = True
            add_name(names, 'verbosity_level', '0')
            if to > i + 1 and not is_rkt_file(argv[i+1]) and argv[i+1] not in all_opts:
                vl = argv[i+1]
                if not vl.isdigit() or int(vl) < 0:
                    print("--verbose <level> can only be a natural number, given : %s" % vl)
                    retval = MISSING_ARG
                    break
                add_name(names, 'verbosity_level', vl, replace=True)
                i += 1

        elif argv[i] == "--jit":
            if to <= i + 1:
                print "missing argument after --jit"
                retval = 2
                break
            i += 1
            jitarg = argv[i]
            jit.set_user_param(None, jitarg)

        elif argv[i] == "--":
            i += 1
            break

        elif argv[i] in ["-h", "--help", "/?", "-?", "/h", "/help"]:
            print_help(argv)
            if retval == INIT:
                retval = JUST_EXIT

        else:
            if '.rkt' in argv[i]:
                add_name(names, 'req-file', argv[i])
                config['no-lib'] = True

                i += 1
                retval = RETURN_OK
            else:
                print "Bad switch : %s" % argv[i]
                retval = MISSING_ARG
            break
        i += 1

    if retval == INIT:
        config['repl'] = True
        config['version'] = True
        retval = RETURN_OK

    if i <= to: #pending args
        args = argv[i:to]

    return config, names, args, retval

# EOF
<|endoftext|>