From 52327528d2c56130b233f84e11b680725a26b998 Mon Sep 17 00:00:00 2001
From: Juho Kettunen <juho.kettunen@nordcloud.com>
Date: Tue, 22 Oct 2024 11:39:13 +0300
Subject: [PATCH] Add more type hints and docstrings

---
 leasing/report/forms.py                         | 10 +++++++++-
 leasing/report/lease/lease_statistic_report.py  |  4 ++--
 leasing/report/lease/lease_statistic_report2.py |  4 ++--
 leasing/report/report_base.py                   | 11 ++++++++---
 4 files changed, 21 insertions(+), 8 deletions(-)

diff --git a/leasing/report/forms.py b/leasing/report/forms.py
index 4f0b2334..79fb3feb 100644
--- a/leasing/report/forms.py
+++ b/leasing/report/forms.py
@@ -5,7 +5,15 @@ class ReportFormBase(forms.Form):
     """Dynamic form that initializes its fields from `input_fields` parameter"""
 
     def __init__(self, *args, **kwargs):
-        input_fields = kwargs.pop("input_fields")
+        """
+        args is expected to contain the query parameters, which in turn contains
+        the report settings when the report was requested.
+
+        kwargs is expected to contain a key "input_fields", whose value is a
+        dictionary containing the names of the query parameters and the form
+        field objects they reference.
+        """
+        input_fields: dict[str, forms.Field] = kwargs.pop("input_fields")
         super().__init__(*args, **kwargs)
 
         for field_name, field in input_fields.items():
diff --git a/leasing/report/lease/lease_statistic_report.py b/leasing/report/lease/lease_statistic_report.py
index 028cdb48..a997cc7d 100644
--- a/leasing/report/lease/lease_statistic_report.py
+++ b/leasing/report/lease/lease_statistic_report.py
@@ -4,7 +4,7 @@
 from functools import lru_cache
 
 from django import forms
-from django.db.models import Q
+from django.db.models import Q, QuerySet
 from django.utils import formats
 from django.utils.translation import gettext_lazy as _
 from enumfields.drf import EnumField
@@ -372,7 +372,7 @@ class LeaseStatisticReport(AsyncReportBase):
         },
     }
 
-    def get_data(self, input_data):
+    def get_data(self, input_data) -> QuerySet[Lease]:
         qs = Lease.objects.select_related(
             "identifier__type",
             "identifier__district",
diff --git a/leasing/report/lease/lease_statistic_report2.py b/leasing/report/lease/lease_statistic_report2.py
index b487eb08..c7fc4ba5 100644
--- a/leasing/report/lease/lease_statistic_report2.py
+++ b/leasing/report/lease/lease_statistic_report2.py
@@ -3,7 +3,7 @@
 from decimal import ROUND_HALF_UP, Decimal
 
 from django import forms
-from django.db.models import Q
+from django.db.models import Q, QuerySet
 from django.utils import formats
 from django.utils.translation import gettext_lazy as _
 from enumfields.drf import EnumField
@@ -397,7 +397,7 @@ class LeaseStatisticReport2(AsyncReportBase):
     }
     async_task_timeout = 60 * 30  # 30 minutes
 
-    def get_data(self, input_data):
+    def get_data(self, input_data) -> QuerySet[Lease]:
         qs = Lease.objects.select_related(
             "identifier__type",
             "identifier__district",
diff --git a/leasing/report/report_base.py b/leasing/report/report_base.py
index 30d42870..ebb501bd 100644
--- a/leasing/report/report_base.py
+++ b/leasing/report/report_base.py
@@ -157,14 +157,19 @@ def get_response(self, request):
 
         return Response(serialized_report_data)
 
-    def get_serializer_class(self):
+    def get_data(self, input_data: dict[str, Any]) -> list[dict] | QuerySet:
+        raise NotImplementedError(
+            "Please implement this method in the concrete report class"
+        )
+
+    def get_serializer_class(self) -> Type[ReportOutputSerializer]:
         return ReportOutputSerializer
 
-    def get_filename(self, format):
+    def get_filename(self, file_format: str):
         return "{}_{}.{}".format(
             timezone.localtime(timezone.now()).strftime("%Y-%m-%d_%H-%M"),
             self.slug,
-            format,
+            file_format,
         )
 
     def get_output_field_attr(self, field_name, attr_name, default=None):