Add SplitMediatTransactionData

1 files changed, 50 insertions, 0 deletions

diff --git a/sortedtransaction.py b/sortedtransaction.py
index 2adcf1b31..82edf04fc 100644
--- a/sortedtransaction.py
+++ b/sortedtransaction.py
@@ -3,10 +3,60 @@
 from yum.transactioninfo import TransactionData, TransactionMember
 from yum.Errors import YumBaseError
 
+import urlparse
+urlparse.uses_fragment.append('media')
+
 WHITE = 0
 GREY = 1
 BLACK = 2
 
+class SplitMediaTransactionData(TransactionData):
+    def __init__(self):
+        TransactionData.__init__(self)
+        self.reqmedia = {}
+        self.curmedia = 0 
+
+    def __getMedia(self, po):
+        try:
+            uri = po.returnSimple('basepath')
+            (scheme, netloc, path, query, fragid) = urlparse.urlsplit(uri)
+            if scheme != "media" or not fragid:
+                return 0
+            else:
+                return fragid
+        except KeyError:
+            return 0
+
+     def getMembers(self, pkgtup=None):
+        if not self.curmedia:
+            return TransactionData.getMembers(self, pkgtup)
+        if pkgtup is None:
+            returnlist = []
+            for key in self.reqmedia[self.curmedia].values():
+                returnlist.extend(self.pkgdict[key])
+
+            return returnlist
+
+        if self.reqmedia[self.curmedia].has_key(pkgtup):
+            return self.pkgdict[pkgtup]
+        else:
+            return []
+
+    def add(self, txmember):
+        id = self.__getMedia(txmember.po)
+        if id and id not in self.reqmedia.keys():
+            self.reqmedia[id].append(txmember.pkgtup)
+        TransactionData.add(self, txmember)
+
+    def remove(self, pkgtup):
+        txmember = self.pkgdict[pkgtup]
+        id = self.__getMedia(txmember.po)
+        if id:
+           self.reqmedia[id].remove(pktup)
+        del self.pkgdict[pkgtup]
+        del txmember
+
+
 class SortableTransactionData(TransactionData):
     def __init__(self):
         self._sorted = []