From 64c2d86cd74848dec21c695137f1f830fe1159f6 Mon Sep 17 00:00:00 2001
From: caoqianming <caoqianming@foxmail.com>
Date: Thu, 31 Aug 2023 14:14:00 +0800
Subject: [PATCH] =?UTF-8?q?feat:=20enstat2=E5=A2=9E=E5=8A=A0=E6=9F=A5?=
 =?UTF-8?q?=E8=AF=A2=E6=9D=A1=E4=BB=B6?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 apps/enm/filters.py | 44 +++++++++++++++++++++++++++++++++++++++++++-
 apps/enm/views.py   |  4 ++--
 2 files changed, 45 insertions(+), 3 deletions(-)

diff --git a/apps/enm/filters.py b/apps/enm/filters.py
index 811638d9..a1151fa3 100644
--- a/apps/enm/filters.py
+++ b/apps/enm/filters.py
@@ -1,5 +1,5 @@
 from django_filters import rest_framework as filters
-from apps.enm.models import MpointStat, EnStat
+from apps.enm.models import MpointStat, EnStat, EnStat2
 from datetime import datetime
 from django.db.models import F, Q
 
@@ -156,4 +156,46 @@ class EnStatFilter(filters.FilterSet):
         elif query['type'] in ['year_s']:
             queryset = queryset.filter(
             Q(year_s__lte=end_year))
+        return queryset
+    
+class EnStat2Filter(filters.FilterSet):
+    start_time = filters.CharFilter(method='filter_start_time')
+    end_time = filters.CharFilter(method='filter_end_time')
+    class Meta:
+        model = EnStat2
+        fields = {
+            "type": ["exact"],
+            "day_s": ["exact"],
+            "month_s": ["exact"],
+            "year_s": ["exact"],
+        }
+
+    def filter_start_time(self, queryset, name, value):
+        query = self.data
+        start_time = datetime.strptime(value, '%Y-%m-%d %H:%M:%S')
+        start_year, start_month, start_day, start_hour = start_time.year, start_time.month, start_time.day, start_time.hour
+        if query['type'] in ['day_s']:
+            queryset = queryset.filter(
+            Q(year_s__gt=start_year) | 
+            Q(year_s=start_year, month_s__gt=start_month) |
+            Q(year_s=start_year, month_s=start_month, day_s__gte=start_day))
+        elif query['type'] in ['month_s']:
+            queryset = queryset.filter(
+            Q(year_s__gt=start_year) | 
+            Q(year_s=start_year, month_s__gte=start_month))
+        return queryset
+
+    def filter_end_time(self, queryset, name, value):
+        query = self.data
+        end_time = datetime.strptime(value, '%Y-%m-%d %H:%M:%S')
+        end_year, end_month, end_day, end_hour = end_time.year, end_time.month, end_time.day, end_time.hour
+        if query['type'] in ['day_s']:
+            queryset = queryset.filter(
+            Q(year_s__lt=end_year) | 
+            Q(year_s=end_year, month_s__lt=end_month) |
+            Q(year_s=end_year, month_s=end_month, day_s__lte=end_day))
+        elif query['type'] in ['month_s']:
+            queryset = queryset.filter(
+            Q(year_s__lt=end_year) | 
+            Q(year_s=end_year, month_s__lte=end_month))
         return queryset
\ No newline at end of file
diff --git a/apps/enm/views.py b/apps/enm/views.py
index efd789af..77a6b807 100644
--- a/apps/enm/views.py
+++ b/apps/enm/views.py
@@ -4,7 +4,7 @@ from apps.utils.viewsets import CustomModelViewSet, CustomGenericViewSet
 from rest_framework.mixins import ListModelMixin
 from apps.utils.mixins import BulkCreateModelMixin, BulkDestroyModelMixin
 from apps.enm.serializers import (MpointSerializer, MpLogSerializer, MpointStatSerializer, EnStatSerializer, EnStat2Serializer, ReCalSerializer)
-from apps.enm.filters import MpointStatFilter, EnStatFilter
+from apps.enm.filters import MpointStatFilter, EnStatFilter, EnStat2Filter
 from apps.enm.tasks import cal_mpointstat_manual
 from rest_framework.response import Response
 from rest_framework.decorators import action
@@ -94,4 +94,4 @@ class EnStat2ViewSet(ListModelMixin, CustomGenericViewSet):
     perms_map = {'get': '*'}
     queryset = EnStat2.objects.all()
     serializer_class = EnStat2Serializer
-    filterset_fields = ['year_s', 'month_s', 'day_s', 'type']
\ No newline at end of file
+    filterset_class = EnStat2Filter
\ No newline at end of file