<commit_before>import re
from os.path import isfile, splitext
from functools import wraps
import argparse
from inspect import getfullargspec

import numpy as np
import pandas as pd

from .gentools import force_df
from .frametools import df_to_list

PICKLE_EXT = ('pkl', 'p')   # First is default for writing to pickle
HDF5_EXT = ('h5', 'hdf5')


# TODO: lob tests
def load_or_build(raw_filepath, copydta=False, path_args=[]):
    """Loads `raw_filepath` as a DataFrame if it exists, otherwise builds the
    data and saves it to `raw_filepath`.

    Args:
        raw_filepath (str): Path to saved DataFrame. If ``raw_filepath``
            includes named replacement fields (e.g., "``'{arg_name}'``") with
            the same name as function arguments, passed values will be inserted
            into the file path.

    Example:

        .. code-block:: python

            @load_or_build('data_for_{year}.pkl')
            def make_data(year):
                <Make the data>

            df = make_data(2018)    # Saves to 'data_for_2018.pkl'

    Keyword Args:
        copydta (bool): Defaults to False. If true, save a copy of the data in
            Stata DTA format if ``raw_filepath`` is not already a DTA file.

        path_args (list-like): DEPRECATED: Use named replacement fields
            instead. A list of ints or strs that point to args or
            kwargs of the build function, respectively. The value of these
            arguments will then be use to format ``raw_filepath``.

    Example:

        .. code-block:: python

            @load_or_build('file_{}_{}.csv', path_args=[0, 'b'])
            def foo(a, b=None):
                return pd.DataFrame([a, b])

            if __name__ == '__main__':
                # Saves `df` to `file_infix_suffix.csv`
                foo('infix', 'suffix')

    Other Parameters:
        These are additional kwargs that can be passed to the wrapped function
            that affect the behavior of ``load_or_build``.

        _rebuild (bool): Defaults to False. If True, build the DataFrame and
            save it to `filepath` even if `filepath` already exists.
        _load (bool): Defaults to True. If True, try loading the data before
            building it. If False, the building function is called and the
            result returned with no data written to disk.
    """
    def actualDecorator(builder):
        @wraps(builder)
        def wrapper(*args, **kwargs):
            # Filter the kwargs from `builder` meant for `load_or_build`
            load = kwargs.pop('_load', True)
            rebuild = kwargs.pop('_rebuild', False)

            # Format the file path
            if not path_args:   # XXX this is preferred/keep old for compat
                filepath = _set_filepath(raw_filepath, args, kwargs, builder)
            else:
                filepath = _set_filepath_old(raw_filepath, path_args, args,
                                             kwargs, builder)

            if load and isfile(filepath) and not rebuild:
                # If it's on disk and there are no objections, read it
                df = read(filepath)
            elif not load:
                # If `load = False` is passed, just return the built `df`
                df = builder(*args, **kwargs)
            else:
                # If it's just not on disk, build it, save it, and copy it
                print("****** Building *******\n\tfile: {}".format(filepath))
                print("\tfunc: {}".format(builder.__name__))
                print("*************")
                df = builder(*args, **kwargs)
                write(df, filepath)
                # Copy to Stata DTA if needed
                fileroot, fileext = splitext(filepath)
                if copydta and fileext != '.dta':
                    force_df(df).to_stata(fileroot + '.dta')

            return df
        return wrapper
    return actualDecorator

def _set_filepath(raw_filepath, args, kwargs, builder):
    argspec = getfullargspec(builder)
    if re.search('{.*}', raw_filepath):
        argspec = getfullargspec(builder)
        arg_names = argspec.args
        arg_dict = dict()
        if argspec.defaults:
            kwarg_defaults = list(argspec.defaults)
            while kwarg_defaults:
                arg_dict[arg_names.pop()] = kwarg_defaults.pop()

        for idx, arg in enumerate(args):
            arg_dict[argspec.args[idx]] = arg
        for k, v in kwargs.items():
            arg_dict[k] = v
        new_filepath = raw_filepath.format(**arg_dict)
    else:
        new_filepath = raw_filepath

    return new_filepath

def _set_filepath_old(raw_path, path_args, args, kwargs, builder):
    """
    Parse `args` and `kwargs` as directed by `path_args` and insert them
    into `raw_path`.
    """
    argspec = getfullargspec(builder)
    format_filepath = _parse_pathargs(path_args, args, kwargs, argspec)
    try:
        filepath = raw_path.format(*format_filepath)
    except IndexError:
        err_str = (
            "Not enough arguments to fill file path template:\n"
            "    Args: {}\n"
            "    Path: {}"
        ).format(raw_path, format_filepath)
        raise ValueError(err_str)
    return filepath

def _parse_pathargs(path_args, args, kwargs, argspec):
    """
    `path_args`, iterable: Which args/kwargs to use, in order. Numbers are
        args, strings are kwargs.
    `args`, `kwargs`: The value of args/kwargs passed to the original builder
        function.
    `argspec`: Stores values of default kwargs from builder function.
    """
    patharg_values = []
    # Handle default kwargs
    argnames = argspec.args
    defaults = argspec.defaults

    for arg in path_args:
        arg_type = type(arg)
        if arg_type is int:
            patharg_values.append(args[arg])
        elif arg_type is str:
            try:
                patharg_values.append(kwargs[arg])
            except KeyError:
                nargs = len(argnames) - len(defaults)
                kwarg_idx = argnames.index(arg) - nargs
                patharg_values.append(defaults[kwarg_idx])
        else:<commit_msg>BUG: workaround for 'getfullargspec' in Python 2<commit_after>import re
from os.path import isfile, splitext
from functools import wraps
import argparse

import sys
import numpy as np
import pandas as pd

from .gentools import force_df
from .frametools import df_to_list

IS_PY2 = sys.version_info[0] < 3
if IS_PY2:
    from inspect import getargspec as getfullargspec
else:
    from inspect import getfullargspec


PICKLE_EXT = ('pkl', 'p')   # First is default for writing to pickle
HDF5_EXT = ('h5', 'hdf5')


# TODO: lob tests
def load_or_build(raw_filepath, copydta=False, path_args=[]):
    """Loads `raw_filepath` as a DataFrame if it exists, otherwise builds the
    data and saves it to `raw_filepath`.

    Args:
        raw_filepath (str): Path to saved DataFrame. If ``raw_filepath``
            includes named replacement fields (e.g., "``'{arg_name}'``") with
            the same name as function arguments, passed values will be inserted
            into the file path.

    Example:

        .. code-block:: python

            @load_or_build('data_for_{year}.pkl')
            def make_data(year):
                <Make the data>

            df = make_data(2018)    # Saves to 'data_for_2018.pkl'

    Keyword Args:
        copydta (bool): Defaults to False. If true, save a copy of the data in
            Stata DTA format if ``raw_filepath`` is not already a DTA file.

        path_args (list-like): DEPRECATED: Use named replacement fields
            instead. A list of ints or strs that point to args or
            kwargs of the build function, respectively. The value of these
            arguments will then be use to format ``raw_filepath``.

    Example:

        .. code-block:: python

            @load_or_build('file_{}_{}.csv', path_args=[0, 'b'])
            def foo(a, b=None):
                return pd.DataFrame([a, b])

            if __name__ == '__main__':
                # Saves `df` to `file_infix_suffix.csv`
                foo('infix', 'suffix')

    Other Parameters:
        These are additional kwargs that can be passed to the wrapped function
            that affect the behavior of ``load_or_build``.

        _rebuild (bool): Defaults to False. If True, build the DataFrame and
            save it to `filepath` even if `filepath` already exists.
        _load (bool): Defaults to True. If True, try loading the data before
            building it. If False, the building function is called and the
            result returned with no data written to disk.
    """
    def actualDecorator(builder):
        @wraps(builder)
        def wrapper(*args, **kwargs):
            # Filter the kwargs from `builder` meant for `load_or_build`
            load = kwargs.pop('_load', True)
            rebuild = kwargs.pop('_rebuild', False)

            # Format the file path
            if not path_args:   # XXX this is preferred/keep old for compat
                filepath = _set_filepath(raw_filepath, args, kwargs, builder)
            else:
                filepath = _set_filepath_old(raw_filepath, path_args, args,
                                             kwargs, builder)

            if load and isfile(filepath) and not rebuild:
                # If it's on disk and there are no objections, read it
                df = read(filepath)
            elif not load:
                # If `load = False` is passed, just return the built `df`
                df = builder(*args, **kwargs)
            else:
                # If it's just not on disk, build it, save it, and copy it
                print("****** Building *******\n\tfile: {}".format(filepath))
                print("\tfunc: {}".format(builder.__name__))
                print("*************")
                df = builder(*args, **kwargs)
                write(df, filepath)
                # Copy to Stata DTA if needed
                fileroot, fileext = splitext(filepath)
                if copydta and fileext != '.dta':
                    force_df(df).to_stata(fileroot + '.dta')

            return df
        return wrapper
    return actualDecorator

def _set_filepath(raw_filepath, args, kwargs, builder):
    argspec = argspec(builder)
    if re.search('{.*}', raw_filepath):
        argspec = getfullargspec(builder)
        arg_names = argspec.args
        arg_dict = dict()
        if argspec.defaults:
            kwarg_defaults = list(argspec.defaults)
            while kwarg_defaults:
                arg_dict[arg_names.pop()] = kwarg_defaults.pop()

        for idx, arg in enumerate(args):
            arg_dict[argspec.args[idx]] = arg
        for k, v in kwargs.items():
            arg_dict[k] = v
        new_filepath = raw_filepath.format(**arg_dict)
    else:
        new_filepath = raw_filepath

    return new_filepath

def _set_filepath_old(raw_path, path_args, args, kwargs, builder):
    """
    Parse `args` and `kwargs` as directed by `path_args` and insert them
    into `raw_path`.
    """
    argspec = getfullargspec(builder)
    format_filepath = _parse_pathargs(path_args, args, kwargs, argspec)
    try:
        filepath = raw_path.format(*format_filepath)
    except IndexError:
        err_str = (
            "Not enough arguments to fill file path template:\n"
            "    Args: {}\n"
            "    Path: {}"
        ).format(raw_path, format_filepath)
        raise ValueError(err_str)
    return filepath

def _parse_pathargs(path_args, args, kwargs, argspec):
    """
    `path_args`, iterable: Which args/kwargs to use, in order. Numbers are
        args, strings are kwargs.
    `args`, `kwargs`: The value of args/kwargs passed to the original builder
        function.
    `argspec`: Stores values of default kwargs from builder function.
    """
    patharg_values = []
    # Handle default kwargs
    if IS_PY2:
        argnames, __, __, defaults = argspec
    else:
        argnames = argspec.args
        defaults = argspec.defaults

    for arg in path_args:
        arg_type = type(arg)
        if arg_type is int:
            patharg_values.append(args[arg])
        elif arg_type is str:
            try:
                patharg_values.append(kwargs[arg])
            except KeyError:
                nargs = len(argnames) - len(defaults)
                kwarg_idx = argnames.index(arg) - nargs
                patharg_values.append(defaults[kwarg_idx])
        else:
{'fraction_non_alphanumeric': 0.0927290198590647, 'fraction_numerical': 0.002402306213965407, 'mean_word_length': 2.0723247232472324, 'pattern_counts': {'":': 0, '<': 6, '<?xml version=': 0, '>': 5, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 4, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '2089069', 'n_tokens_mistral': 3636, 'n_tokens_neox': 3541, 'n_words': 1165}