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