<commit_before>import os

import pandas

from ulmo import util


GHCN_DAILY_DIR = os.path.join(util.get_ulmo_dir(), 'ncdc/ghcn_daily')
GHCN_DAILY_STATIONS_FILE = os.path.join(GHCN_DAILY_DIR, 'ghcnd-stations.txt')


def get_stations(country=None, state=None, update=True, as_dataframe=False):
    """returns a collection of stations

    The stations can be represented as dict of station dicts keyed to their station ID
    or a pandas.dataframe (see the as_dataframe parameter)

    Parameters
    ----------
    country : The country code to use to limit station results. If set to None
            (default), then stations from all countries are returned.
    state : The state code to use to limit station results. If set to None
            (default), then stations from all states are returned.
    update : If False, tries to use a cached copy of the stations file. If one
             can't be found or if update is True, then a new copy of the
             stations file is pulled from the web. If update is True, but the
             cached stations file is still good then a new file won't be pulled.
             Default is True.
    as_dataframe : If True, a pandas.DataFrame object will be returned,
            otherwise a dict is of stations dicts is returned. The pandas
            dataframe is used internally, so setting this to True is a little
            bit faster as it skips a serialization step. Default is False.,
    """
    if update or not os.path.exists(GHCN_DAILY_STATIONS_FILE):
        url = 'http://www1.ncdc.noaa.gov/pub/data/ghcn/daily/ghcnd-stations.txt'
        util.download_if_new(url, GHCN_DAILY_STATIONS_FILE, check_modified=True)
        print 'Saved station list {0}'.format(GHCN_DAILY_STATIONS_FILE)

    columns = [
        ('country', 0, 2),
        ('network', 2, 3),
        ('network_id', 3, 11),
        ('latitude', 12, 20),
        ('longitude', 21, 30),
        ('elevation', 31, 37),
        ('state', 38, 40),
        ('name', 41, 71),
        ('gsn_flag', 72, 75),
        ('hcn_flag', 76, 79),
        ('wm_oid', 80, 85),
    ]
    colspecs = [(start, end) for name, start, end in columns]
    names = [name for name, start, end in columns]

    stations = pandas.io.parsers.read_fwf(GHCN_DAILY_STATIONS_FILE, colspecs=colspecs,
            header=None, names=names)

    if not country is None:
        stations = stations[stations['country'] == country]
    if not state is None:
        stations = stations[stations['state'] == state]

    # wm_oid gets converted as a float, so cast it to str manually
    stations['wm_oid'] = stations['wm_oid'].astype('|S5')
    stations['wm_oid'][stations['wm_oid'] == 'nan'] = None

    for column_name in stations.columns:
        stations[column_name][pandas.isnull(stations[column_name])] = None

    # set id
    stations['id'] = stations[['country', 'network', 'network_id']].T.apply(''.join)

    if as_dataframe:
        # return dataframe with a meaningful index (id)
        return stations.set_index('id', drop=False)
    else:
        return {
            station_series['id']: station_series.to_dict()
            for index, station_series in stations.T.iteritems()
        }
<commit_msg>use dataframe.T.to_dict() rather than iterating to create dicts<commit_after>import os

import pandas

from ulmo import util


GHCN_DAILY_DIR = os.path.join(util.get_ulmo_dir(), 'ncdc/ghcn_daily')
GHCN_DAILY_STATIONS_FILE = os.path.join(GHCN_DAILY_DIR, 'ghcnd-stations.txt')


def get_stations(country=None, state=None, update=True, as_dataframe=False):
    """returns a collection of stations

    The stations can be represented as dict of station dicts keyed to their station ID
    or a pandas.dataframe (see the as_dataframe parameter)

    Parameters
    ----------
    country : The country code to use to limit station results. If set to None
            (default), then stations from all countries are returned.
    state : The state code to use to limit station results. If set to None
            (default), then stations from all states are returned.
    update : If False, tries to use a cached copy of the stations file. If one
             can't be found or if update is True, then a new copy of the
             stations file is pulled from the web. If update is True, but the
             cached stations file is still good then a new file won't be pulled.
             Default is True.
    as_dataframe : If True, a pandas.DataFrame object will be returned,
            otherwise a dict is of stations dicts is returned. The pandas
            dataframe is used internally, so setting this to True is a little
            bit faster as it skips a serialization step. Default is False.,
    """
    if update or not os.path.exists(GHCN_DAILY_STATIONS_FILE):
        url = 'http://www1.ncdc.noaa.gov/pub/data/ghcn/daily/ghcnd-stations.txt'
        util.download_if_new(url, GHCN_DAILY_STATIONS_FILE, check_modified=True)
        print 'Saved station list {0}'.format(GHCN_DAILY_STATIONS_FILE)

    columns = [
        ('country', 0, 2),
        ('network', 2, 3),
        ('network_id', 3, 11),
        ('latitude', 12, 20),
        ('longitude', 21, 30),
        ('elevation', 31, 37),
        ('state', 38, 40),
        ('name', 41, 71),
        ('gsn_flag', 72, 75),
        ('hcn_flag', 76, 79),
        ('wm_oid', 80, 85),
    ]
    colspecs = [(start, end) for name, start, end in columns]
    names = [name for name, start, end in columns]

    stations = pandas.io.parsers.read_fwf(GHCN_DAILY_STATIONS_FILE, colspecs=colspecs,
            header=None, names=names)

    if not country is None:
        stations = stations[stations['country'] == country]
    if not state is None:
        stations = stations[stations['state'] == state]

    # wm_oid gets converted as a float, so cast it to str manually
    stations['wm_oid'] = stations['wm_oid'].astype('|S5')
    stations['wm_oid'][stations['wm_oid'] == 'nan'] = None

    for column_name in stations.columns:
        stations[column_name][pandas.isnull(stations[column_name])] = None

    # set id
    stations['id'] = stations[['country', 'network', 'network_id']].T.apply(''.join)
    stations = stations.set_index('id', drop=False)

    if as_dataframe:
        return stations
    else:
        for column_name in stations.columns:
            stations[column_name][pandas.isnull(stations[column_name])] = None

        return stations.T.to_dict()

{'fraction_non_alphanumeric': 0.10620432281138237, 'fraction_numerical': 0.013061732234489193, 'mean_word_length': 2.792452830188679, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 3, 'https://': 0, 'lorem ipsum': 0, 'www.': 2, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 2, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '789769', 'n_tokens_mistral': 2080, 'n_tokens_neox': 1950, 'n_words': 660}