| 1 |  |  | from functools import partial | 
            
                                                                                                            
                            
            
                                    
            
            
                | 2 |  |  |  | 
            
                                                                                                            
                            
            
                                    
            
            
                | 3 |  |  | import arrow | 
            
                                                                                                            
                            
            
                                    
            
            
                | 4 |  |  | import telegram | 
            
                                                                                                            
                            
            
                                    
            
            
                | 5 |  |  |  | 
            
                                                                                                            
                            
            
                                    
            
            
                | 6 |  |  | from . import BotCommand | 
            
                                                                                                            
                            
            
                                    
            
            
                | 7 |  |  |  | 
            
                                                                                                            
                            
            
                                    
            
            
                | 8 |  |  |  | 
            
                                                                                                            
                            
            
                                    
            
            
                | 9 |  |  | class SplitCommand(BotCommand): | 
            
                                                                                                            
                            
            
                                    
            
            
                | 10 |  |  |     command = '/split' | 
            
                                                                                                            
                            
            
                                    
            
            
                | 11 |  |  |  | 
            
                                                                                                            
                            
            
                                    
            
            
                | 12 |  |  |     def __init__(self, *args, **kwargs): | 
            
                                                                                                            
                            
            
                                    
            
            
                | 13 |  |  |         self.commands = [ | 
            
                                                                                                            
                            
            
                                    
            
            
                | 14 |  |  |             {'text': "Split All", | 
            
                                                                                                            
                            
            
                                    
            
            
                | 15 |  |  |              'function': self.get_all_split}, | 
            
                                                                                                            
                            
            
                                    
            
            
                | 16 |  |  |             {'text': 'Split this Month', | 
            
                                                                                                            
                            
            
                                    
            
            
                | 17 |  |  |              'function': self.get_this_month_split}, | 
            
                                                                                                            
                            
            
                                    
            
            
                | 18 |  |  |             {'text': 'Split last Month', | 
            
                                                                                                            
                            
            
                                    
            
            
                | 19 |  |  |              'function': self.get_last_month_split}, | 
            
                                                                                                            
                            
            
                                    
            
            
                | 20 |  |  |             ] | 
            
                                                                                                            
                            
            
                                    
            
            
                | 21 |  |  |         super(SplitCommand, self).__init__(*args, **kwargs) | 
            
                                                                                                            
                                                                
            
                                    
            
            
                | 22 |  |  |  | 
            
                                                                        
                            
            
                                    
            
            
                | 23 |  |  |     def default(self, message): | 
            
                                                                        
                            
            
                                    
            
            
                | 24 |  |  |         keyboard = [["Split All"], | 
            
                                                                        
                            
            
                                    
            
            
                | 25 |  |  |                     ["Split this Month"], | 
            
                                                                        
                            
            
                                    
            
            
                | 26 |  |  |                     ["Split last Month"] | 
            
                                                                        
                            
            
                                    
            
            
                | 27 |  |  |                     ] | 
            
                                                                        
                            
            
                                    
            
            
                | 28 |  |  |         reply_markup = telegram.ReplyKeyboardMarkup( | 
            
                                                                        
                            
            
                                    
            
            
                | 29 |  |  |             keyboard, resize_keyboard=True, | 
            
                                                                        
                            
            
                                    
            
            
                | 30 |  |  |             one_time_keyboard=True, selective=True) | 
            
                                                                        
                            
            
                                    
            
            
                | 31 |  |  |         msg = self.bot.say(message, 'Choose', reply_markup=reply_markup) | 
            
                                                                        
                            
            
                                    
            
            
                | 32 |  |  |         self.queue(msg, partial(self.process_which_split)) | 
            
                                                                                                            
                            
            
                                    
            
            
                | 33 |  |  |  | 
            
                                                                                                            
                            
            
                                    
            
            
                | 34 |  |  |     def process_which_split(self, message): | 
            
                                                                                                            
                            
            
                                    
            
            
                | 35 |  |  |         if not message.text: | 
            
                                                                                                            
                            
            
                                    
            
            
                | 36 |  |  |             self.bot.say(message, 'Nope') | 
            
                                                                                                            
                            
            
                                    
            
            
                | 37 |  |  |             return | 
            
                                                                                                            
                            
            
                                    
            
            
                | 38 |  |  |  | 
            
                                                                                                            
                            
            
                                    
            
            
                | 39 |  |  |         for command in self.commands: | 
            
                                                                                                            
                            
            
                                    
            
            
                | 40 |  |  |             if message.text == command['text']: | 
            
                                                                                                            
                            
            
                                    
            
            
                | 41 |  |  |                 command['function'](message) | 
            
                                                                                                            
                            
            
                                    
            
            
                | 42 |  |  |  | 
            
                                                                                                            
                            
            
                                    
            
            
                | 43 |  |  |     def _split(self, entries): | 
            
                                                                                                            
                            
            
                                    
            
            
                | 44 |  |  |         total = 0 | 
            
                                                                                                            
                            
            
                                    
            
            
                | 45 |  |  |         users = {} | 
            
                                                                                                            
                            
            
                                    
            
            
                | 46 |  |  |         if not entries: | 
            
                                                                                                            
                            
            
                                    
            
            
                | 47 |  |  |             return 'No entries!' | 
            
                                                                                                            
                            
            
                                    
            
            
                | 48 |  |  |  | 
            
                                                                                                            
                            
            
                                    
            
            
                | 49 |  |  |         for entry in entries: | 
            
                                                                                                            
                            
            
                                    
            
            
                | 50 |  |  |             total += entry.amount | 
            
                                                                                                            
                            
            
                                    
            
            
                | 51 |  |  |             users[entry.user_id] = users.get(entry.user_id, 0) + entry.amount | 
            
                                                                                                            
                            
            
                                    
            
            
                | 52 |  |  |         per_person = total / len(users) | 
            
                                                                                                            
                            
            
                                    
            
            
                | 53 |  |  |  | 
            
                                                                                                            
                            
            
                                    
            
            
                | 54 |  |  |         text = '' | 
            
                                                                                                            
                            
            
                                    
            
            
                | 55 |  |  |         for user_id, amount in users.items(): | 
            
                                                                                                            
                            
            
                                    
            
            
                | 56 |  |  |             user = self._db.root.users[user_id] | 
            
                                                                                                            
                            
            
                                    
            
            
                | 57 |  |  |             text += u'{}: {:.2f}\n'.format(user.first_name, per_person - amount) | 
            
                                                                                                            
                            
            
                                    
            
            
                | 58 |  |  |         return text | 
            
                                                                                                            
                            
            
                                    
            
            
                | 59 |  |  |  | 
            
                                                                                                            
                            
            
                                    
            
            
                | 60 |  |  |     def get_all_split(self, message): | 
            
                                                                                                            
                            
            
                                    
            
            
                | 61 |  |  |         tab = self.get_tab(message.chat.id) | 
            
                                                                                                            
                            
            
                                    
            
            
                | 62 |  |  |         if not tab.users: | 
            
                                                                                                            
                            
            
                                    
            
            
                | 63 |  |  |             return | 
            
                                                                                                            
                            
            
                                    
            
            
                | 64 |  |  |         text = self._split(tab.entries) | 
            
                                                                                                            
                            
            
                                    
            
            
                | 65 |  |  |         self.bot.say(message, text, reply_markup=telegram.ReplyKeyboardHide()) | 
            
                                                                                                            
                            
            
                                    
            
            
                | 66 |  |  |  | 
            
                                                                                                            
                            
            
                                    
            
            
                | 67 |  |  |     def get_this_month_split(self, message): | 
            
                                                                                                            
                            
            
                                    
            
            
                | 68 |  |  |         tab = self.get_tab(message.chat.id) | 
            
                                                                                                            
                            
            
                                    
            
            
                | 69 |  |  |         if not tab.users: | 
            
                                                                                                            
                            
            
                                    
            
            
                | 70 |  |  |             return | 
            
                                                                                                            
                            
            
                                    
            
            
                | 71 |  |  |         from_date = arrow.now(tab.tz).floor('month') | 
            
                                                                                                            
                            
            
                                    
            
            
                | 72 |  |  |         entries = tab.get_entries(from_date) | 
            
                                                                                                            
                            
            
                                    
            
            
                | 73 |  |  |         text = self._split(entries) | 
            
                                                                                                            
                            
            
                                    
            
            
                | 74 |  |  |         self.bot.say(message, text, reply_markup=telegram.ReplyKeyboardHide()) | 
            
                                                                                                            
                            
            
                                    
            
            
                | 75 |  |  |  | 
            
                                                                                                            
                            
            
                                    
            
            
                | 76 |  |  |     def get_last_month_split(self, message): | 
            
                                                                                                            
                            
            
                                    
            
            
                | 77 |  |  |         tab = self.get_tab(message.chat.id) | 
            
                                                                                                            
                            
            
                                    
            
            
                | 78 |  |  |         if not tab.users: | 
            
                                                                                                            
                            
            
                                    
            
            
                | 79 |  |  |             return | 
            
                                                                                                            
                            
            
                                    
            
            
                | 80 |  |  |         from_date, to_date = arrow.now(tab.tz).replace(months=-1).span('month') | 
            
                                                                                                            
                            
            
                                    
            
            
                | 81 |  |  |         entries = tab.get_entries(from_date, to_date) | 
            
                                                                                                            
                            
            
                                    
            
            
                | 82 |  |  |         text = self._split(entries) | 
            
                                                                                                            
                                                                
            
                                    
            
            
                | 83 |  |  |         self.bot.say(message, text, reply_markup=telegram.ReplyKeyboardHide()) | 
            
                                                        
            
                                    
            
            
                | 84 |  |  |  |