Skip to content

Commit

Permalink
Merge pull request #67 from Doist/proxi/improve-types
Browse files Browse the repository at this point in the history
chore: Improve types.
  • Loading branch information
proxi authored Oct 29, 2024
2 parents 8d23662 + 4153640 commit 4cbd5d5
Show file tree
Hide file tree
Showing 2 changed files with 42 additions and 32 deletions.
22 changes: 16 additions & 6 deletions bitmapist/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -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:
Expand All @@ -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:
Expand Down Expand Up @@ -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)


Expand Down
52 changes: 26 additions & 26 deletions bitmapist/cohort/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -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.
Expand Down Expand Up @@ -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.
Expand All @@ -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 (
Expand All @@ -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.
Expand Down Expand Up @@ -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":
Expand All @@ -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":
Expand All @@ -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":
Expand All @@ -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)
Expand Down

0 comments on commit 4cbd5d5

Please sign in to comment.