Subclassing a dictionary

[foxfirefey]

Makes me love Python. I needed an class for a test suite for the DWMinion module I'm making to make it easy to create DWAccount objects on the fly. I figured out that I could do this by subclassing the dict class.

Basically, this subclassed dict *always* returns a DWAccount for a given account name. And I have a function to have it return a set of accounts, all at once.

class DWAccountMaker(dict):
    """A cheap class extending a dictionary that will always give back an account.  
    Start the account names with p, i, c, and f to get different account types."""

    def make_account(self, key):
        """Make an account type depending on the starting letter.
        Personal is the default otherwise."""
        if key.startswith("p"):
            return e.DWAccount(key, "personal")
        elif key.startswith("i"):
            return e.DWAccount(key, "identity")
        elif key.startswith("c"):
            return e.DWAccount(key, "community")
        elif key.startswith("f"):
            return e.DWAccount(key, "feed")
        else:
            return e.DWAccount(key, "personal")

    def get_set(self, *keys):
        """Get a set of accounts."""
        return set([self.__getitem__(key) for key in keys])

    def __getitem__( self, key ):
        # Set the key to a new account if it doesn't exist
        if not self.has_key( key ):
            super(DWAccountMaker,self).__setitem__( key, self.make_account(key) )
        return super(DWAccountMaker,self).__getitem__( key )

So, now my testing becomes a lot simpler, for example:

class DWUserBasicsTest(unittest.TestCase):
    """Test DWUser"""
    
    def setUp(self):
        self.user = e.DWUser("testuser")
        self.am = DWAccountMaker()
        
        self.user.watched_by = self.am.get_set("pa", "pb", "pc", "pd", "ia", "ib", "ic")
        self.user.watch = self.am.get_set("pa", "pc", "pe", "ca", "cb")
        self.user.trust = self.am.get_set("pa", "ib")
        self.user.trusted_by = self.am.get_set("pa", "pc", "pd", "ic")
        self.user.member_of = self.am.get_set("ca")
        
    def testName(self):
        self.assertEqual(self.user.account_name, "testuser")
        
    def testMutualWatchers(self):
        
        should_be = self.am.get_set("pa", "pc")
        self.assertEqual(self.user.get_mutual_watch(), should_be)

The subclassed DWAccountMaker dict is useful because instead of having to create lots of individual DWAccount classes, I can just let the DWAccountMaker manage the classes instead and retrieve sets from it at will.