From bb4b4d20d43cdae42df06553de2908e035d566cf Mon Sep 17 00:00:00 2001
From: Glenn Matthews <glenn.matthews@networktocode.com>
Date: Wed, 29 Nov 2023 14:54:26 -0500
Subject: [PATCH] Handle tree queryset .values() even more correctly

---
 tests/testapp/test_queries.py | 21 +++++++++++++++++++++
 tree_queries/compiler.py      |  5 ++---
 2 files changed, 23 insertions(+), 3 deletions(-)

diff --git a/tests/testapp/test_queries.py b/tests/testapp/test_queries.py
index 6f4d875..b4aeca0 100644
--- a/tests/testapp/test_queries.py
+++ b/tests/testapp/test_queries.py
@@ -149,6 +149,20 @@ def test_update_aggregate(self):
         )
 
     def test_values(self):
+        tree = self.create_tree()
+        self.assertEqual(
+            list(Model.objects.with_tree_fields().values("name")),
+            [
+                {"name": "root"},
+                {"name": "1"},
+                {"name": "1-1"},
+                {"name": "2"},
+                {"name": "2-1"},
+                {"name": "2-2"},
+            ],
+        )
+
+    def test_values_ancestors(self):
         tree = self.create_tree()
         self.assertEqual(
             list(Model.objects.ancestors(tree.child2_1).values()),
@@ -169,6 +183,13 @@ def test_values(self):
         )
 
     def test_values_list(self):
+        tree = self.create_tree()
+        self.assertEqual(
+            list(Model.objects.with_tree_fields().values_list("name", flat=True)),
+            ["root", "1", "1-1", "2", "2-1", "2-2"],
+        )
+
+    def test_values_list_ancestors(self):
         tree = self.create_tree()
         self.assertEqual(
             list(
diff --git a/tree_queries/compiler.py b/tree_queries/compiler.py
index 37b6172..8d0579c 100644
--- a/tree_queries/compiler.py
+++ b/tree_queries/compiler.py
@@ -211,7 +211,6 @@ def as_sql(self, *args, **kwargs):
         # problem but I just gotta stop worrying and trust the testsuite.
         skip_tree_fields = (
             (self.query.distinct and self.query.subquery)
-            or self.query.values_select
             or any(  # pragma: no branch
                 # OK if generator is not consumed completely
                 annotation.is_summary
@@ -244,9 +243,9 @@ def as_sql(self, *args, **kwargs):
 
             self.query.add_extra(
                 # Do not add extra fields to the select statement when it is a
-                # summary query
+                # summary query or when using .values() or .values_list()
                 select={}
-                if skip_tree_fields
+                if skip_tree_fields or self.query.values_select
                 else {
                     "tree_depth": "__tree.tree_depth",
                     "tree_path": "__tree.tree_path",