Managing and searching objects using tags

Question

I wonder

1. Is it appropriate to hide imported classes (collections and UserDict in this case) from Python IDE (e.g. IPython)?

2. Is there a more efficient algorithm/implementation?

Please feel free to comment on how you would improve this class.

import collections as _collections
import UserDict as _UserDict

class _IdDict(_UserDict.IterableUserDict):
    def __missing__(self,key):
        raise KeyError("The item requested is not in the TagDict. "+\
                       "Perhaps more than one item were requested.")

class TagDict(object):
    ''' 
    TagDict is similar to a dictionary except 

    Keys are unique tags/attributes of all the items
    Each key can be mapped onto multiple items that have
    that key as a tag

    TagDict[*tags] returns a list of items that share the same tags
    TagDict["*"] returns all the items
    '''
    def __init__(self):
        # Keys are tags. Values are sets of ids
        self.data = _collections.defaultdict(set)
        # Keys are ids of the objects. Values are (object,tags)
        self._ids = _IdDict()

    def add(self,item,tags):
        '''  Add an item with a list of tags
        if tags is empty, the item will not be added to
        the TagDict
        Input:
        item   - an object
        tags   - a string or a list of strings
        '''
        if type(tags) is str: tags = [tags,]
        tags = set(tags)
        self._ids[id(item)] = (item,tags)
        for tag in tags:
            self.data[tag].add(id(item))

    def __getitem__(self,tags):
        ''' Get the items that share the tags

        Return:
        A list of object
        If the list contains only one object, return the object

        Example:
        TagDict["a","b"] returns all items that have both "a" 
                         and "b" as tags
        TagDict["*"] returns all the items in the TagDict

        '''
        if tags[0] == '*':
            return [ value[0] for value in self._ids.values() ]
        if type(tags) is str: tags = [tags,]
        collected_ids = [ self.data[tag] 
                          if tag in self.data else set()
                          for tag in tags ]
        unique_ids = list(set.intersection(*collected_ids))
        if len(unique_ids) == 1:
            return self._ids[unique_ids[0]][0]
        else:
            return tuple([ self._ids[itemid][0] for itemid in unique_ids ])

    def __repr__(self):
        return "{"+', '.join(
            ['\''+tag +"\' : "+str(len(self.data[tag]))+" items" 
             for tag in self.data.keys()])+"}"

    def remove(self,item):
        ''' Remove an item from the TagDict.
        Any tag that points to no object will be removed too.
        Input:
        item   - an object inside the TagDict
        '''
        self.replace_tags(item,[])

    def add_tag(self,item,tag):
        ''' Add a tag to an item in the TagDict 
        Input:
        item   - an object inside the TagDict
        tag    - a string
        '''
        assert type(tag) is str
        self._ids[id(item)][1].add(tag)
        self.data[tag].add(id(item))

    def remove_tag(self,item,tag):
        ''' Remove a tag from an item in the TagDict
        Input:
        item   - an object inside the TagDict
        tag    - a string
        '''
        assert type(tag) is str
        self._ids[id(item)][1].remove(tag)
        self.data[tag].remove(id(item))
        if len(self.data[tag]) == 0:
            del self.data[tag]

    def replace_tags(self,item,newtags):
        ''' Replace the tags of an item.
        Any tag that points to no object will be removed too.
        Input:
        item     - an object inside the TagDict
        newtags  - a string or a list of strings
        '''
        if type(newtags) is str: newtags = {newtags}
        if type(newtags) is not set: newtags = set(newtags)
        oldtags = self._ids[id(item)][1]
        tags2rm = oldtags - newtags
        tags2add = newtags - oldtags
        for tag in tags2rm:
            self.remove_tag(item,tag)
        for tag in tags2add:
            self.add_tag(item,tag)
        if len(newtags) == 0:
            del self._ids[id(item)]

    def view_all(self):
        ''' Show all the items and their tags
        '''
        return tuple(self._ids.values())


if __name__ == '__main__':
    data = TagDict()
    # 4 unhashable items
    data.add({'Name':'Ben'},tags=['Male','Student'])
    data.add({'Name':'Tom','Age':40},tags=['Male','Teacher'])
    data.add({'Name':'Tina','Age':30},tags=['Female','Teacher'])
    data.add({'Name':'Ann'},tags=['Female','Student'])
    print data["*"]                     # --> All the items
    print data['Teacher','Female']      # --> {'Age': 30, 'Name': 'Tina'}
    print data['Student']               # --> ({'Name': 'Ben'}, {'Name': 'Ann'})
    # Add one more tag for one of the items
    data.add_tag(data['Teacher','Female'],'Mother')
    print data['Mother']                # --> {'Age': 30, 'Name': 'Tina'}
    # Remove a tag for one of the items
    data.remove_tag(data['Male','Student'],'Student')
    print data['Student']               # --> {'Name': 'Ann'}
    # Replace all the tags of one of the items
    data.replace_tags(data['Female','Mother','Teacher'],'Human')
    print data['Mother']                # --> ()
    print data['Human']                 # --> {'Age': 30, 'Name': 'Tina'}
    # Change the content of one of the items
    data['Human']['Age']=31
    print data['Human']                 # --> {'Age': 31, 'Name': 'Tina'}
    # Remove an item
    data.remove(data['Student'])
    print data['Student']               # ()
    print data['*']                     # Only 3 items now
    # Updating a list of items will fail
    try:
        data.add_tag(data['Male'],'Martian')
    except KeyError:
        print "More than one item were requested."
        for man in data['Male']:
            data.add_tag(man,'Martian')
    print data['Martian']              # --> ({'Name': 'Ben'}, {'Age': 40, 'Name': 'Tom'})

Update: Thanks to Sean Perry, I modified repr() by quite a bit as well:

def __repr__(self):
    pairs = [ "{tag} : {count} item{plural}".\
              format(tag=tag,
                     count=len(self.data[tag]),
                     plural='s' if len(self.data[tag]) > 1 else '')
              for tag in self.data.keys() ]
    return "<TagDict {"+", ".join(pairs)+"}>"

Answer 1

Sean Perry has made several good points, which I'll not duplicate. Though, I think your imports with leading underscores are fine! A leading underscore in a global name suggests that the item is not part of the module's public interface.

First off, There is usually no need to use UserDict in new code. As the docs for that module say:

The need for this class has been largely supplanted by the ability to subclass directly from dict (a feature that became available starting with Python version 2.2). Prior to the introduction of dict, the UserDict class was used to create dictionary-like sub-classes that obtained new behaviors by overriding existing methods or adding new ones.

So, your _IdDict class should probably inherit from dict directly, rather than from UserDict unless you need to support Python versions older than 2.2! This will also improve your forward compatibility, as the UserDict module has been removed in Python 3.

Or, you could probably do without the special dict subclass entirely, and handle the exception raising in the TagDict yourself. Just catch whatever exception gets raised by a normal dictionary, and raise your own (in Python 3, you'd want to use raise Whatever() from None to suppress the previous exception context, but in Python 2 that's neither possible nor necessary).

I see a few things that could be improved in your TagDict class itself.

You should probably add a check in add to make sure the item being added isn't in the dictionary already. If it is and the tags it's being added under are not the same as the ones it was under previously, you may end up with inconsistent information in your data and _ids dicts.

In __getitem__ you have the expression self.data[tag] if tag in self.data else set() in your list comprehension. You can write this more concisely as self.data.get(tag, set()).

But you might need to think about whether that is what you actually want to happen. If a requested tag is not found in the data dictionary, the intersection of the sets is going to be empty. This means you'll end up returning an empty tuple. Perhaps you should raise an exception instead?

Answer 2

The message to KeyError could be defined with a triple quote string to avoid the addition.

Instead of hiding the modules you should use explicit imports like

from collections import defaultdict

Instead of using type to check if something is a string you should use isinstance like if isinstance(foo, str).

I would use format in __repr__ to make the output easier to read.