diff --git a/django_pandas/io.py b/django_pandas/io.py index a9c289b..b75e047 100644 --- a/django_pandas/io.py +++ b/django_pandas/io.py @@ -3,15 +3,13 @@ def to_fields(qs, fieldnames): - fields = [] for fieldname in fieldnames: model = qs.model for fieldname_part in fieldname.split('__'): field = model._meta.get_field(fieldname_part) if field.get_internal_type() == 'ForeignKey': model = field.rel.to - fields.append(field) - return fields + yield field def read_frame(qs, fieldnames=(), index_col=None, coerce_float=False, diff --git a/django_pandas/managers.py b/django_pandas/managers.py index a74174c..5d1e16a 100644 --- a/django_pandas/managers.py +++ b/django_pandas/managers.py @@ -125,10 +125,8 @@ def to_timeseries(self, fieldnames=(), verbose=True, defined in the ``__unicode__`` or ``__str__`` methods of the related class definition """ - if index is None: - raise AssertionError('You must supply an index field') - if storage not in ('wide', 'long'): - raise AssertionError('storage must be wide or long') + assert index is not None, 'You must supply an index field' + assert storage in ('wide', 'long'), 'storage must be wide or long' if rs_kwargs is None: rs_kwargs = {} @@ -136,11 +134,8 @@ def to_timeseries(self, fieldnames=(), verbose=True, df = self.to_dataframe(fieldnames, verbose=verbose, index=index) else: df = self.to_dataframe(fieldnames, verbose=verbose) - if values is None: - raise AssertionError('You must specify a values field') - - if pivot_columns is None: - raise AssertionError('You must specify pivot_columns') + assert values is not None, 'You must specify a values field' + assert pivot_columns is not None, 'You must specify pivot_columns' if isinstance(pivot_columns, (tuple, list)): df['combined_keys'] = '' @@ -190,11 +185,8 @@ def to_dataframe(self, fieldnames=(), verbose=True, index=None, """ - df = read_frame(self, fieldnames=fieldnames, verbose=verbose, - index_col=index, - coerce_float=coerce_float) - - return df + return read_frame(self, fieldnames=fieldnames, verbose=verbose, + index_col=index, coerce_float=coerce_float) class DataFrameManager(PassThroughManager): diff --git a/django_pandas/utils.py b/django_pandas/utils.py index 2a87efa..8bd4b50 100644 --- a/django_pandas/utils.py +++ b/django_pandas/utils.py @@ -36,41 +36,29 @@ def inner(pk_list): for pk in pk_list] out_dict = cache.get_many(frozenset(cache_keys)) try: - out_list = [None if k is None else out_dict[k] - for k in cache_keys] + return [None if k is None else out_dict[k] for k in cache_keys] except KeyError: out_dict = { base_cache_key % obj.pk: force_text(obj) for obj in model.objects.filter(pk__in={pk for pk in pk_list if not isnan(pk)})} cache.set_many(out_dict) - out_list = list(map(out_dict.get, cache_keys)) - return out_list + return list(map(out_dict.get, cache_keys)) return inner def build_update_functions(fieldnames, fields): - update_functions = [] - for fieldname, field in zip(fieldnames, fields): if field.choices: choices = {k: force_text(v) for k, v in field.flatchoices} - update_functions.append((fieldname, - replace_from_choices(choices))) + yield fieldname, replace_from_choices(choices) elif field.get_internal_type() == 'ForeignKey': - update_functions.append((fieldname, replace_pk(field.rel.to))) - - return update_functions + yield fieldname, replace_pk(field.rel.to) def update_with_verbose(df, fieldnames, fields): - update_functions = build_update_functions(fieldnames, fields) - - if not update_functions: - return - - for fieldname, function in update_functions: + for fieldname, function in build_update_functions(fieldnames, fields): df[fieldname] = function(df[fieldname])