diff --git a/django_pandas/io.py b/django_pandas/io.py index 8764594..2e3811f 100644 --- a/django_pandas/io.py +++ b/django_pandas/io.py @@ -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() diff --git a/django_pandas/tests/test_io.py b/django_pandas/tests/test_io.py index 95434a2..07d3700 100644 --- a/django_pandas/tests/test_io.py +++ b/django_pandas/tests/test_io.py @@ -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 @@ -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() @@ -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()