<commit_before>from vega_datasets._dataset import Dataset


class DataLoader(object):
    """Load a dataset from a local file or remote URL.

    There are two ways to call this; for example to load the iris dataset, you
    can call this object and pass the dataset name by string:

        >>> from vega_datasets import data
        >>> df = data('iris')

    or you can call the associated named method:

        >>> df = data.iris()

    Optionally, additional parameters can be passed to either of these

    Optional parameters
    -------------------
    return_raw : boolean
        If True, then return the raw string or bytes.
        If False (default), then return a pandas dataframe.
    use_local : boolean
        If True (default), then attempt to load the dataset locally. If
        False or if the dataset is not available locally, then load the
        data from an external URL.
    **kwargs :
        additional keyword arguments are passed to the pandas parsing function,
        either ``read_csv()`` or ``read_json()`` depending on the data format.
    """
    _datasets = {name.replace('-', '_'): name for name in Dataset.list_datasets()}

    def _method_factory(self, name):
        def get_dataset(return_raw=False, use_local=True, **kwargs):
            return self._load_dataset(self._datasets[name],
                                      return_raw=return_raw,
                                      use_local=use_local,
                                      **kwargs)
        return get_dataset

    def _load_dataset(self, name, **kwargs):
        return Dataset(name).load(**kwargs)

    def __call__(self, name, return_raw=False, use_local=True, **kwargs):
        loader = getattr(self, name.replace('-', '_'))
        return loader(return_raw=False, use_local=True, **kwargs)

    def __getattr__(self, attr):
        if attr in self._datasets:
            return self._method_factory(attr)
        else:
            raise AttributeError("No dataset named '{0}'".format(attr))

    def __dir__(self):
        return list(self._datasets.keys())

    #--------------------------------------------------------------------------
    # specialized methods for individual datasets that need specific processing
    def stocks(self, pivoted=False, return_raw=False, use_local=True,
               parse_dates=('date',), **kwargs):
        """A time-series of stock prices from several tech companies

        Parameters
        ----------
        pivoted : boolean
            If True, then pivot the data so each company is in its own column
        """
        if type(parse_dates) is tuple:
            parse_dates=list(parse_dates)
        data = self._load_dataset('stocks',
                                  return_raw=return_raw,
                                  use_local=use_local,
                                  parse_dates=parse_dates, **kwargs)
        if pivoted:
            data = data.pivot(index='date', columns='symbol', values='price')
        return data<commit_msg>BUG: don't pivot raw data<commit_after>from vega_datasets._dataset import Dataset


class DataLoader(object):
    """Load a dataset from a local file or remote URL.

    There are two ways to call this; for example to load the iris dataset, you
    can call this object and pass the dataset name by string:

        >>> from vega_datasets import data
        >>> df = data('iris')

    or you can call the associated named method:

        >>> df = data.iris()

    Optionally, additional parameters can be passed to either of these

    Optional parameters
    -------------------
    return_raw : boolean
        If True, then return the raw string or bytes.
        If False (default), then return a pandas dataframe.
    use_local : boolean
        If True (default), then attempt to load the dataset locally. If
        False or if the dataset is not available locally, then load the
        data from an external URL.
    **kwargs :
        additional keyword arguments are passed to the pandas parsing function,
        either ``read_csv()`` or ``read_json()`` depending on the data format.
    """
    _datasets = {name.replace('-', '_'): name for name in Dataset.list_datasets()}

    def _method_factory(self, name):
        def get_dataset(return_raw=False, use_local=True, **kwargs):
            return self._load_dataset(self._datasets[name],
                                      return_raw=return_raw,
                                      use_local=use_local,
                                      **kwargs)
        return get_dataset

    def _load_dataset(self, name, **kwargs):
        return Dataset(name).load(**kwargs)

    def __call__(self, name, return_raw=False, use_local=True, **kwargs):
        loader = getattr(self, name.replace('-', '_'))
        return loader(return_raw=False, use_local=True, **kwargs)

    def __getattr__(self, attr):
        if attr in self._datasets:
            return self._method_factory(attr)
        else:
            raise AttributeError("No dataset named '{0}'".format(attr))

    def __dir__(self):
        return list(self._datasets.keys())

    #--------------------------------------------------------------------------
    # specialized methods for individual datasets that need specific processing
    def stocks(self, pivoted=False, return_raw=False, use_local=True,
               parse_dates=('date',), **kwargs):
        """A time-series of stock prices from several tech companies

        Parameters
        ----------
        pivoted : boolean
            If True, then pivot the data so each company is in its own column.
            Not referenced if return_raw is True.
        """
        if type(parse_dates) is tuple:
            parse_dates=list(parse_dates)
        data = self._load_dataset('stocks',
                                  return_raw=return_raw,
                                  use_local=use_local,
                                  parse_dates=parse_dates, **kwargs)
        if pivoted and not return_raw:
            data = data.pivot(index='date', columns='symbol', values='price')
        return data
{'fraction_non_alphanumeric': 0.11695524338451486, 'fraction_numerical': 0.0003266906239790918, 'mean_word_length': 1.9608317214700193, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 21, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 12, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '5362589', 'n_tokens_mistral': 1557, 'n_tokens_neox': 1501, 'n_words': 530}