Added Balance by month, year or total

budget.py

@@ -16,7 +16,8 @@ menu_options = {
     2: 'List Bills',
     3: 'Add Donation',
     4: 'List Donations',
-    5: 'Exit',
+    5: 'Balance',
+    6: 'Exit',
 }
 list_show_options = {
     1: 'Month',
@@ -381,6 +382,126 @@ class DonationShow:
 
         return (month_date, month_donation)
 
+class BalanceShow:
+
+    name = "BalanceShow"
+
+    def __init__(self):
+
+        while(True):
+
+            print('\n')
+
+            list_show_menu()
+
+            show_option = ''
+
+            now = datetime.now(tz)
+
+            try:
+
+                show_option = int(input('\nEnter your choice: '))
+
+            except:
+
+                print('\nWrong input. Please enter a number between 1 and 4.\n')
+
+            if show_option == 1:
+
+                self.option = 1
+                self.criteria = input('Month (enter = current month)? ') or str(now.month)
+                break
+
+            elif show_option == 2:
+
+                self.option = 2
+                self.criteria = input('Year (enter = current year)? ') or str(now.year)
+                break
+
+            elif show_option == 3:
+
+                self.option = 3
+                self.criteria = None
+                break
+
+            elif show_option == 4:
+
+                self.option = 4
+                self.criteria = None
+                break
+
+    def show(self, criteria):
+
+        balance = 0
+
+        if self.option == 1:
+
+            incomes_sql = "select sum(donations) from incomes where date_part('year', datetime) = date_part('year', CURRENT_DATE) and date_part('month', datetime) = (%s)"
+
+            bills_sql  = "select sum(coalesce(domain) + coalesce(server) + coalesce(backup) + coalesce(fileserver) + coalesce(setup)) from bills where date_part('year', datetime) = date_part('year', CURRENT_DATE) and date_part('month', datetime) = (%s)"
+
+        elif self.option == 2:
+
+            incomes_sql  = "select sum(donations) from incomes where date_part('year', datetime) = (%s)"
+
+            bills_sql  = "select sum(coalesce(domain) + coalesce(server) + coalesce(backup) + coalesce(fileserver) + coalesce(setup)) from bills where date_part('year', datetime) = (%s)"
+
+        elif self.option == 3:
+
+            incomes_sql  = "select sum(donations) from incomes order by 1 asc"
+
+            bills_sql = "select sum(coalesce(domain) + coalesce(server) + coalesce(backup) + coalesce(fileserver) + coalesce(setup)) from bills"
+
+        try:
+
+            conn = psycopg2.connect(database = budget_db, user = budget_db_user, password = "", host = "/var/run/postgresql", port = "5432")
+
+            cur = conn.cursor()
+
+            if self.option == 1 or self.option == 2:
+
+                cur.execute(incomes_sql, (self.criteria,))
+
+                row = cur.fetchone()
+
+                incomes = row[0]
+
+                cur.execute(bills_sql, (self.criteria,))
+
+                row = cur.fetchone()
+
+                bills = row[0]
+
+            else:
+
+                cur.execute(incomes_sql)
+
+                row = cur.fetchone()
+
+                incomes = row[0]
+
+                cur.execute(bills_sql)
+
+                row = cur.fetchone()
+
+                bills = row[0]
+
+            balance = incomes - bills
+
+            cur.close()
+
+        except (Exception, psycopg2.DatabaseError) as error:
+
+            print(error)
+
+        finally:
+
+            if conn is not None:
+
+                conn.close()
+
+        return (balance)
+
 def get_config():
 
     global budget_db
@@ -537,6 +658,28 @@ if __name__ == '__main__':
 
         elif option == 5:
 
+            balance = BalanceShow()
+
+            if balance.option == 1 or balance.option == 2 or balance.option == 3: 
+
+                balance_result  = balance.show(balance.criteria)
+
+            if balance.option != 4:
+
+                print_table = PrettyTable()
+
+                print_table.field_names = ['Balance']
+
+                print_table.align['Balance'] = "r"
+
+                b_amount = float(balance_result)
+
+                print_table.add_row([b_amount])
+
+                print(print_table)
+
+        elif option == 6:
+
             print('Bye!')
 
             exit()