<commit_before># vim: set et ts=4 sw=4 fileencoding=utf-8:
'''
yalp_grok.yalp_grok
===================
'''
import os
import sys
import copy
from collections import namedtuple
import regex as re

DEFAULT_PATTERNS_DIRS = [os.path.join(os.path.dirname(os.path.abspath(__file__)), 'patterns')]

# Attributes used to determine type on groupdict post processing
# Just covering floats and ints
INT = ('INT', 'POSINT', 'NONNEGINT')
FLOAT = ('BASE10NUM')
TYPES = {'int': INT, 'float': FLOAT}

# GROK pattern/format data abstracted from logic
NAMED_PATTERN = {'pattern': r'%{(\w+):(\w+)(:\w+)?}', 'format': '(?P<{key}>{regex})'}
UNNAMED_PATTERN = {'pattern': r'%{(\w+)}', 'format': '({regex})'}
PATTERN_FORMATS = (UNNAMED_PATTERN, NAMED_PATTERN)

# Magic pattern for capturing named grok keys in various states of expansion to
# identify final grok key of named attributes
NAMED_KEY_PATTERN = re.compile(
    r'(?:%{(?P<grok>\w+):(?P<key>\w+)(?::(?P<type>\w+))?}'
    r'|\(\?P<(?P<key>\w+)>(?:%{(?P<grok>\w+)}'
    r'|\(\?:%{(?P<grok>\w+)}\))\))'
)

Pattern = namedtuple('Pattern', 'name regex_str')

def grok_match(text, pattern, custom_patterns=None, custom_patterns_dir=None, auto_map=False):
    '''
    Search for pattern in text.

    If text is matched with pattern, return variable names specified
    (%{pattern:variable name}) in pattern and their corresponding
    values. If not matched, return None. custom patterns can be passed
    in by custom_patterns(pattern name, pattern regular expression pair)
    or custom_patterns_dir.

    Data type conversion supported. For example %{NUMBER:num:int} which
    converts the num semantic from a string to an integer. Currently the
    only supported conversions are int and float.

    If auto_map set to True then GROK key will be used to auto determine
    data type conversion. For example %{NUMBER:num} which converts the
    num semantic from a string to a float. A data type defined in grok
    pattern will take precedence over auto determined type.

    If type conversion fails then value left as a string.
    '''
    return grok_search(text, compile_pattern(
        pattern, custom_patterns, custom_patterns_dir, auto_map=auto_map))

def compile_pattern(pattern, custom_patterns=None, custom_patterns_dir=None, auto_map=False):
    '''
    Compile pattern before use for better performance when matching.

    Returns a regex string that can be passed as a pattern to
    grok_search() for matching. Custom patterns can be passed in by
    custom_patterns(pattern name, pattern regular expression pair) or
    custom_patterns_dir, and will then be used in addition to the
    built-in ones.
    '''
    all_patterns = copy.deepcopy(PREDEFINED_PATTERNS)

    custom_pats = {}

    if custom_patterns_dir is not None:
        custom_pats = _reload_patterns([custom_patterns_dir])

    if custom_patterns:
        for pat_name, regex_str in custom_patterns.items():
            custom_pats[pat_name] = Pattern(pat_name, regex_str)

    all_patterns.update(custom_pats)
    py_regex_pattern = pattern
    type_map = _map_types(py_regex_pattern, auto_map)

    while True:

        # replace %{pattern_name:custom_name} with regex and regex group name
        # replace %{pattern_name} with regex
        py_regex_pattern, type_map_temp = _sub_pattern_name(
            py_regex_pattern, all_patterns, type_map, auto_map)
        type_map.update(type_map_temp)

        if re.search(r'%{\w+(:\w+)?}', py_regex_pattern) is None:
            break

    return re.compile(py_regex_pattern), type_map

def grok_search(text, pattern):
    '''
    Search for pattern in text.

    Return dictionary with named fields in pattern as keys, or None if
    no match found.
    '''
    match_obj = pattern[0].search(text)

    if match_obj is not None:
        match_dict = match_obj.groupdict()
        return _apply_map(match_dict, pattern[1])
    return None

def _reload_patterns(patterns_dirs):
    '''
    Load patters from all files in a directory.
    '''
    all_patterns = {}
    for dir_ in patterns_dirs:
        for pat_file in os.listdir(dir_):
            patterns = _load_patterns_from_file(os.path.join(dir_, pat_file))
            all_patterns.update(patterns)

    return all_patterns

def _load_patterns_from_file(pat_file):
    '''
    Load patterns from a text file.
    '''
    patterns = {}
    with open(pat_file, 'r') as pfh:
        for line in pfh:
            line = line.strip()
            if line == '' or line.startswith('#'):
                continue

            sep = line.find(' ')
            pat_name = line[:sep]
            regex_str = line[sep:].strip()
            pat = Pattern(pat_name, regex_str)
            patterns[pat.name] = pat

    return patterns

# Substitute keys with regex
def _sub_pattern_name(py_regex_pattern, all_patterns, type_map, auto_map):
    for pattern_format in PATTERN_FORMATS:
        pattern = pattern_format['pattern']
        py_regex_pattern = re.sub(
            pattern, lambda match, pat_format=pattern_format: _format_pattern_name(
                match, pat_format, all_patterns),
            py_regex_pattern)
        type_map.update(_map_types(py_regex_pattern, auto_map, type_map))

    return py_regex_pattern, type_map

# Expand keys with regex using generic string format
def _format_pattern_name(match, pattern_format, all_patterns):
    return pattern_format['format'].format(
        key=_get_group_key(match),
        regex=all_patterns[match.group(1)].regex_str
    )

# Return None if no more than 1 index in group
# Probably a better way to do this...
def _get_group_key(match):
    try:
        return match.group(2)
    except IndexError:
        return None

# Follow conditions in correct order to assign data type for a named Grok key.
# For now this will only affect Grok keys listed in INT and FLOAT variables
# in order to assign int or float primitive types to attribute names using
# matching Grok keys. A defined type using ES Grok type semantic takes precedence.
def _map_types(py_regex_pattern, auto_map, type_map=None):
    if not type_map:
        type_map = {}

    named_grok_keys = NAMED_KEY_PATTERN.findall(py_regex_pattern)
    for grok_key, name, defined_type in named_grok_keys:
        if auto_map:
            detected_type = _type_match(grok_key)
            if detected_type and name not in type_map.keys():
                type_map[name] = detected_type

        if defined_type in TYPES.keys():
            type_map[name] = defined_type

    return type_map

def _type_match(grok_key):
    for type_item, grok_keys in TYPES.items():
        if grok_key in grok_keys:
            return type_item
    return None

# Cast values of any attribute names inside map to their identified type.
def _apply_map(match_dict, type_map):
    if type_map:
        for name, detected_type in type_map.items():
            try:
                match_dict[name] = _convert(match_dict[name], detected_type)
            except ValueError:
                sys.exc_clear()
    return match_dict

# The power of Christ compels you!!!
def _convert(value, detected_type):
    if value:
        if detected_type == 'int':
            return int(value)
        if detected_type == 'float':
            return float(value)
    return value

PREDEFINED_PATTERNS = _reload_patterns(DEFAULT_PATTERNS_DIRS)
<commit_msg>Update yalp_grok.py<commit_after># vim: set et ts=4 sw=4 fileencoding=utf-8:
'''
yalp_grok.yalp_grok
===================
'''
import os
import sys
import copy
from collections import namedtuple
import regex as re

DEFAULT_PATTERNS_DIRS = [
    os.path.join(os.path.dirname(os.path.abspath(__file__)), 'patterns')
]

# Attributes used to determine type on groupdict post processing
# Just covering floats and ints
INT = ('INT', 'POSINT', 'NONNEGINT')
FLOAT = ('BASE10NUM')
TYPES = {'int': INT, 'float': FLOAT}

# GROK pattern/format data abstracted from logic
NAMED_PATTERN = {
    'pattern': r'%{(\w+):(\w+)(:\w+)?}',
    'format': '(?P<{key}>{regex})'
}

UNNAMED_PATTERN = {'pattern': r'%{(\w+)}', 'format': '({regex})'}
PATTERN_FORMATS = (UNNAMED_PATTERN, NAMED_PATTERN)

# Magic pattern for capturing named grok keys in various states
# of expansion to identify final grok key of named attributes
NAMED_KEY_PATTERN = re.compile(
    r'(?:%{(?P<grok>\w+):(?P<key>\w+)(?::(?P<type>\w+))?}'
    r'|\(\?P<(?P<key>\w+)>(?:%{(?P<grok>\w+)}'
    r'|\(\?:%{(?P<grok>\w+)}\))\))'
)

Pattern = namedtuple('Pattern', 'name regex_str')

def grok_match(text, pattern, custom_patterns=None,
               custom_patterns_dir=None, auto_map=False):
    '''
    Search for pattern in text.

    If text is matched with pattern, return variable names specified
    (%{pattern:variable name}) in pattern and their corresponding
    values. If not matched, return None. custom patterns can be passed
    in by custom_patterns(pattern name, pattern regular expression pair)
    or custom_patterns_dir.

    Data type conversion supported. For example %{NUMBER:num:int} which
    converts the num semantic from a string to an integer. Currently the
    only supported conversions are int and float.

    If auto_map set to True then GROK key will be used to auto determine
    data type conversion. For example %{NUMBER:num} which converts the
    num semantic from a string to a float. A data type defined in grok
    pattern will take precedence over auto determined type.

    If type conversion fails then value left as a string.
    '''
    return grok_search(text, compile_pattern(
        pattern, custom_patterns, custom_patterns_dir, auto_map=auto_map))

def compile_pattern(pattern, custom_patterns=None,
                    custom_patterns_dir=None, auto_map=False):
    '''
    Compile pattern before use for better performance when matching.

    Returns a regex string that can be passed as a pattern to
    grok_search() for matching. Custom patterns can be passed in by
    custom_patterns(pattern name, pattern regular expression pair) or
    custom_patterns_dir, and will then be used in addition to the
    built-in ones.
    '''
    all_patterns = copy.deepcopy(PREDEFINED_PATTERNS)

    custom_pats = {}

    if custom_patterns_dir is not None:
        custom_pats = _reload_patterns([custom_patterns_dir])

    if custom_patterns:
        for pat_name, regex_str in custom_patterns.items():
            custom_pats[pat_name] = Pattern(pat_name, regex_str)

    all_patterns.update(custom_pats)
    py_regex_pattern = pattern
    type_map = _map_types(py_regex_pattern, auto_map)

    while True:

        # replace %{pattern_name:custom_name} with regex and regex group name
        # replace %{pattern_name} with regex
        py_regex_pattern, type_map_temp = _sub_pattern_name(
            py_regex_pattern, all_patterns, type_map, auto_map)
        type_map.update(type_map_temp)

        if re.search(r'%{\w+(:\w+)?}', py_regex_pattern) is None:
            break

    return re.compile(py_regex_pattern), type_map

def grok_search(text, pattern):
    '''
    Search for pattern in text.

    Return dictionary with named fields in pattern as keys, or None if
    no match found.
    '''
    match_obj = pattern[0].search(text)

    if match_obj is not None:
        match_dict = match_obj.groupdict()
        return _apply_map(match_dict, pattern[1])
    return None

def _reload_patterns(patterns_dirs):
    '''
    Load patters from all files in a directory.
    '''
    all_patterns = {}
    for dir_ in patterns_dirs:
        for pat_file in os.listdir(dir_):
            patterns = _load_patterns_from_file(os.path.join(dir_, pat_file))
            all_patterns.update(patterns)

    return all_patterns

def _load_patterns_from_file(pat_file):
    '''
    Load patterns from a text file.
    '''
    patterns = {}
    with open(pat_file, 'r') as pfh:
        for line in pfh:
            line = line.strip()
            if line == '' or line.startswith('#'):
                continue

            sep = line.find(' ')
            pat_name = line[:sep]
            regex_str = line[sep:].strip()
            pat = Pattern(pat_name, regex_str)
            patterns[pat.name] = pat

    return patterns

# Substitute keys with regex
def _sub_pattern_name(py_regex_pattern, all_patterns, type_map, auto_map):
    '''
    Expand grok pattern with regex substitution
    '''
    for pattern_format in PATTERN_FORMATS:
        pattern = pattern_format['pattern']
        py_regex_pattern = re.sub(
            pattern, lambda match, pat_format=pattern_format:
            _format_pattern_name(
                match, pat_format, all_patterns),
            py_regex_pattern)
        type_map.update(_map_types(py_regex_pattern, auto_map, type_map))

    return py_regex_pattern, type_map

# Expand keys with regex using generic string format
def _format_pattern_name(match, pattern_format, all_patterns):
    '''
    Generate a new pattern based on capture results
    '''
    return pattern_format['format'].format(
        key=_get_group_key(match),
        regex=all_patterns[match.group(1)].regex_str
    )

# Return None if no more than 1 index in group
# Probably a better way to do this...
def _get_group_key(match):
    '''
    Return None on IndexError
    '''
    try:
        return match.group(2)
    except IndexError:
        return None

# Follow conditions in correct order to assign data type for a named Grok key.
# For now this will only affect Grok keys listed in INT and FLOAT variables
# in order to assign int or float primitive types to
# attribute names using matching Grok keys.
# A defined type using ES Grok type semantic takes precedence.
def _map_types(py_regex_pattern, auto_map, type_map=None):
    '''
    Generate type map against regex pattern
    '''
    if not type_map:
        type_map = {}

    named_grok_keys = NAMED_KEY_PATTERN.findall(py_regex_pattern)
    for grok_key, name, defined_type in named_grok_keys:
        if auto_map:
            detected_type = _type_match(grok_key)
            if detected_type and name not in type_map.keys():
                type_map[name] = detected_type

        if defined_type in TYPES.keys():
            type_map[name] = defined_type

    return type_map

def _type_match(grok_key):
    '''
    Attempt to match grok key with types associated with it
    '''
    for type_item, grok_keys in TYPES.items():
        if grok_key in grok_keys:
            return type_item
    return None

# Cast values of any attribute names inside map to their identified type.
def _apply_map(match_dict, type_map):
    '''
    Apply generated type map to regex group dict
    '''
    if type_map:
        for name, detected_type in type_map.items():
            try:
                match_dict[name] = _convert(match_dict[name], detected_type)
            except ValueError:
                sys.exc_clear()
    return match_dict

# The power of Christ compels you!!!
def _convert(value, detected_type):
    '''
    Attempts conversion if value is not None.
    '''
    if value:
        if detected_type == 'int':
            return int(value)
        if detected_type == 'float':
            return float(value)
    return value

PREDEFINED_PATTERNS = _reload_patterns(DEFAULT_PATTERNS_DIRS)
<|endoftext|>


<commit_before>import _data_operations as dataops
from ._decorators import timeit
from ._exceptions import BadTickersFormat

@timeit
def current(tickers, write_to_csv=False, result_csv_path=None):
    """Retrieves realtime stock data from Yahoo Finance.

    Arguments:
        tickers --> list of tickers or path to ticker csv file
        write_to_csv --> boolean, current writes data to a csv file if set to True, returns a pandas dataframe
                         if set to False (default)
        result_csv_path --> string, specifies path to csv file to write the data (overwrites the file)

    Returns:
        a pandas dataframe or None (if write_to_csv is set to False)
    """

    # create parameter string for the request
    api_dict = dataops.read_api_dict()
    param_list = dataops.get_param_list_from_api_dict(api_dict)
    param_string = dataops.get_param_string_from_list(param_list)

    # list of tickers or path to a csv file with tickers?
    try:
        if type(tickers) == str:
            ticker_list = sorted(dataops.get_ticker_list_from_file(tickers))
            ticker_string = dataops.get_ticker_string_from_list(ticker_list)

        elif type(tickers) == list:
            ticker_list = tickers
            ticker_string = dataops.get_ticker_string_from_list(ticker_list)
        else:
            raise BadTickersFormat('Please provide either a csv file or a list of tickers.')
    except BadTickersFormat as err:
        quit(err.message)
    else:
        # make request and create a pandas dataframe from the response
        answer_string = dataops.get_current_answer_string(ticker_string, param_string)
        answer_list = dataops.get_answer_list_from_string(answer_string)
        pandas_dataframe = dataops.current_pd_dataframe(api_dict, answer_list, param_list)

        if write_to_csv:
            pandas_dataframe.to_csv(result_csv_path)

        return pandas_dataframe

@timeit
def historical(ticker, from_date=None, to_date=None):
    """Retrieves historical stock price data from Yahoo Finance.

    Arguments:
        ticker --> one ticker symbol
        from_date --> lower bound for timeframe
        to_date --> upper bound for timeframe

        from_date and to_date can be of the following types:
            - string of form 'YYYY-MM-DD'
            - Python datetime
            - NumPy datetime64
            - Pandas Timestamp

    Returns:
        Pandas DataFrame
    """

    answer_string = dataops.get_historical_answer_string(ticker, from_date, to_date)
    answer_list = dataops.get_answer_list_from_string(answer_string)
    pandas_dataframe = dataops.historical_pd_dataframe(answer_list)

    # TODO Finish this
    return pandas_dataframe

@timeit
def correlation_matrix(tickers, heatmap=False):
    """Calculates a correlation matrix for the stocks in ticker_list."""

    try:
        if type(tickers) == str:
            ticker_list = sorted(dataops.get_ticker_list_from_file(tickers))

        elif type(tickers) == list:
            ticker_list = tickers
        else:
            raise BadTickersFormat('Please provide either a csv file or a list of tickers.')
    except BadTickersFormat as err:
        quit(err.message)
    else:

        # iterate through tuples of tickers and corresponding dataframes
        full_dfs = zip(ticker_list, map(historical, ticker_list))

        # throw out the Nones
        not_nones = filter(lambda (ticker, df): df is not None, full_dfs)

        # use only the 'Close' column
        close_only = [(ticker, df[['Close']]) for ticker, df in not_nones]

        # replace 'Close' in all dataframes with their tickers
        renamed = [df.rename(columns={'Close': ticker}) for ticker, df in close_only]

        # inner-join all the dataframes by index ('Date')
        joined = reduce(lambda df1, df2: df1.join(df2, how='inner'), renamed)

        corrmat = joined.corr()

        if heatmap:
            import seaborn as sns
            # Draw the heatmap using seaborn
            sns.heatmap(corrmat, vmax=.8, square=True)

            sns.plt.show()

        return corrmat
<commit_msg>added write to csv option for historical<commit_after>import _data_operations as dataops
from ._decorators import timeit
from ._exceptions import BadTickersFormat

@timeit
def current(tickers, write_to_csv=False, result_csv_path=None):
    """Retrieves realtime stock data from Yahoo Finance.

    Arguments:
        tickers --> list of tickers or path to ticker csv file
        write_to_csv --> boolean, current writes data to a csv file if set to True, returns a pandas dataframe
                         if set to False (default)
        result_csv_path --> string, specifies path to csv file to write the data (overwrites the file)

    Returns:
        a pandas dataframe or None (if write_to_csv is set to False)
    """

    # create parameter string for the request
    api_dict = dataops.read_api_dict()
    param_list = dataops.get_param_list_from_api_dict(api_dict)
    param_string = dataops.get_param_string_from_list(param_list)

    # list of tickers or path to a csv file with tickers?
    try:
        if type(tickers) == str:
            ticker_list = sorted(dataops.get_ticker_list_from_file(tickers))
            ticker_string = dataops.get_ticker_string_from_list(ticker_list)

        elif type(tickers) == list:
            ticker_list = tickers
            ticker_string = dataops.get_ticker_string_from_list(ticker_list)
        else:
            raise BadTickersFormat('Please provide either a csv file or a list of tickers.')
    except BadTickersFormat as err:
        quit(err.message)
    else:
        # make request and create a pandas dataframe from the response
        answer_string = dataops.get_current_answer_string(ticker_string, param_string)
        answer_list = dataops.get_answer_list_from_string(answer_string)
        pandas_dataframe = dataops.current_pd_dataframe(api_dict, answer_list, param_list)

        if write_to_csv:
            pandas_dataframe.to_csv(result_csv_path)

        return pandas_dataframe

@timeit
def historical(ticker, from_date=None, to_date=None, write_to_csv=False, result_csv_path=None):
    """Retrieves historical stock price data from Yahoo Finance.

    Arguments:
        ticker --> one ticker symbol
        from_date --> lower bound for timeframe
        to_date --> upper bound for timeframe

        from_date and to_date can be of the following types:
            - string of form 'YYYY-MM-DD'
            - Python datetime
            - NumPy datetime64
            - Pandas Timestamp

    Returns:
        Pandas DataFrame
    """

    answer_string = dataops.get_historical_answer_string(ticker, from_date, to_date)
    answer_list = dataops.get_answer_list_from_string(answer_string)
    pandas_dataframe = dataops.historical_pd_dataframe(answer_list)

    if write_to_csv:
        pandas_dataframe.to_csv(result_csv_path)

    return pandas_dataframe

@timeit
def correlation_matrix(tickers, heatmap=False):
    """Calculates a correlation matrix for the stocks in ticker_list."""

    try:
        if type(tickers) == str:
            ticker_list = sorted(dataops.get_ticker_list_from_file(tickers))

        elif type(tickers) == list:
            ticker_list = tickers
        else:
            raise BadTickersFormat('Please provide either a csv file or a list of tickers.')
    except BadTickersFormat as err:
        quit(err.message)
    else:

        # iterate through tuples of tickers and corresponding dataframes
        full_dfs = zip(ticker_list, map(historical, ticker_list))

        # throw out the Nones
        not_nones = filter(lambda (ticker, df): df is not None, full_dfs)

        # use only the 'Close' column
        close_only = [(ticker, df[['Close']]) for ticker, df in not_nones]

        # replace 'Close' in all dataframes with their tickers
        renamed = [df.rename(columns={'Close': ticker}) for ticker, df in close_only]

        # inner-join all the dataframes by index ('Date')
        joined = reduce(lambda df1, df2: df1.join(df2, how='inner'), renamed)

        corrmat = joined.corr()

        if heatmap:
            import seaborn as sns
            # Draw the heatmap using seaborn
            sns.heatmap(corrmat, vmax=.8, square=True)

            sns.plt.show()

        return corrmat
<|endoftext|>


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

import os
import shutil
import subprocess
import sys
#import tarfile

from inspectShell import archOfBinary, testDbgOrOpt, testJsShellOrXpcshell

p0 = os.path.dirname(__file__)
lithiumpy = os.path.abspath(os.path.join(p0, os.pardir, 'lithium', 'lithium.py'))
autobisectpy = os.path.abspath(os.path.join(p0, os.pardir, 'js-autobisect', 'autoBisect.py'))
shellBeautificationpy = os.path.join(p0, 'shellBeautification.py')

path0 = os.path.dirname(__file__)
path2 = os.path.abspath(os.path.join(path0, os.pardir, 'util'))
sys.path.append(path2)
from subprocesses import captureStdout

def tempdir(path):
    os.mkdir(path)
    return "--tempdir=" + path

def pinpoint(itest, logPrefix, jsEngine, engineFlags, infilename, bisectRepo, alsoRunChar=True, alsoReduceEntireFile=False):
    """
       Run Lithium and autobisect.

       itest must be an array of the form [module, ...] where module is an interestingness module.
       The module's "interesting" function must accept [...] + [jsEngine] + engineFlags + infilename
       (If it's not prepared to accept engineFlags, engineFlags must be empty.)
    """

    valgrindSupport = "--valgrind" in itest
    valgrindX = ["--valgrind"] if valgrindSupport else []

    lithArgs = itest + [jsEngine] + engineFlags + [infilename]
    print ' '.join([lithiumpy] + lithArgs)
    subprocess.call([sys.executable, lithiumpy, tempdir(logPrefix + "-lith1-tmp")] + lithArgs, stdout=open(logPrefix + "-lith1-out", "w"))

    if alsoRunChar:
        lith2Args = ["--char"] + lithArgs
        print ' '.join([lithiumpy] + lith2Args)
        subprocess.call([sys.executable, lithiumpy, tempdir(logPrefix + "-lith2-tmp")] + lith2Args, stdout=open(logPrefix + "-lith2-out", "w"))

    print
    print "Done running Lithium on the part in between DDBEGIN and DDEND. To reproduce, run:"
    print ' '.join([lithiumpy, "--strategy=check-only"] + lithArgs)
    print

    unbeautifiedOutput = captureStdout([sys.executable, lithiumpy, "--strategy=check-only", tempdir(logPrefix + "-lith-b-tmp")] + lithArgs)[0]
    # Check that the testcase is interesting.
    if False and alsoReduceEntireFile and 'not interesting' not in unbeautifiedOutput:
        assert 'interesting' in unbeautifiedOutput
        # Beautify the output. This will remove DDBEGIN and DDEND as they are comments.
        # This will output a file with the '-beautified' suffix.
        # Reduce once using toString decompile method.
        subprocess.call([sys.executable, shellBeautificationpy, '--shell=' + jsEngine, "--decompilationType=uneval", infilename])

        print 'Operating on the beautified testcase for the n-th time where n =',
        # iterNum starts from 3 because lith1 and lith2 are already used above.
        iterNumStart = 3 if alsoRunChar is True else 2
        iterNum = iterNumStart
        # Run Lithium on the testcase 10 more times, but run it using char only for 3 tries of toString and uneval reduction each.
        # Generally, lines don't get significantly reduced after the 3rd try of line reduction.
        MAX_BEAUTIFIED_LITHIUM_RUNS = iterNumStart + 10
        while(iterNum < MAX_BEAUTIFIED_LITHIUM_RUNS):
            print iterNum - 2,
            # Operate on the beautified version first.
            lithArgs = lithArgs[0:-1]
            beautifiedFilename = infilename + '-beautified'
            lithArgs = lithArgs + [beautifiedFilename]
            # We must still be operating on the beautified version.
            assert beautifiedFilename in lithArgs
            # Check that the testcase is still interesting.
            beautifiedOutput = captureStdout([sys.executable, lithiumpy, "--strategy=check-only", tempdir(logPrefix + "-lith" + str(iterNum) + "-b-tmp")] + lithArgs)[0]
            if 'not interesting' not in beautifiedOutput:
                assert 'interesting' in beautifiedOutput
                # Overwrite the original -reduced file with the beautified version since it is interesting.
                shutil.move(beautifiedFilename, infilename)
                # Operate on the original -reduced file.
                lithArgs = lithArgs[0:-1]
                lithArgs = lithArgs + [infilename]
                assert beautifiedFilename not in lithArgs
                #print ' '.join([lithiumpy] + lithArgs)
                subprocess.call([sys.executable, lithiumpy, tempdir(logPrefix + '-lith' + str(iterNum) + '-tmp')] + lithArgs, stdout=open(logPrefix + "-lith" + str(iterNum) + "-out", "w"))

                # Run it using char only after 3 tries of toString and uneval reduction each.
                if alsoRunChar and (iterNum - 2) > ((MAX_BEAUTIFIED_LITHIUM_RUNS - iterNumStart) // 2):
                    iterNum += 1
                    print iterNum - 2,
                    print '(operating on chars..)',
                    assert iterNum in (9, 11, 13)  # Refer to reduction method below
                    lith2Args = ["--char"] + lithArgs
                    assert beautifiedFilename not in lith2Args
                    #print ' '.join([lithiumpy] + lith2Args)
                    subprocess.call([sys.executable, lithiumpy, tempdir(logPrefix + '-lith' + str(iterNum) + '-tmp')] + lith2Args, stdout=open(logPrefix + "-lith" + str(iterNum) + "-out", "w"))
            else:
                # Beautified testcase is no longer interesting.
                print 'Beautified testcase is no longer interesting!'
                break
            iterNum += 1
            # We want the following reduction method:
            # iterNum 3: uneval
            # iterNum 4: toString
            # iterNum 5: uneval
            # iterNum 6: toString
            # iterNum 7: uneval
            # iterNum 8: toString
            # iterNum 9: char
            # iterNum 10: uneval
            # iterNum 11: char
            # iterNum 12: toString
            # iterNum 13: char
            if iterNum < MAX_BEAUTIFIED_LITHIUM_RUNS:
                # This will output a file with the '-beautified' suffix.
                # Rotate between reducing using the toString and uneval decompile method
                if alsoRunChar:
                    if iterNum % 2 == 0 and iterNum != 10:
                        # toString
                        assert iterNum in (4, 6, 8, 12)
                        subprocess.call([sys.executable, shellBeautificationpy, '--shell=' + jsEngine, "--decompilationType=toString", infilename])
                    else:
                        # uneval
                        # iterNum 3 has already occurred prior to the increment of iterNum above.
                        assert iterNum in (5, 7, 10)
                        subprocess.call([sys.executable, shellBeautificationpy, '--shell=' + jsEngine, "--decompilationType=uneval", infilename])
                else:
                    # If alsoRunChar is false, which occurs when jsInteresting.py is operating at JS_DID_NOT_FINISH and below. iterNumStart also starts from 2.
                    if iterNum % 2 != 0:
                        # toString
                        assert iterNum in (3, 5, 7, 9, 11, 13)
                        subprocess.call([sys.executable, shellBeautificationpy, '--shell=' + jsEngine, "--decompilationType=toString", infilename])
                    else:
                        # uneval
                        # iterNum 3 has already occurred prior to the increment of iterNum above.
                        assert iterNum in (4, 6, 8, 10, 12)
                        subprocess.call([sys.executable, shellBeautificationpy, '--shell=' + jsEngine, "--decompilationType=uneval", infilename])
            else:
                print

        # Operate on the original -reduced file.
        lithArgs = lithArgs[0:-1]
        lithArgs = lithArgs + [infilename]
        assert beautifiedFilename not in lithArgs
        # Check that the testcase is still interesting after the extra beautified lithium reductions.
        finalBeautifiedOutput = captureStdout([sys.executable, lithiumpy, "--strategy=check-only", tempdir(logPrefix + "-lith-final-tmp")] + lithArgs)[0]
        if 'not interesting' not in finalBeautifiedOutput:
            assert 'interesting' in finalBeautifiedOutput

        # Archive all wXX-lith*-tmp directories in a tarball.
        #lithTmpDirTarball = tarfile.open('tempName.tar.bz2', 'w:bz2')
        #for loop from 1 to iterNum
            #lithTmpDirTarball.add(logPrefix + '-lith' + str(iterNum) + '-tmp')
        #lithTmpDirTarball.close()

    # FIXME: We should do the build slave detection logic later, and maybe somewhere else.
    isBuildSlave = False
    jsEngineName = os.path.basename(jsEngine)
    if bisectRepo is not "none":
        # We cannot test that it is not xpcshell with Python 2.5 since testJsShellOrXpcshell relies
        # on 'delete' being a keyword argument in NamedTemporaryFile(). The testing functions in
        # inspectShell in general need at least Python 2.6 because of this.
        if not isBuildSlave and sys.version_info >= (2, 6) and testJsShellOrXpcshell(jsEngine) != "xpcshell":
            autobisectCmd = [sys.executable, autobisectpy] + valgrindX + ["-d", bisectRepo, "-i", "-p", "-a", archOfBinary(jsEngine), "-c", testDbgOrOpt(jsEngine), "--flags=," + ','.join(engineFlags)] + [infilename] + itest
            print ' '.join(autobisectCmd)
            subprocess.call(autobisectCmd, stdout=open(logPrefix + "-autobisect", "w"), stderr=subprocess.STDOUT)
            print "Done running autobisect. Log: " + logPrefix + "-autobisect"
        elif not isBuildSlave and sys.version_info < (2, 6):
            print 'Not pinpointing to exact changeset, please use a Python version >= 2.6.'
<commit_msg>Pinpoint.py should point to autoBisect in autoBisectJs directory and not the old one.<commit_after>#!/usr/bin/env python

import os
import shutil
import subprocess
import sys
#import tarfile

from inspectShell import archOfBinary, testDbgOrOpt, testJsShellOrXpcshell

p0 = os.path.dirname(__file__)
lithiumpy = os.path.abspath(os.path.join(p0, os.pardir, 'lithium', 'lithium.py'))
autobisectpy = os.path.abspath(os.path.join(p0, os.pardir, 'autoBisectJs', 'autoBisect.py'))
shellBeautificationpy = os.path.join(p0, 'shellBeautification.py')

path0 = os.path.dirname(__file__)
path2 = os.path.abspath(os.path.join(path0, os.pardir, 'util'))
sys.path.append(path2)
from subprocesses import captureStdout

def tempdir(path):
    os.mkdir(path)
    return "--tempdir=" + path

def pinpoint(itest, logPrefix, jsEngine, engineFlags, infilename, bisectRepo, alsoRunChar=True, alsoReduceEntireFile=False):
    """
       Run Lithium and autobisect.

       itest must be an array of the form [module, ...] where module is an interestingness module.
       The module's "interesting" function must accept [...] + [jsEngine] + engineFlags + infilename
       (If it's not prepared to accept engineFlags, engineFlags must be empty.)
    """

    valgrindSupport = "--valgrind" in itest
    valgrindX = ["--valgrind"] if valgrindSupport else []

    lithArgs = itest + [jsEngine] + engineFlags + [infilename]
    print ' '.join([lithiumpy] + lithArgs)
    subprocess.call([sys.executable, lithiumpy, tempdir(logPrefix + "-lith1-tmp")] + lithArgs, stdout=open(logPrefix + "-lith1-out", "w"))

    if alsoRunChar:
        lith2Args = ["--char"] + lithArgs
        print ' '.join([lithiumpy] + lith2Args)
        subprocess.call([sys.executable, lithiumpy, tempdir(logPrefix + "-lith2-tmp")] + lith2Args, stdout=open(logPrefix + "-lith2-out", "w"))

    print
    print "Done running Lithium on the part in between DDBEGIN and DDEND. To reproduce, run:"
    print ' '.join([lithiumpy, "--strategy=check-only"] + lithArgs)
    print

    unbeautifiedOutput = captureStdout([sys.executable, lithiumpy, "--strategy=check-only", tempdir(logPrefix + "-lith-b-tmp")] + lithArgs)[0]
    # Check that the testcase is interesting.
    if False and alsoReduceEntireFile and 'not interesting' not in unbeautifiedOutput:
        assert 'interesting' in unbeautifiedOutput
        # Beautify the output. This will remove DDBEGIN and DDEND as they are comments.
        # This will output a file with the '-beautified' suffix.
        # Reduce once using toString decompile method.
        subprocess.call([sys.executable, shellBeautificationpy, '--shell=' + jsEngine, "--decompilationType=uneval", infilename])

        print 'Operating on the beautified testcase for the n-th time where n =',
        # iterNum starts from 3 because lith1 and lith2 are already used above.
        iterNumStart = 3 if alsoRunChar is True else 2
        iterNum = iterNumStart
        # Run Lithium on the testcase 10 more times, but run it using char only for 3 tries of toString and uneval reduction each.
        # Generally, lines don't get significantly reduced after the 3rd try of line reduction.
        MAX_BEAUTIFIED_LITHIUM_RUNS = iterNumStart + 10
        while(iterNum < MAX_BEAUTIFIED_LITHIUM_RUNS):
            print iterNum - 2,
            # Operate on the beautified version first.
            lithArgs = lithArgs[0:-1]
            beautifiedFilename = infilename + '-beautified'
            lithArgs = lithArgs + [beautifiedFilename]
            # We must still be operating on the beautified version.
            assert beautifiedFilename in lithArgs
            # Check that the testcase is still interesting.
            beautifiedOutput = captureStdout([sys.executable, lithiumpy, "--strategy=check-only", tempdir(logPrefix + "-lith" + str(iterNum) + "-b-tmp")] + lithArgs)[0]
            if 'not interesting' not in beautifiedOutput:
                assert 'interesting' in beautifiedOutput
                # Overwrite the original -reduced file with the beautified version since it is interesting.
                shutil.move(beautifiedFilename, infilename)
                # Operate on the original -reduced file.
                lithArgs = lithArgs[0:-1]
                lithArgs = lithArgs + [infilename]
                assert beautifiedFilename not in lithArgs
                #print ' '.join([lithiumpy] + lithArgs)
                subprocess.call([sys.executable, lithiumpy, tempdir(logPrefix + '-lith' + str(iterNum) + '-tmp')] + lithArgs, stdout=open(logPrefix + "-lith" + str(iterNum) + "-out", "w"))

                # Run it using char only after 3 tries of toString and uneval reduction each.
                if alsoRunChar and (iterNum - 2) > ((MAX_BEAUTIFIED_LITHIUM_RUNS - iterNumStart) // 2):
                    iterNum += 1
                    print iterNum - 2,
                    print '(operating on chars..)',
                    assert iterNum in (9, 11, 13)  # Refer to reduction method below
                    lith2Args = ["--char"] + lithArgs
                    assert beautifiedFilename not in lith2Args
                    #print ' '.join([lithiumpy] + lith2Args)
                    subprocess.call([sys.executable, lithiumpy, tempdir(logPrefix + '-lith' + str(iterNum) + '-tmp')] + lith2Args, stdout=open(logPrefix + "-lith" + str(iterNum) + "-out", "w"))
            else:
                # Beautified testcase is no longer interesting.
                print 'Beautified testcase is no longer interesting!'
                break
            iterNum += 1
            # We want the following reduction method:
            # iterNum 3: uneval
            # iterNum 4: toString
            # iterNum 5: uneval
            # iterNum 6: toString
            # iterNum 7: uneval
            # iterNum 8: toString
            # iterNum 9: char
            # iterNum 10: uneval
            # iterNum 11: char
            # iterNum 12: toString
            # iterNum 13: char
            if iterNum < MAX_BEAUTIFIED_LITHIUM_RUNS:
                # This will output a file with the '-beautified' suffix.
                # Rotate between reducing using the toString and uneval decompile method
                if alsoRunChar:
                    if iterNum % 2 == 0 and iterNum != 10:
                        # toString
                        assert iterNum in (4, 6, 8, 12)
                        subprocess.call([sys.executable, shellBeautificationpy, '--shell=' + jsEngine, "--decompilationType=toString", infilename])
                    else:
                        # uneval
                        # iterNum 3 has already occurred prior to the increment of iterNum above.
                        assert iterNum in (5, 7, 10)
                        subprocess.call([sys.executable, shellBeautificationpy, '--shell=' + jsEngine, "--decompilationType=uneval", infilename])
                else:
                    # If alsoRunChar is false, which occurs when jsInteresting.py is operating at JS_DID_NOT_FINISH and below. iterNumStart also starts from 2.
                    if iterNum % 2 != 0:
                        # toString
                        assert iterNum in (3, 5, 7, 9, 11, 13)
                        subprocess.call([sys.executable, shellBeautificationpy, '--shell=' + jsEngine, "--decompilationType=toString", infilename])
                    else:
                        # uneval
                        # iterNum 3 has already occurred prior to the increment of iterNum above.
                        assert iterNum in (4, 6, 8, 10, 12)
                        subprocess.call([sys.executable, shellBeautificationpy, '--shell=' + jsEngine, "--decompilationType=uneval", infilename])
            else:
                print

        # Operate on the original -reduced file.
        lithArgs = lithArgs[0:-1]
        lithArgs = lithArgs + [infilename]
        assert beautifiedFilename not in lithArgs
        # Check that the testcase is still interesting after the extra beautified lithium reductions.
        finalBeautifiedOutput = captureStdout([sys.executable, lithiumpy, "--strategy=check-only", tempdir(logPrefix + "-lith-final-tmp")] + lithArgs)[0]
        if 'not interesting' not in finalBeautifiedOutput:
            assert 'interesting' in finalBeautifiedOutput

        # Archive all wXX-lith*-tmp directories in a tarball.
        #lithTmpDirTarball = tarfile.open('tempName.tar.bz2', 'w:bz2')
        #for loop from 1 to iterNum
            #lithTmpDirTarball.add(logPrefix + '-lith' + str(iterNum) + '-tmp')
        #lithTmpDirTarball.close()

    # FIXME: We should do the build slave detection logic later, and maybe somewhere else.
    isBuildSlave = False
    jsEngineName = os.path.basename(jsEngine)
    if bisectRepo is not "none":
        # We cannot test that it is not xpcshell with Python 2.5 since testJsShellOrXpcshell relies
        # on 'delete' being a keyword argument in NamedTemporaryFile(). The testing functions in
        # inspectShell in general need at least Python 2.6 because of this.
        if not isBuildSlave and sys.version_info >= (2, 6) and testJsShellOrXpcshell(jsEngine) != "xpcshell":
            autobisectCmd = [sys.executable, autobisectpy] + valgrindX + ["-d", bisectRepo, "-i", "-p", "-a", archOfBinary(jsEngine), "-c", testDbgOrOpt(jsEngine), "--flags=," + ','.join(engineFlags)] + [infilename] + itest
            print ' '.join(autobisectCmd)
            subprocess.call(autobisectCmd, stdout=open(logPrefix + "-autobisect", "w"), stderr=subprocess.STDOUT)
            print "Done running autobisect. Log: " + logPrefix + "-autobisect"
        elif not isBuildSlave and sys.version_info < (2, 6):
            print 'Not pinpointing to exact changeset, please use a Python version >= 2.6.'
<|endoftext|>


<commit_before>import logging
import random

from django.apps import apps
from django.conf import settings
from django.contrib.contenttypes.models import ContentType
from django.db import transaction, IntegrityError
from django.db.utils import OperationalError
from django.utils import timezone

from share import exceptions
from share.change import ChangeGraph
from share.harvest.exceptions import HarvesterConcurrencyError
from share.models import (
    AbstractCreativeWork,
    ChangeSet,
    HarvestJob,
    IngestJob,
    NormalizedData,
    RawDatum,
    Source,
)
from share.models.ingest import RawDatumJob
from share.regulate import Regulator
from share.search import SearchIndexer
from share.util import chunked

logger = logging.getLogger(__name__)

class JobConsumer:
    Job = None
    lock_field = None

    def __init__(self, task=None):
        if self.Job is None or self.lock_field is None:
            raise NotImplementedError
        self.task = task

    def _consume_job(self, job, **kwargs):
        raise NotImplementedError

    def _current_versions(self, job):
        """Get up-to-date values for the job's `*_version` fields

        Dict from field name to version number
        """
        raise NotImplementedError

    def consume(self, job_id=None, exhaust=True, ignore_disabled=False, superfluous=False, force=False, **kwargs):
        """Consume the given job, or consume an available job if no job is specified.

        Parameters:
            job_id (int, optional): Consume the given job. Defaults to None.
                If the given job cannot be locked, the task will retry indefinitely.
                If the given job belongs to a disabled or deleted Source or SourceConfig, the task will fail.
            exhaust (bool, optional): If True and there are queued jobs, start another task. Defaults to True.
                Used to prevent a backlog. If we have a valid job, spin off another task to eat through
                the rest of the queue.
            ignore_disabled (bool, optional): Consume jobs from disabled source configs and/or deleted sources. Defaults to False.
            superfluous (bool, optional): Re-ingest Rawdata that we've already collected. Defaults to False.
            force (bool, optional):
        Additional keyword arguments passed to _consume_job, along with superfluous and force
        """
        with self._locked_job(job_id, ignore_disabled) as job:
            if job is None:
                if job_id is None:
                    logger.info('No %ss are currently available', self.Job.__name__)
                    return
                else:
                    # If an id was given to us, we should have gotten a job
                    job = self.Job.objects.get(id=job_id)  # Force the failure
                    raise Exception('Failed to load {} but then found {!r}.'.format(job_id, job))  # Should never be reached

            assert self.task or not exhaust, 'Cannot pass exhaust=True unless running in an async context'
            if exhaust and job_id is None:
                if force:
                    logger.warning('propagating force=True until queue exhaustion')

                logger.debug('Spawning another task to consume %s', self.Job.__name__)
                res = self.task.apply_async(self.task.request.args, self.task.request.kwargs)
                logger.info('Spawned %r', res)

            if self._prepare_job(job, superfluous=superfluous):
                logger.info('Consuming %r', job)
                with job.handle():
                    self._consume_job(job, **kwargs, superfluous=superfluous, force=force)

    def _prepare_job(self, job, superfluous):
        if self.task and self.task.request.id:
            # Additional attributes for the celery backend
            # Allows for better analytics of currently running tasks
            self.task.update_state(meta={
                'job_id': job.id,
                'source': job.source_config.source.long_title,
                'source_config': job.source_config.label,
            })

            job.task_id = self.task.request.id
            job.save(update_fields=('task_id',))

        if job.completions > 0 and job.status == self.Job.STATUS.succeeded:
            if not superfluous:
                job.skip(job.SkipReasons.duplicated)
                logger.warning('%r has already been consumed. Force a re-run with superfluous=True', job)
                return False
            logger.info('%r has already been consumed. Re-running superfluously', job)

        if not self._update_versions(job):
            job.skip(job.SkipReasons.obsolete)
            return False

        return True

    def _filter_ready(self, qs):
        return qs.filter(
            status__in=self.Job.READY_STATUSES,
        ).exclude(
            claimed=True
        )

    def _locked_job(self, job_id, ignore_disabled=False):
        qs = self.Job.objects.all()
        if job_id is not None:
            logger.debug('Loading %s %d', self.Job.__name__, job_id)
            qs = qs.filter(id=job_id)
        else:
            logger.debug('job_id was not specified, searching for an available job.')

            if not ignore_disabled:
                qs = qs.exclude(
                    source_config__disabled=True,
                ).exclude(
                    source_config__source__is_deleted=True
                )
            qs = self._filter_ready(qs).unlocked(self.lock_field)

        return qs.lock_first(self.lock_field)

    def _update_versions(self, job):
        """Update version fields to the values from self.current_versions

        Return True if successful, else False.
        """
        current_versions = self._current_versions(job)
        if all(getattr(job, f) == v for f, v in current_versions.items()):
            # No updates required
            return True

        if job.completions > 0:
            logger.warning('%r is outdated but has previously completed, skipping...', job)
            return False

        try:
            with transaction.atomic():
                for f, v in current_versions.items():
                    setattr(job, f, v)
                job.save()
            logger.warning('%r has been updated to the versions: %s', job, current_versions)
            return True
        except IntegrityError:
            logger.warning('A newer version of %r already exists, skipping...', job)
            return False

class HarvestJobConsumer(JobConsumer):
    Job = HarvestJob
    lock_field = 'source_config'

    def _filter_ready(self, qs):
        qs = super()._filter_ready(qs)
        return qs.filter(
            end_date__lte=timezone.now().date(),
            source_config__harvest_after__lte=timezone.now().time(),
        )

    def _current_versions(self, job):
        return {
            'source_config_version': job.source_config.version,
            'harvester_version': job.source_config.harvester.version,
        }

    def _consume_job(self, job, force, superfluous, limit=None, ingest=True):
        try:
            if ingest:
                datum_gen = (datum for datum in self._harvest(job, force, limit) if datum.created or superfluous)
                for chunk in chunked(datum_gen, 500):
                    self._bulk_schedule_ingest(job, chunk)
            else:
                for _ in self._harvest(job, force, limit):
                    pass
        except HarvesterConcurrencyError as e:
            if not self.task:
                raise
            # If job_id was specified there's a chance that the advisory lock was not, in fact, acquired.
            # If so, retry indefinitely to preserve existing functionality.
            # Use random to add jitter to help break up locking issues
            # Kinda hacky, allow a stupidly large number of retries as there is no options for infinite
            raise self.task.retry(
                exc=e,
                max_retries=99999,
                countdown=(random.random() + 1) * min(settings.CELERY_RETRY_BACKOFF_BASE ** self.task.request.retries, 60 * 15)
            )

    def _harvest(self, job, force, limit):
        error = None
        datum_ids = []
        logger.info('Harvesting %r', job)
        harvester = job.source_config.get_harvester()

        try:
            for datum in harvester.harvest_date_range(job.start_date, job.end_date, limit=limit, force=force):
                datum_ids.append(datum.id)
                yield datum
        except Exception as e:
            error = e
            raise error
        finally:
            try:
                RawDatumJob.objects.bulk_create([
                    RawDatumJob(job=job, datum_id=datum_id)
                    for datum_id in datum_ids
                ])
            except Exception as e:
                logger.exception('Failed to connect %r to raw data', job)
                # Avoid shadowing the original error
                if not error:
                    raise e

    def _bulk_schedule_ingest(self, job, datums):
        job_kwargs = {
            'source_config': job.source_config,
            'source_config_version': job.source_config.version,
            'transformer_version': job.source_config.transformer.version,
            'regulator_version': Regulator.VERSION,
        }
        IngestJob.objects.bulk_get_or_create(
            [IngestJob(raw_id=datum.id, suid_id=datum.suid_id, **job_kwargs) for datum in datums]
        )

class IngestJobConsumer(JobConsumer):
    Job = IngestJob
    lock_field = 'suid'

    MAX_RETRIES = 5

    def _current_versions(self, job):
        return {
            'source_config_version': job.source_config.version,
            'transformer_version': job.source_config.transformer.version,
            'regulator_version': Regulator.VERSION,
        }

    def _prepare_job(self, job, *args, **kwargs):
        # TODO think about getting rid of these triangles -- infer source config from suid?
        assert job.suid_id == job.raw.suid_id
        assert job.source_config_id == job.suid.source_config_id

        if job.raw.datestamp and self.Job.objects.filter(
                status__in=job.READY_STATUSES,
                suid=job.suid,
                raw__datestamp__gt=job.raw.datestamp
        ).exists():
            job.skip(job.SkipReasons.pointless)
            return False

        return super()._prepare_job(job, *args, **kwargs)

    def _consume_job(self, job, superfluous, force, apply_changes=True, index=True, urgent=False):
        datum = None

        # Check whether we've already done transform/regulate
        if not superfluous:
            datum = job.ingested_normalized_data.order_by('-created_at').first()

        if superfluous or datum is None:
            graph = self._transform(job)
            if not graph:
                return
            graph = self._regulate(job, graph)
            if not graph:
                return
            if settings.SHARE_LEGACY_PIPELINE:
                datum = NormalizedData.objects.create(
                    data={'@graph': graph.to_jsonld()},
                    source=job.suid.source_config.source.user,
                    raw=job.raw,
                )
                job.ingested_normalized_data.add(datum)

        if apply_changes and settings.SHARE_LEGACY_PIPELINE:
            # TODO make this pipeline actually legacy by implementing a new one
            updated_work_ids = self._apply_changes(job, datum)
            if index and updated_work_ids:
                self._update_index(updated_work_ids, urgent)

    def _transform(self, job):
        transformer = job.suid.source_config.get_transformer()
        try:
            graph = transformer.transform(job.raw)
        except exceptions.TransformError as e:
            job.fail(e)
            return None

        if not graph:
            if not job.raw.normalizeddata_set.exists():
                logger.warning('Graph was empty for %s, setting no_output to True', job.raw)
                RawDatum.objects.filter(id=job.raw_id).update(no_output=True)
            else:
                logger.warning('Graph was empty for %s, but a normalized data already exists for it', job.raw)
            return None

        return graph

    def _regulate(self, job, graph):
        try:
            Regulator(job).regulate(graph)
            return graph
        except exceptions.RegulateError as e:
            job.fail(e)
            return None

    def _apply_changes(self, job, normalized_datum):
        updated = None

        try:
            # Load all relevant ContentTypes in a single query
            ContentType.objects.get_for_models(*apps.get_models('share'), for_concrete_models=False)

            with transaction.atomic():
                cg = ChangeGraph(normalized_datum.data['@graph'], namespace=normalized_datum.source.username)
                cg.process()
                cs = ChangeSet.objects.from_graph(cg, normalized_datum.id)
                if cs and (normalized_datum.source.is_robot or normalized_datum.source.is_trusted or Source.objects.filter(user=normalized_datum.source).exists()):
                    updated = cs.accept()

        # Retry if it was just the wrong place at the wrong time
        except (exceptions.IngestConflict, OperationalError) as e:
            job.retries = (job.retries or 0) + 1
            job.save(update_fields=('retries',))
            if job.retries > self.MAX_RETRIES:
                raise
            job.reschedule()
            return

        if not updated:
            return  # Nothing to index

        # Index works that were added or directly updated
        updated_works = set(x.id for x in updated if isinstance(x, AbstractCreativeWork))
        # and works that matched, even if they didn't change, in case any related objects did
        existing_works = set(n.instance.id for n in cg.nodes if isinstance(n.instance, AbstractCreativeWork))

        return list(updated_works | existing_works)

    def _update_index(self, work_ids, urgent):
        indexer = SearchIndexer(self.task.app) if self.task else SearchIndexer()
        indexer.index('creativework', *work_ids, urgent=urgent)
<commit_msg>Don't retry "skipped" jobs<commit_after>import logging
import random

from django.apps import apps
from django.conf import settings
from django.contrib.contenttypes.models import ContentType
from django.db import transaction, IntegrityError
from django.db.utils import OperationalError
from django.utils import timezone

from share import exceptions
from share.change import ChangeGraph
from share.harvest.exceptions import HarvesterConcurrencyError
from share.models import (
    AbstractCreativeWork,
    ChangeSet,
    HarvestJob,
    IngestJob,
    NormalizedData,
    RawDatum,
    Source,
)
from share.models.ingest import RawDatumJob
from share.regulate import Regulator
from share.search import SearchIndexer
from share.util import chunked

logger = logging.getLogger(__name__)

class JobConsumer:
    Job = None
    lock_field = None

    def __init__(self, task=None):
        if self.Job is None or self.lock_field is None:
            raise NotImplementedError
        self.task = task

    def _consume_job(self, job, **kwargs):
        raise NotImplementedError

    def _current_versions(self, job):
        """Get up-to-date values for the job's `*_version` fields

        Dict from field name to version number
        """
        raise NotImplementedError

    def consume(self, job_id=None, exhaust=True, ignore_disabled=False, superfluous=False, force=False, **kwargs):
        """Consume the given job, or consume an available job if no job is specified.

        Parameters:
            job_id (int, optional): Consume the given job. Defaults to None.
                If the given job cannot be locked, the task will retry indefinitely.
                If the given job belongs to a disabled or deleted Source or SourceConfig, the task will fail.
            exhaust (bool, optional): If True and there are queued jobs, start another task. Defaults to True.
                Used to prevent a backlog. If we have a valid job, spin off another task to eat through
                the rest of the queue.
            ignore_disabled (bool, optional): Consume jobs from disabled source configs and/or deleted sources. Defaults to False.
            superfluous (bool, optional): Re-ingest Rawdata that we've already collected. Defaults to False.
            force (bool, optional):
        Additional keyword arguments passed to _consume_job, along with superfluous and force
        """
        with self._locked_job(job_id, ignore_disabled) as job:
            if job is None:
                if job_id is None:
                    logger.info('No %ss are currently available', self.Job.__name__)
                    return
                else:
                    # If an id was given to us, we should have gotten a job
                    job = self.Job.objects.get(id=job_id)  # Force the failure
                    raise Exception('Failed to load {} but then found {!r}.'.format(job_id, job))  # Should never be reached

            assert self.task or not exhaust, 'Cannot pass exhaust=True unless running in an async context'
            if exhaust and job_id is None:
                if force:
                    logger.warning('propagating force=True until queue exhaustion')

                logger.debug('Spawning another task to consume %s', self.Job.__name__)
                res = self.task.apply_async(self.task.request.args, self.task.request.kwargs)
                logger.info('Spawned %r', res)

            if self._prepare_job(job, superfluous=superfluous):
                logger.info('Consuming %r', job)
                with job.handle():
                    self._consume_job(job, **kwargs, superfluous=superfluous, force=force)

    def _prepare_job(self, job, superfluous):
        if job.status == self.Job.STATUS.skipped:
            # Need some way to short-circuit a superfluous retry loop
            logger.warning('%r has been marked skipped. Change its status to allow re-running it', job)
            return False

        if self.task and self.task.request.id:
            # Additional attributes for the celery backend
            # Allows for better analytics of currently running tasks
            self.task.update_state(meta={
                'job_id': job.id,
                'source': job.source_config.source.long_title,
                'source_config': job.source_config.label,
            })

            job.task_id = self.task.request.id
            job.save(update_fields=('task_id',))

        if job.completions > 0 and job.status == self.Job.STATUS.succeeded:
            if not superfluous:
                job.skip(job.SkipReasons.duplicated)
                logger.warning('%r has already been consumed. Force a re-run with superfluous=True', job)
                return False
            logger.info('%r has already been consumed. Re-running superfluously', job)

        if not self._update_versions(job):
            job.skip(job.SkipReasons.obsolete)
            return False

        return True

    def _filter_ready(self, qs):
        return qs.filter(
            status__in=self.Job.READY_STATUSES,
        ).exclude(
            claimed=True
        )

    def _locked_job(self, job_id, ignore_disabled=False):
        qs = self.Job.objects.all()
        if job_id is not None:
            logger.debug('Loading %s %d', self.Job.__name__, job_id)
            qs = qs.filter(id=job_id)
        else:
            logger.debug('job_id was not specified, searching for an available job.')

            if not ignore_disabled:
                qs = qs.exclude(
                    source_config__disabled=True,
                ).exclude(
                    source_config__source__is_deleted=True
                )
            qs = self._filter_ready(qs).unlocked(self.lock_field)

        return qs.lock_first(self.lock_field)

    def _update_versions(self, job):
        """Update version fields to the values from self.current_versions

        Return True if successful, else False.
        """
        current_versions = self._current_versions(job)
        if all(getattr(job, f) == v for f, v in current_versions.items()):
            # No updates required
            return True

        if job.completions > 0:
            logger.warning('%r is outdated but has previously completed, skipping...', job)
            return False

        try:
            with transaction.atomic():
                for f, v in current_versions.items():
                    setattr(job, f, v)
                job.save()
            logger.warning('%r has been updated to the versions: %s', job, current_versions)
            return True
        except IntegrityError:
            logger.warning('A newer version of %r already exists, skipping...', job)
            return False

class HarvestJobConsumer(JobConsumer):
    Job = HarvestJob
    lock_field = 'source_config'

    def _filter_ready(self, qs):
        qs = super()._filter_ready(qs)
        return qs.filter(
            end_date__lte=timezone.now().date(),
            source_config__harvest_after__lte=timezone.now().time(),
        )

    def _current_versions(self, job):
        return {
            'source_config_version': job.source_config.version,
            'harvester_version': job.source_config.harvester.version,
        }

    def _consume_job(self, job, force, superfluous, limit=None, ingest=True):
        try:
            if ingest:
                datum_gen = (datum for datum in self._harvest(job, force, limit) if datum.created or superfluous)
                for chunk in chunked(datum_gen, 500):
                    self._bulk_schedule_ingest(job, chunk)
            else:
                for _ in self._harvest(job, force, limit):
                    pass
        except HarvesterConcurrencyError as e:
            if not self.task:
                raise
            # If job_id was specified there's a chance that the advisory lock was not, in fact, acquired.
            # If so, retry indefinitely to preserve existing functionality.
            # Use random to add jitter to help break up locking issues
            # Kinda hacky, allow a stupidly large number of retries as there is no options for infinite
            raise self.task.retry(
                exc=e,
                max_retries=99999,
                countdown=(random.random() + 1) * min(settings.CELERY_RETRY_BACKOFF_BASE ** self.task.request.retries, 60 * 15)
            )

    def _harvest(self, job, force, limit):
        error = None
        datum_ids = []
        logger.info('Harvesting %r', job)
        harvester = job.source_config.get_harvester()

        try:
            for datum in harvester.harvest_date_range(job.start_date, job.end_date, limit=limit, force=force):
                datum_ids.append(datum.id)
                yield datum
        except Exception as e:
            error = e
            raise error
        finally:
            try:
                RawDatumJob.objects.bulk_create([
                    RawDatumJob(job=job, datum_id=datum_id)
                    for datum_id in datum_ids
                ])
            except Exception as e:
                logger.exception('Failed to connect %r to raw data', job)
                # Avoid shadowing the original error
                if not error:
                    raise e

    def _bulk_schedule_ingest(self, job, datums):
        job_kwargs = {
            'source_config': job.source_config,
            'source_config_version': job.source_config.version,
            'transformer_version': job.source_config.transformer.version,
            'regulator_version': Regulator.VERSION,
        }
        IngestJob.objects.bulk_get_or_create(
            [IngestJob(raw_id=datum.id, suid_id=datum.suid_id, **job_kwargs) for datum in datums]
        )

class IngestJobConsumer(JobConsumer):
    Job = IngestJob
    lock_field = 'suid'

    MAX_RETRIES = 5

    def _current_versions(self, job):
        return {
            'source_config_version': job.source_config.version,
            'transformer_version': job.source_config.transformer.version,
            'regulator_version': Regulator.VERSION,
        }

    def _prepare_job(self, job, *args, **kwargs):
        # TODO think about getting rid of these triangles -- infer source config from suid?
        assert job.suid_id == job.raw.suid_id
        assert job.source_config_id == job.suid.source_config_id

        if job.raw.datestamp and self.Job.objects.filter(
                status__in=job.READY_STATUSES,
                suid=job.suid,
                raw__datestamp__gt=job.raw.datestamp
        ).exists():
            job.skip(job.SkipReasons.pointless)
            return False

        return super()._prepare_job(job, *args, **kwargs)

    def _consume_job(self, job, superfluous, force, apply_changes=True, index=True, urgent=False):
        datum = None

        # Check whether we've already done transform/regulate
        if not superfluous:
            datum = job.ingested_normalized_data.order_by('-created_at').first()

        if superfluous or datum is None:
            graph = self._transform(job)
            if not graph:
                return
            graph = self._regulate(job, graph)
            if not graph:
                return
            if settings.SHARE_LEGACY_PIPELINE:
                datum = NormalizedData.objects.create(
                    data={'@graph': graph.to_jsonld()},
                    source=job.suid.source_config.source.user,
                    raw=job.raw,
                )
                job.ingested_normalized_data.add(datum)

        if apply_changes and settings.SHARE_LEGACY_PIPELINE:
            # TODO make this pipeline actually legacy by implementing a new one
            updated_work_ids = self._apply_changes(job, datum)
            if index and updated_work_ids:
                self._update_index(updated_work_ids, urgent)

    def _transform(self, job):
        transformer = job.suid.source_config.get_transformer()
        try:
            graph = transformer.transform(job.raw)
        except exceptions.TransformError as e:
            job.fail(e)
            return None

        if not graph:
            if not job.raw.normalizeddata_set.exists():
                logger.warning('Graph was empty for %s, setting no_output to True', job.raw)
                RawDatum.objects.filter(id=job.raw_id).update(no_output=True)
            else:
                logger.warning('Graph was empty for %s, but a normalized data already exists for it', job.raw)
            return None

        return graph

    def _regulate(self, job, graph):
        try:
            Regulator(job).regulate(graph)
            return graph
        except exceptions.RegulateError as e:
            job.fail(e)
            return None

    def _apply_changes(self, job, normalized_datum):
        updated = None

        try:
            # Load all relevant ContentTypes in a single query
            ContentType.objects.get_for_models(*apps.get_models('share'), for_concrete_models=False)

            with transaction.atomic():
                cg = ChangeGraph(normalized_datum.data['@graph'], namespace=normalized_datum.source.username)
                cg.process()
                cs = ChangeSet.objects.from_graph(cg, normalized_datum.id)
                if cs and (normalized_datum.source.is_robot or normalized_datum.source.is_trusted or Source.objects.filter(user=normalized_datum.source).exists()):
                    updated = cs.accept()

        # Retry if it was just the wrong place at the wrong time
        except (exceptions.IngestConflict, OperationalError) as e:
            job.retries = (job.retries or 0) + 1
            job.save(update_fields=('retries',))
            if job.retries > self.MAX_RETRIES:
                raise
            job.reschedule()
            return

        if not updated:
            return  # Nothing to index

        # Index works that were added or directly updated
        updated_works = set(x.id for x in updated if isinstance(x, AbstractCreativeWork))
        # and works that matched, even if they didn't change, in case any related objects did
        existing_works = set(n.instance.id for n in cg.nodes if isinstance(n.instance, AbstractCreativeWork))

        return list(updated_works | existing_works)

    def _update_index(self, work_ids, urgent):
        indexer = SearchIndexer(self.task.app) if self.task else SearchIndexer()
        indexer.index('creativework', *work_ids, urgent=urgent)
<|endoftext|>


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

#   Copyright (c) 2010-2014, MIT Probabilistic Computing Project
#
#   Licensed under the Apache License, Version 2.0 (the "License");
#   you may not use this file except in compliance with the License.
#   You may obtain a copy of the License at
#
#       http://www.apache.org/licenses/LICENSE-2.0
#
#   Unless required by applicable law or agreed to in writing, software
#   distributed under the License is distributed on an "AS IS" BASIS,
#   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
#   See the License for the specific language governing permissions and
#   limitations under the License.

def pp_cursor(out, cursor):
    if cursor.description is None:
        return
    labels = [d[0] for d in cursor.description]
    table = list(cursor)
    # XXX Consider quotation/escapes.
    colwidths = [len(label) for label in labels]
    for row in table:
        for colno, v in enumerate(row):
            # XXX Consider quotation/escapes.
            # XXX Combining characters?
            colwidths[colno] = max(colwidths[colno], len(unicode(v)))
    first = True
    for colno, label in enumerate(labels):
        if first:
            first = False
        else:
            out.write(' | ')
        # XXX Quote/escape.
        out.write('%*s' % (colwidths[colno], label))
    out.write('\n')
    first = True
    for colno, label in enumerate(labels):
        if first:
            first = False
        else:
            out.write('-+-')
        # XXX Quote/escape.
        out.write('%s' % ('-' * colwidths[colno]))
    out.write('\n')
    for row in table:
        first = True
        for colno, v in enumerate(row):
            if first:
                first = False
            else:
                out.write(' | ')
            # XXX Quote/escape.
            out.write('%*s' % (colwidths[colno], unicode(v)))
        out.write('\n')
<commit_msg>Break pretty printing away from cursors so I can pretty print the results of multiple queries together.<commit_after># -*- coding: utf-8 -*-

#   Copyright (c) 2010-2014, MIT Probabilistic Computing Project
#
#   Licensed under the Apache License, Version 2.0 (the "License");
#   you may not use this file except in compliance with the License.
#   You may obtain a copy of the License at
#
#       http://www.apache.org/licenses/LICENSE-2.0
#
#   Unless required by applicable law or agreed to in writing, software
#   distributed under the License is distributed on an "AS IS" BASIS,
#   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
#   See the License for the specific language governing permissions and
#   limitations under the License.

def pp_cursor(out, cursor):
    if cursor.description is None:
        return
    labels = [d[0] for d in cursor.description]
    table = list(cursor)
    pp_list(out, table, labels)

def pp_list(out, table, labels):
    # XXX Consider quotation/escapes.
    colwidths = [len(label) for label in labels]
    for row in table:
        for colno, v in enumerate(row):
            # XXX Consider quotation/escapes.
            # XXX Combining characters?
            colwidths[colno] = max(colwidths[colno], len(unicode(v)))
    first = True
    for colno, label in enumerate(labels):
        if first:
            first = False
        else:
            out.write(' | ')
        # XXX Quote/escape.
        out.write('%*s' % (colwidths[colno], label))
    out.write('\n')
    first = True
    for colno, label in enumerate(labels):
        if first:
            first = False
        else:
            out.write('-+-')
        # XXX Quote/escape.
        out.write('%s' % ('-' * colwidths[colno]))
    out.write('\n')
    for row in table:
        first = True
        for colno, v in enumerate(row):
            if first:
                first = False
            else:
                out.write(' | ')
            # XXX Quote/escape.
            out.write('%*s' % (colwidths[colno], unicode(v)))
        out.write('\n')
<|endoftext|>