From c2a1ae661db07e9db7c88caa20bb469b49d99b56 Mon Sep 17 00:00:00 2001
From: caoqianming <caoqianming@foxmail.com>
Date: Thu, 7 Nov 2024 17:50:48 +0800
Subject: [PATCH] =?UTF-8?q?feat:=20base=20=E6=94=AF=E6=8C=81=E6=A0=91?=
 =?UTF-8?q?=E5=9E=8B=E7=9A=84=E9=80=92=E5=BD=92=E5=88=86=E9=A1=B5=E6=9F=A5?=
 =?UTF-8?q?=E8=AF=A2=E4=BB=A5=E4=BC=A0with=5Fchildren=E5=8F=82=E6=95=B0?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 apps/utils/mixins.py      |  2 ++
 apps/utils/serializers.py | 25 ++++++++++++++++++++++++-
 apps/utils/viewsets.py    |  4 ++++
 3 files changed, 30 insertions(+), 1 deletion(-)

diff --git a/apps/utils/mixins.py b/apps/utils/mixins.py
index 93f747a3..ab9632a8 100755
--- a/apps/utils/mixins.py
+++ b/apps/utils/mixins.py
@@ -179,6 +179,8 @@ class CustomListModelMixin(ListModelMixin):
     @swagger_auto_schema(manual_parameters=[
         openapi.Parameter(name="query", in_=openapi.IN_QUERY, description="定制返回数据",
                           type=openapi.TYPE_STRING, required=False),
+        openapi.Parameter(name="with_children", in_=openapi.IN_QUERY, description="带有children(yes/no/count)",
+                          type=openapi.TYPE_STRING, required=False),
     ])
     def list(self, request, *args, **kwargs):
         """
diff --git a/apps/utils/serializers.py b/apps/utils/serializers.py
index 247438f1..77db28cd 100755
--- a/apps/utils/serializers.py
+++ b/apps/utils/serializers.py
@@ -19,7 +19,30 @@ class GenSignatureSerializer(serializers.Serializer):
     path = serializers.CharField(label="图片地址")
 
 
-class CustomModelSerializer(DynamicFieldsMixin, serializers.ModelSerializer):
+class TreeSerializerMixin:
+
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        request = self.context.get('request', None)
+        self.with_children = request.query_params.get('with_children', 'no') if request else 'no'
+        if self.with_children in ['yes', 'count']:
+            if 'children' not in self.fields:
+                self.fields['children'] = serializers.SerializerMethodField()
+            if 'children_count' not in self.fields:
+                self.fields['children_count'] = serializers.SerializerMethodField()
+
+    def get_children(self, obj):
+        if hasattr(obj, 'parent') and self.with_children == 'yes':
+            serializer_class = self.__class__
+            return serializer_class(obj.__class__.objects.filter(parent=obj), many=True, context=self.context).data
+        return []
+    
+    def get_children_count(self, obj):
+        if hasattr(obj, 'parent') and self.with_children in ['yes', 'count']:
+            return obj.__class__.objects.filter(parent=obj).count()
+        return 0
+
+class CustomModelSerializer(DynamicFieldsMixin, TreeSerializerMixin, serializers.ModelSerializer):
     """
     自定义serializer/包含创建和新增字段处理
     """
diff --git a/apps/utils/viewsets.py b/apps/utils/viewsets.py
index b8bb0785..986336c1 100755
--- a/apps/utils/viewsets.py
+++ b/apps/utils/viewsets.py
@@ -102,6 +102,10 @@ class CustomGenericViewSet(MyLoggingMixin, GenericViewSet):
 
     def filter_queryset(self, queryset):
         queryset = super().filter_queryset(queryset)
+        # 如果带有with_children查询, 出于优化需要应自动过滤掉一些内容
+        if (self.request.query_params.get("with_children", "no") in ["yes", "count"] 
+            and self.request.query_params.get("parent", None) is None):
+            queryset = queryset.filter(parent=None)
         # 用于性能优化
         if self.select_related_fields:
             queryset = queryset.select_related(*self.select_related_fields)