Skip to content

Commit

Permalink
Test: read_frame doesn't populate cache, handles values_list correctly
Browse files Browse the repository at this point in the history
Fixed a crash when the query set was the result of qs.values_list()
  • Loading branch information
wkschwartz committed Mar 28, 2018
1 parent 4a28342 commit cfd45c8
Show file tree
Hide file tree
Showing 2 changed files with 12 additions and 3 deletions.
2 changes: 1 addition & 1 deletion django_pandas/io.py
Original file line number Diff line number Diff line change
Expand Up @@ -200,7 +200,7 @@ def read_frame(qs, fieldnames=(), index_col=None, coerce_float=False,
fields = qs.model._meta.fields
fieldnames = [f.name for f in fields]

if not is_values_queryset(qs):
if not issubclass(qs._iterable_class, django.db.models.query.ValuesListIterable):
qs = qs.values_list(*fieldnames)
recs = qs.iterator()

Expand Down
13 changes: 11 additions & 2 deletions django_pandas/tests/test_io.py
Original file line number Diff line number Diff line change
Expand Up @@ -47,8 +47,7 @@ def test_basic(self):
df1 = read_frame(qs, ['col1', 'col2'])
self.assertEqual(df1.shape, (qs.count(), 2))

def test_compress_basic(self):
qs = MyModel.objects.all()
def assert_compress_basic(self, qs):
df = read_frame(qs, compress=True)

# Test automatic inference of dtypes
Expand All @@ -59,6 +58,14 @@ def test_compress_basic(self):

# Compress should use less memory
self.assertLess(df.memory_usage().sum(), read_frame(qs).memory_usage().sum())
# Uses qs.iterator() rather than for x in qs.
self.assertFalse(qs._result_cache)

def test_compress_basic(self):
qs = MyModel.objects.all()
self.assert_compress_basic(qs)
self.assert_compress_basic(qs.values())
self.assert_compress_basic(qs.values_list())

def test_compress_bad_argument(self):
qs = MyModel.objects.all()
Expand Down Expand Up @@ -103,6 +110,8 @@ def test_compress_custom_field(self):
# Memory usage is ordered as df1 < df2 < read_frame(qs, compress=False)
self.assertLess(df2.memory_usage().sum(), read_frame(qs).memory_usage().sum())
self.assertLess(df1.memory_usage().sum(), df2.memory_usage().sum())
# Uses qs.iterator() rather than for x in qs.
self.assertFalse(qs._result_cache)

def test_values(self):
qs = MyModel.objects.all()
Expand Down

0 comments on commit cfd45c8

Please sign in to comment.