--- usbspindownd/diskstats.py	2008/01/29 19:48:35	29
+++ usbspindownd/diskstats.py	2008/01/31 02:21:39	30
@@ -7,20 +7,26 @@
 # Copyright: (c) 2008 Andrew Pollock
 #
 
-class DiskStat:
+from UserDict import UserDict
+
+class Error(Exception):
+	pass
+
+class DiskStat(UserDict):
 
 	def __init__(self, nrread, nrmergedread, nrsectorsread, msread, nrwrite, nrmergedwrite, nrsectorswrite, mswrite, nrio, msio, weightedmsio):
-		self.nrread = nrread
-		self.nrmergedread = nrmergedread
-		self.nrsectorsread = nrsectorsread
-		self.msread = msread
-		self.nrwrite = nrwrite
-		self.nrmergedwrite = nrmergedwrite
-		self.nrsectorswrite = nrsectorswrite
-		self.mswrite = mswrite
-		self.nrio = nrio
-		self.msio = msio
-		self.weightedmsio = weightedmsio
+		UserDict.__init__(self)
+		self["nrread"] = int(nrread)
+		self["nrmergedread"] = int(nrmergedread)
+		self["nrsectorsread"] = int(nrsectorsread)
+		self["msread"] = int(msread)
+		self["nrwrite"] = int(nrwrite)
+		self["nrmergedwrite"] = int(nrmergedwrite)
+		self["nrsectorswrite"] = int(nrsectorswrite)
+		self["mswrite"] = int(mswrite)
+		self["nrio"] = int(nrio)
+		self["msio"] = int(msio)
+		self["weightedmsio"] = int(weightedmsio)
 
 
 class DiskStats:
@@ -28,7 +34,7 @@
 	def __init__(self, disk):
 		self.disk = disk.lstrip("/dev/")
 
-	def diskstat(self):
+	def diskstat(self, key=None):
 		proc_diskstats = open("/proc/diskstats")
 		stat = None
 		while True:
@@ -43,11 +49,14 @@
 		proc_diskstats.close()
 		if stat:
 			ds = DiskStat(stat[0], stat[1], stat[2], stat[3], stat[4], stat[5], stat[6], stat[7], stat[8], stat[9], stat[10])
-			return ds
+			if not key:
+				return ds
+			else:
+				return ds[key]
 		else:
 			# Need to raise an exception here
-			pass
+			raise Error("No such disk: %s" % (self.disk))
 
 if __name__ == "__main__":
-	ds = DiskStats("/dev/hda")
-	print ds.diskstat()
+	ds = DiskStats("/dev/sda")
+	print ds.diskstat("msio")