Coverage for polar/kit/time_queries.py: 100%

1from datetime import date, datetime 

2from enum import StrEnum 

3 

4from sqlalchemy import ( 

5 CTE, 

6 Function, 

7 SQLColumnExpression, 

8 TextClause, 

9 cte, 

10 func, 

11 select, 

12 text, 

13) 

14 

15 

16class TimeInterval(StrEnum): 

17 year = "year" 

18 month = "month" 

19 week = "week" 

20 day = "day" 

21 hour = "hour" 

22 

23 def sql_interval(self) -> TextClause: 

24 return text(f"'1 {self.value}'::interval") 

25 

26 def sql_date_trunc( 

27 self, column: SQLColumnExpression[datetime] | datetime 

28 ) -> Function[datetime]: 

29 return func.date_trunc(self.value, column) 

30 

31 

32def get_timestamp_series_cte( 

33 start_timestamp: datetime, end_timestamp: datetime, interval: TimeInterval 

34) -> CTE: 

35 return cte( 

36 select( 

37 func.generate_series( 

38 start_timestamp, end_timestamp, interval.sql_interval() 

39 ).column_valued("timestamp") 

40 ) 

41 ) 

42 

43 

44MIN_DATETIME = datetime(2023, 1, 1) # Before that, Polar didn't even exist! 🚀 

45MIN_DATE = MIN_DATETIME.date() 

46 

47MAX_INTERVAL_DAYS: dict[TimeInterval, int] = { 

48 TimeInterval.hour: 7, 

49 TimeInterval.day: 366, 

50 TimeInterval.week: 7 * 53, 

51 TimeInterval.month: 365 * 4, 

52 TimeInterval.year: 365 * 10, 

53} 

54 

55MIN_INTERVAL_DAYS: dict[TimeInterval, int] = { 

56 TimeInterval.hour: 0, 

57 TimeInterval.day: 0, 

58 TimeInterval.week: 14, 

59 TimeInterval.month: 60, 

60 TimeInterval.year: 366, 

61} 

62 

63 

64def is_under_limits(start_date: date, end_date: date, interval: TimeInterval) -> bool: 

65 return end_date.toordinal() - start_date.toordinal() <= MAX_INTERVAL_DAYS[interval]