From 4153640aa572b4fe2d3aaea319801ce56685c2c5 Mon Sep 17 00:00:00 2001
From: Janusz Gregorczyk <janusz@doist.com>
Date: Tue, 29 Oct 2024 14:39:09 +0100
Subject: [PATCH] chore: Improve types.

---
 bitmapist/__init__.py        | 22 ++++++++++-----
 bitmapist/cohort/__init__.py | 52 ++++++++++++++++++------------------
 2 files changed, 42 insertions(+), 32 deletions(-)

diff --git a/bitmapist/__init__.py b/bitmapist/__init__.py
index 677ffd9..e276fea 100644
--- a/bitmapist/__init__.py
+++ b/bitmapist/__init__.py
@@ -198,13 +198,13 @@ def unmark_event(
 
 def _mark(
     event_name,
-    uuid,
+    uuid: int,
     system="default",
     now=None,
     track_hourly=None,
     track_unique=None,
-    use_pipeline=True,
-    value=1,
+    use_pipeline: bool = True,
+    value: int = 1,
 ):
     if track_hourly is None:
         track_hourly = TRACK_HOURLY
@@ -214,7 +214,17 @@ def _mark(
     if not now:
         now = datetime.utcnow()
 
-    obj_classes = [MonthEvents, WeekEvents, DayEvents]
+    obj_classes: list[
+        type[MonthEvents]
+        | type[WeekEvents]
+        | type[DayEvents]
+        | type[HourEvents]
+        | type[UniqueEvents]
+    ] = [
+        MonthEvents,
+        WeekEvents,
+        DayEvents,
+    ]
     if track_hourly:
         obj_classes.append(HourEvents)
     if track_unique:
@@ -228,7 +238,7 @@ def _mark(
         client.setbit(obj_class.from_date(event_name, now).redis_key, uuid, value)
 
     if use_pipeline:
-        client.execute()
+        client.execute()  # type: ignore[attr-defined] # type dependent on conditional
 
 
 def mark_unique(event_name: str, uuid: int, system: str = "default") -> None:
@@ -271,7 +281,7 @@ def unmark_unique(event_name: str, uuid: int, system: str = "default") -> None:
     _mark_unique(event_name, uuid, system, value=0)
 
 
-def _mark_unique(event_name, uuid, system="default", value=1):
+def _mark_unique(event_name, uuid: int, system="default", value: int = 1):
     get_redis(system).setbit(UniqueEvents(event_name).redis_key, uuid, value)
 
 
diff --git a/bitmapist/cohort/__init__.py b/bitmapist/cohort/__init__.py
index 3682888..d0b0706 100644
--- a/bitmapist/cohort/__init__.py
+++ b/bitmapist/cohort/__init__.py
@@ -64,7 +64,7 @@
 
 from datetime import date, datetime, timedelta
 from os import path
-from typing import Any, Callable, Optional
+from typing import Any, Callable, Literal, Optional
 
 from dateutil.relativedelta import relativedelta
 from mako.lookup import TemplateLookup
@@ -88,15 +88,15 @@ def render_html_form(
     selections1b=None,
     selections2=None,
     selections2b=None,
-    time_group="days",
+    time_group: Literal["days", "weeks", "months", "years"] = "days",
     select1=None,
     select1b=None,
     select2=None,
     select2b=None,
-    as_percent=1,
-    num_results=31,
-    num_of_rows=12,
-    start_date=None,
+    as_percent: bool = True,
+    num_results: int = 31,
+    num_of_rows: int = 12,
+    start_date: str | None = None,
 ):
     """
     Render a HTML form that can be used to query the data in bitmapist.
@@ -147,11 +147,11 @@ def render_html_form(
 
 def render_html_data(
     dates_data,
-    as_percent=True,
-    time_group="days",
+    as_percent: bool = True,
+    time_group: Literal["days", "weeks", "months", "years"] = "days",
     num_results: int = 31,
     num_of_rows: int = 12,
-    start_date=None,
+    start_date: str | None = None,
 ):
     """
     Render's data as HTML, inside a TABLE element.
@@ -176,11 +176,11 @@ def render_html_data(
 
 def render_csv_data(
     dates_data,
-    as_percent=True,
-    time_group="days",
+    as_percent: bool = True,
+    time_group: Literal["days", "weeks", "months", "years"] = "days",
     num_results: int = 31,
     num_of_rows: int = 12,
-    start_date=None,
+    start_date: str | None = None,
 ):
     """Render's data as CSV."""
     return (
@@ -205,12 +205,12 @@ def get_dates_data(
     select2,
     select1b=None,
     select2b=None,
-    time_group="days",
+    time_group: Literal["days", "weeks", "months", "years"] = "days",
     system="default",
-    as_percent=1,
-    num_results=31,
-    num_of_rows=12,
-    start_date=None,
+    as_percent: bool = True,
+    num_results: int = 31,
+    num_of_rows: int = 12,
+    start_date: str | None = None,
 ):
     """
     Fetch the data from bitmapist.
@@ -240,8 +240,8 @@ def get_dates_data(
         date_range = num_results
         now = now - timedelta(days=num_results - 1)
 
-        def timedelta_inc(d):
-            return timedelta(days=d)
+        def timedelta_inc(delta: int) -> relativedelta | timedelta:
+            return timedelta(days=delta)
 
     # Weeks
     elif time_group == "weeks":
@@ -250,8 +250,8 @@ def timedelta_inc(d):
         date_range = num_results
         now = now - relativedelta(weeks=num_results - 1)
 
-        def timedelta_inc(w):
-            return relativedelta(weeks=w)
+        def timedelta_inc(delta: int) -> relativedelta | timedelta:
+            return relativedelta(weeks=delta)
 
     # Months
     elif time_group == "months":
@@ -261,8 +261,8 @@ def timedelta_inc(w):
         now = now - relativedelta(months=num_results - 1)
         now -= timedelta(days=now.day - 1)
 
-        def timedelta_inc(m):
-            return relativedelta(months=m)
+        def timedelta_inc(delta: int) -> relativedelta | timedelta:
+            return relativedelta(months=delta)
 
     # Years
     elif time_group == "years":
@@ -273,13 +273,13 @@ def timedelta_inc(m):
         date_range = num_results
         now = now - relativedelta(years=num_results - 1)
 
-        def timedelta_inc(m):
-            return relativedelta(years=m)
+        def timedelta_inc(delta: int) -> relativedelta | timedelta:
+            return relativedelta(years=delta)
 
     dates = []
 
     for _i in range(date_range):
-        result = [now]
+        result: list[Any] = [now]
 
         # events for select1 (+select1b)
         select1_events = fn_get_events(select1, now, system)