1#!/usr/bin/env python
2from __future__ import division, print_function
3"""A dictionary that stores a list of values for each key.
4
5Note: one could subclass dict but this requires writing
6explicit methods for setdefault, copy, and possibly others.
7Only do this if extra performance is required.
8
9History
101-1 Russell Owen 8/8/00
111-2 corrected an indentation error
122003-05-06 ROwen    Test copy and setdefault; renamed to ListDict
13                    and added SetDict.
142010-05-18 ROwen    Modified SetDict to use sets
152015-09-24 ROwen    Replace "== None" with "is None" to future-proof array tests and modernize the code.
16"""
17__all__ = ["ListDict", "SetDict"]
18
19import UserDict
20
21class ListDict(UserDict.UserDict):
22    """A dictionary whose values are a list of items.
23    """
24    def __setitem__(self, key, val):
25        """Add a value to the list of values for a given key, creating a new entry if necessary.
26
27        Supports the notation: aListDict[key] = val
28        """
29        if key in self.data:
30            self.data[key].append(val)
31        else:
32            self.data[key] = [val]
33
34    def addList(self, key, valList):
35        """Append values to the list of values for a given key, creating a new entry if necessary.
36
37        Inputs:
38        - valList: an iterable collection (preferably ordered) of values
39        """
40        valList = list(valList)
41        if key in self.data:
42            self.data[key] += valList
43        else:
44            self.data[key] = valList
45
46    def remove(self, key, val):
47        """removes the specified value from the list of values for the specified key;
48
49        raise KeyError if key not found
50        raise ValueError if val not found
51        """
52        self.data.get(key).remove(val)
53
54class SetDict(ListDict):
55    """A dictionary whose values are a set of items, meaning
56    a list of unique items. Duplicate items are silently not added.
57    """
58
59    def __setitem__(self, key, val):
60        """Add a value to the set of values for a given key, creating a new entry if necessary.
61
62        Duplicate values are silently ignored.
63
64        Supports the notation: aListDict[key] = val
65        """
66        valSet = self.data.get(key)
67        if valSet is None:
68            self.data[key] = set([val])
69        else:
70            valSet.add(val)
71
72    def addList(self, key, valList):
73        """Add values to the set of values for a given key, creating a new entry if necessary.
74
75        Duplicate values are silently ignored.
76
77        Inputs:
78        - valList: an iterable collection of values
79        """
80        valSet = self.data.get(key)
81        if valSet is None:
82            self.data[key] = set(valList)
83        else:
84            valSet.update(valList)
85
86
87if __name__ == "__main__":
88    import RO.StringUtil
89
90    ad = ListDict()
91    ad["a"] = "foo a"
92    ad["a"] = "foo a"
93    ad["a"] = "bar a"
94    ad[1] = "foo 1"
95    ad[1] = "foo 2"
96    ad[1] = "foo 2"
97    ad2 = ad.copy()
98    ad2.setdefault("a", "foo")
99    ad2.setdefault("a", "foo")
100    ad2.setdefault("b", "bar")
101    ad2.setdefault("b", "bar")
102    print("listdict:")
103    print((RO.StringUtil.prettyDict(ad)))
104    print("listdict copy (modified):")
105    print((RO.StringUtil.prettyDict(ad2)))
106
107
108    ad = SetDict()
109    ad["a"] = "foo a"
110    ad["a"] = "foo a"
111    ad["a"] = "bar a"
112    ad[1] = "foo 1"
113    ad[1] = "foo 2"
114    ad[1] = "foo 2"
115    ad2 = ad.copy()
116    ad2.setdefault("a", "foo")
117    ad2.setdefault("a", "foo")
118    ad2.setdefault("b", "bar")
119    ad2.setdefault("b", "bar")
120    print("setdict:")
121    print((RO.StringUtil.prettyDict(ad)))
122    print("setdict copy (modified):")
123    print((RO.StringUtil.prettyDict(ad2)))
124