from django_filters import rest_framework as filter


class VisitRecordFilter(filter.FilterSet):
    year_month = filter.CharFilter(method="get_year_month")

    def get_year_month(self, queryset, name, value):
        if value:
            year = value.split("-")[0]
            month = value.split("-")[1]
            queryset = queryset.filter(visit_time__year=year, visit_time__month=int(month))
        return queryset
