<s> from datetime import date , datetime , time , timedelta <EOL> from django . contrib . auth . models import User <EOL> from django_lean . lean_retention . models import DailyActivity <EOL> def sort_retention_periods ( retention_periods ) : <EOL> result = list ( sorted ( set ( p + <NUM_LIT:0> for p in retention_periods ) ) ) <EOL> if result and result [ <NUM_LIT:0> ] < <NUM_LIT:1> : <EOL> raise ValueError ( '<STR_LIT>' <EOL> '<STR_LIT>' % result [ <NUM_LIT:0> ] ) <EOL> return result <EOL> class Period ( object ) : <EOL> def __init__ ( self , cohort , start_day , end_day ) : <EOL> self . cohort = cohort <EOL> self . start_day = start_day <EOL> self . end_day = end_day <EOL> if self . start_day < <NUM_LIT:1> : <EOL> raise ValueError ( "<STR_LIT>" % self . start_day ) <EOL> if self . start_day >= self . end_day : <EOL> raise ValueError ( "<STR_LIT>" % <EOL> ( self . start_day , self . end_day ) ) <EOL> self . _activities = None <EOL> self . _users = None <EOL> def length ( self ) : <EOL> return self . end_day - self . start_day <EOL> @ property <EOL> def activities ( self ) : <EOL> if self . _activities is None : <EOL> self . _activities = DailyActivity . objects . filter ( <EOL> user__in = self . cohort . users , <EOL> days__range = ( self . start_day , self . end_day - <NUM_LIT:1> ) <EOL> ) <EOL> return self . _activities <EOL> @ property <EOL> def users ( self ) : <EOL> if self . _users is None : <EOL> self . _users = User . objects . filter ( <EOL> id__in = self . activities . values ( '<STR_LIT:user>' ) <EOL> ) <EOL> return self . _users <EOL> @ classmethod <EOL> def periods ( cls , cohort , retention_periods ) : <EOL> last = <NUM_LIT:1> <EOL> for period in sort_retention_periods ( retention_periods ) : <EOL> yield cls ( cohort = cohort , start_day = last , end_day = period ) <EOL> last = period <EOL> class Cohort ( object ) : <EOL> def __init__ ( self , start_date , end_date , retention_periods , <EOL> period_class = Period ) : <EOL> if hasattr ( start_date , '<STR_LIT:date>' ) : <EOL> start_date = start_date . date ( ) <EOL> if hasattr ( end_date , '<STR_LIT:date>' ) : <EOL> end_date = end_date . date ( ) <EOL> self . start_date = start_date <EOL> self . end_date = end_date <EOL> if self . start_date > self . end_date : <EOL> raise ValueError ( "<STR_LIT>" % <EOL> ( self . start_date , self . end_date ) ) <EOL> self . retention_periods = sort_retention_periods ( retention_periods ) <EOL> self . _Period = Period <EOL> self . _periods = None <EOL> self . _users = None <EOL> @ property <EOL> def periods ( self ) : <EOL> if self . _periods is None : <EOL> self . _periods = list ( <EOL> self . _Period . periods ( cohort = self , <EOL> retention_periods = self . retention_periods ) <EOL> ) <EOL> return self . _periods <EOL> @ property <EOL> def users ( self ) : <EOL> if self . _users is None : <EOL> start = datetime . combine ( self . start_date , time ( <NUM_LIT:0> , <NUM_LIT:0> , <NUM_LIT:0> ) ) <EOL> end = datetime . combine ( self . end_date , time ( <NUM_LIT> , <NUM_LIT> , <NUM_LIT> ) ) <EOL> self . _users = User . objects . filter ( date_joined__range = ( start , end ) ) <EOL> return self . _users <EOL> @ classmethod <EOL> def cohorts ( cls , end_date , length , retention_periods , <EOL> period_class = Period ) : <EOL> if hasattr ( end_date , '<STR_LIT:date>' ) : <EOL> end_date = end_date . date ( ) <EOL> retention_periods = sort_retention_periods ( retention_periods ) <EOL> min_period = retention_periods [ <NUM_LIT:0> ] if retention_periods else <NUM_LIT:0> <EOL> end_date -= timedelta ( days = min_period ) <EOL> one_day = timedelta ( days = <NUM_LIT:1> ) <EOL> start_date = end_date - timedelta ( days = length ) + one_day <EOL> try : <EOL> while True : <EOL> yield cls ( start_date = start_date , end_date = end_date , <EOL> retention_periods = retention_periods ) <EOL> start_date -= one_day <EOL> end_date -= one_day <EOL> except OverflowError : <EOL> raise StopIteration </s>
