BankAccount.py

   1 import pickle
   2
   3 ACCOUNT_FILE = 'bank.pck'
   4
   5 class BankAccount:
   6         def __init__(self):
   7                 f = None
   8                 self.bank = {}
   9                 try:
  10                         f = open(ACCOUNT_FILE)
  11                 except IOError:
  12                         pass
  13                 if f != None:
  14                         self.bank = pickle.load(f)
  15                         self.sanity_check(self.bank)
  16                         f.close()
  17                 self.save()
  18
  19         def sanity_check_user(self, user):
  20                 wanted_fields = ['balance']
  21                 for f in user:
  22                         if f in wanted_fields:
  23                                 if f == 'balance':
  24                                         if not isinstance(user['balance'], tuple):
  25                                                 return False
  26                                         if len(user['balance']) != 2:
  27                                                 return False
  28                                         if not isinstance(user[0], int) or \
  29                                                         not isinstance(user[1], int):
  30                                                 return False
  31                                 del wanted_fields[f]
  32                         else:
  33                                 return False
  34                 if len(wanted_fields) != 0:
  35                         return False
  36                 return True
  37
  38         def sanity_check(self, bank):
  39                 for u in bank:
  40                         if not self.sanity_check_user(u):
  41                                 return False
  42                 return True
  43
  44         def save(self):
  45                 f = open(ACCOUNT_FILE, 'w')
  46                 pickle.dump(self.bank, f)
  47                 f.close()
  48
  49         def ensure_user_exists(self, username):
  50                 if self.bank.has_key(username):
  51                         return
  52                 self.bank[username] = {}
  53                 self.bank[username]['balance'] = (0,0)
  54
  55         def get_balance(self, username):
  56                 if self.bank.has_key(username):
  57                         return self.bank[username]['balance']
  58                 return (0,0)
  59
  60         def add_amount(self, username, amount):
  61                 self.ensure_user_exists(username)
  62                 (cur_cents, cur_bytes) = self.bank[username]['balance']
  63                 (add_cents, add_bytes) = amount
  64                 self.bank[username]['balance'] = (cur_cents+add_cents,
  65                         cur_bytes+add_bytes)
  66                 self.save()