Add new Sectors unit, enhance Bytes unit, add unit tests for both

bootstrapvz/common/bytes.py

@@ -1,3 +1,14 @@
+from exceptions import UnitError
+
+
+def onlybytes(msg):
+	def decorator(func):
+		def check_other(self, other):
+			if not isinstance(other, Bytes):
+				raise UnitError(msg)
+			return func(self, other)
+		return check_other
+	return decorator
 
 
 class Bytes(object):
@@ -61,25 +72,48 @@ class Bytes(object):
 	def __long__(self):
 		return self.qty
 
+	def __abs__(self):
+		return self.qty
+
+	@onlybytes('Can only compare Bytes to Bytes')
+	def __lt__(self, other):
+		return self.qty < other.qty
+
+	@onlybytes('Can only compare Bytes to Bytes')
+	def __le__(self, other):
+		return self.qty <= other.qty
+
+	@onlybytes('Can only compare Bytes to Bytes')
+	def __eq__(self, other):
+		return self.qty == other.qty
+
+	@onlybytes('Can only compare Bytes to Bytes')
+	def __ne__(self, other):
+		return self.qty != other.qty
+
+	@onlybytes('Can only compare Bytes to Bytes')
+	def __ge__(self, other):
+		return self.qty >= other.qty
+
+	@onlybytes('Can only compare Bytes to Bytes')
+	def __gt__(self, other):
+		return self.qty > other.qty
+
+	@onlybytes('Can only add Bytes to Bytes')
 	def __add__(self, other):
-		if not isinstance(other, Bytes):
-			raise UnitError('Can only add Bytes to Bytes')
 		return Bytes(self.qty + other.qty)
 
+	@onlybytes('Can only add Bytes to Bytes')
 	def __iadd__(self, other):
-		if not isinstance(other, Bytes):
-			raise UnitError('Can only add Bytes to Bytes')
 		self.qty += other.qty
 		return self
 
+	@onlybytes('Can only subtract Bytes from Bytes')
 	def __sub__(self, other):
-		if not isinstance(other, Bytes):
-			raise UnitError('Can only subtract Bytes from Bytes')
 		return Bytes(self.qty - other.qty)
 
+	@onlybytes('Can only subtract Bytes from Bytes')
 	def __isub__(self, other):
-		if not isinstance(other, Bytes):
-			raise UnitError('Can only subtract Bytes from Bytes')
 		self.qty -= other.qty
 		return self
 
@@ -110,20 +144,13 @@ class Bytes(object):
 			self.qty /= other
 		return self
 
+	@onlybytes('Can only take modulus of Bytes with Bytes')
 	def __mod__(self, other):
-		if isinstance(other, Bytes):
-			return self.qty % other.qty
-		if not isinstance(other, (int, long)):
-			raise UnitError('Can only take modulus of Bytes with integers or Bytes')
-		return Bytes(self.qty % other)
+		return Bytes(self.qty % other.qty)
 
+	@onlybytes('Can only take modulus of Bytes with Bytes')
 	def __imod__(self, other):
-		if isinstance(other, Bytes):
-			self.qty %= other.qty
-		else:
-			if not isinstance(other, (int, long)):
-				raise UnitError('Can only divide Bytes with integers or Bytes')
-			self.qty %= other
+		self.qty %= other.qty
 		return self
 
 	def __getstate__(self):
@@ -133,7 +160,3 @@ class Bytes(object):
 
 	def __setstate__(self, state):
 		self.qty = state['qty']
-
-
-class UnitError(Exception):
-	pass

bootstrapvz/common/sectors.py

@@ -0,0 +1,154 @@
+from exceptions import UnitError
+from bytes import Bytes
+
+
+def onlysectors(msg):
+	def decorator(func):
+		def check_other(self, other):
+			if not isinstance(other, Sectors):
+				raise UnitError(msg)
+			return func(self, other)
+		return check_other
+	return decorator
+
+
+class Sectors(object):
+
+	def __init__(self, quantity, sector_size):
+		if isinstance(sector_size, Bytes):
+			self.sector_size = sector_size
+		else:
+			self.sector_size = Bytes(sector_size)
+
+		if isinstance(quantity, Bytes):
+			self.bytes = quantity
+		else:
+			if isinstance(quantity, (int, long)):
+				self.bytes = self.sector_size * quantity
+			else:
+				self.bytes = Bytes(quantity)
+
+	def get_sectors(self):
+		return self.bytes / self.sector_size
+
+	def __repr__(self):
+		return str(self.get_sectors()) + 's'
+
+	def __str__(self):
+		return self.__repr__()
+
+	def __int__(self):
+		return self.get_sectors()
+
+	def __long__(self):
+		return self.get_sectors()
+
+	def __abs__(self):
+		return self.get_sectors()
+
+	@onlysectors('Can only compare sectors with sectors')
+	def __lt__(self, other):
+		return self.bytes < other.bytes
+
+	@onlysectors('Can only compare sectors with sectors')
+	def __le__(self, other):
+		return self.bytes <= other.bytes
+
+	@onlysectors('Can only compare sectors with sectors')
+	def __eq__(self, other):
+		return self.bytes == other.bytes
+
+	@onlysectors('Can only compare sectors with sectors')
+	def __ne__(self, other):
+		return self.bytes != other.bytes
+
+	@onlysectors('Can only compare sectors with sectors')
+	def __ge__(self, other):
+		return self.bytes >= other.bytes
+
+	@onlysectors('Can only compare sectors with sectors')
+	def __gt__(self, other):
+		return self.bytes > other.bytes
+
+	def __add__(self, other):
+		if not isinstance(other, Sectors):
+			raise UnitError('Can only add sectors to sectors')
+		if self.sector_size != other.sector_size:
+			raise UnitError('Cannot sum sectors with different sector sizes')
+		return Sectors(self.bytes + other.bytes, self.sector_size)
+
+	def __iadd__(self, other):
+		if not isinstance(other, (Bytes, Sectors)):
+			raise UnitError('Can only add Bytes or sectors to sectors')
+		if isinstance(other, Bytes):
+			self.bytes += other
+		if isinstance(other, Sectors):
+			if self.sector_size != other.sector_size:
+				raise UnitError('Cannot sum sectors with different sector sizes')
+			self.bytes += other.bytes
+		return self
+
+	def __sub__(self, other):
+		if not isinstance(other, (Sectors, int, long)):
+			raise UnitError('Can only subtract sectors or integers from sectors')
+		if isinstance(other, int):
+			return Sectors(self.bytes - self.sector_size * other, self.sector_size)
+		else:
+			if self.sector_size != other.sector_size:
+				raise UnitError('Cannot subtract sectors with different sector sizes')
+			return Sectors(self.bytes - other.bytes, self.sector_size)
+
+	def __isub__(self, other):
+		if not isinstance(other, (Sectors, int, long)):
+			raise UnitError('Can only subtract sectors or integers from sectors')
+		if isinstance(other, int):
+			self.bytes -= self.sector_size * other
+		else:
+			if self.sector_size != other.sector_size:
+				raise UnitError('Cannot subtract sectors with different sector sizes')
+			self.bytes -= other.bytes
+		return self
+
+	def __mul__(self, other):
+		if not isinstance(other, (int, long)):
+			raise UnitError('Can only multiply sectors with integers')
+		return Sectors(self.bytes * other, self.sector_size)
+
+	def __imul__(self, other):
+		if not isinstance(other, (int, long)):
+			raise UnitError('Can only multiply sectors with integers')
+		self.bytes *= other
+		return self
+
+	def __div__(self, other):
+		if isinstance(other, Sectors):
+			if self.sector_size != other.sector_size:
+				raise UnitError('Cannot divide sectors with different sector sizes')
+			return self.bytes / other.bytes
+		if not isinstance(other, (int, long)):
+			raise UnitError('Can only divide sectors with integers or sectors')
+		return Sectors(self.bytes / other, self.sector_size)
+
+	def __idiv__(self, other):
+		if isinstance(other, Sectors):
+			if self.sector_size != other.sector_size:
+				raise UnitError('Cannot divide sectors with different sector sizes')
+			self.bytes /= other.bytes
+		else:
+			if not isinstance(other, (int, long)):
+				raise UnitError('Can only divide sectors with integers or sectors')
+			self.bytes /= other
+		return self
+
+	@onlysectors('Can only take modulus of sectors with sectors')
+	def __mod__(self, other):
+		if self.sector_size != other.sector_size:
+			raise UnitError('Cannot take modulus of sectors with different sector sizes')
+		return Sectors(self.bytes % other.bytes, self.sector_size)
+
+	@onlysectors('Can only take modulus of sectors with sectors')
+	def __imod__(self, other):
+		if self.sector_size != other.sector_size:
+			raise UnitError('Cannot take modulus of sectors with different sector sizes')
+		self.bytes %= other.bytes
+		return self

tests/unit/bytes.py

@@ -0,0 +1,108 @@
+from nose.tools import eq_
+from nose.tools import raises
+from bootstrapvz.common.bytes import Bytes
+from bootstrapvz.common.exceptions import UnitError
+
+
+def test_lt():
+	assert Bytes('1MiB') < Bytes('2MiB')
+
+
+def test_le():
+	assert Bytes('1MiB') <= Bytes('2MiB')
+	assert Bytes('1MiB') <= Bytes('1MiB')
+
+
+def test_eq():
+	eq_(Bytes('1MiB'), Bytes('1MiB'))
+
+
+def test_neq():
+	assert Bytes('15MiB') != Bytes('1MiB')
+
+
+def test_gt():
+	assert Bytes('2MiB') > Bytes('1MiB')
+
+
+def test_ge():
+	assert Bytes('2MiB') >= Bytes('1MiB')
+	assert Bytes('2MiB') >= Bytes('2MiB')
+
+
+def test_eq_unit():
+	eq_(Bytes('1024MiB'), Bytes('1GiB'))
+
+
+def test_add():
+	eq_(Bytes('2GiB'), Bytes('1GiB') + Bytes('1GiB'))
+
+
+def test_iadd():
+	b = Bytes('1GiB')
+	b += Bytes('1GiB')
+	eq_(Bytes('2GiB'), b)
+
+
+def test_sub():
+	eq_(Bytes('1GiB'), Bytes('2GiB') - Bytes('1GiB'))
+
+
+def test_isub():
+	b = Bytes('2GiB')
+	b -= Bytes('1GiB')
+	eq_(Bytes('1GiB'), b)
+
+
+def test_mul():
+	eq_(Bytes('2GiB'), Bytes('1GiB') * 2)
+
+
+@raises(UnitError)
+def test_mul_bytes():
+	Bytes('1GiB') * Bytes('1GiB')
+
+
+def test_imul():
+	b = Bytes('1GiB')
+	b *= 2
+	eq_(Bytes('2GiB'), b)
+
+
+def test_div():
+	eq_(Bytes('1GiB'), Bytes('2GiB') / 2)
+
+
+def test_div_bytes():
+	eq_(2, Bytes('2GiB') / Bytes('1GiB'))
+
+
+def test_idiv():
+	b = Bytes('2GiB')
+	b /= 2
+	eq_(Bytes('1GiB'), b)
+
+
+def test_mod():
+	eq_(Bytes('256MiB'), Bytes('1GiB') % Bytes('768MiB'))
+
+
+@raises(UnitError)
+def test_mod_int():
+	Bytes('1GiB') % 768
+
+
+def test_imod():
+	b = Bytes('1GiB')
+	b %= Bytes('768MiB')
+	eq_(Bytes('256MiB'), b)
+
+
+@raises(UnitError)
+def test_imod_int():
+	b = Bytes('1GiB')
+	b %= 5
+
+
+def test_abs():
+	eq_(pow(1024, 3), abs(Bytes('1GiB')))

tests/unit/sectors.py

@@ -0,0 +1,128 @@
+from nose.tools import eq_
+from nose.tools import raises
+from bootstrapvz.common.sectors import Sectors
+from bootstrapvz.common.bytes import Bytes
+from bootstrapvz.common.exceptions import UnitError
+
+std_secsz = Bytes(512)
+
+
+def test_init_with_int():
+	eq_(4, abs(Sectors(4, std_secsz)))
+	secsize = 4096
+	eq_(Sectors('1MiB', secsize), Sectors(256, secsize))
+
+
+def test_lt():
+	assert Sectors('1MiB', std_secsz) < Sectors('2MiB', std_secsz)
+
+
+def test_le():
+	assert Sectors('1MiB', std_secsz) <= Sectors('2MiB', std_secsz)
+	assert Sectors('1MiB', std_secsz) <= Sectors('1MiB', std_secsz)
+
+
+def test_eq():
+	eq_(Sectors('1MiB', std_secsz), Sectors('1MiB', std_secsz))
+
+
+def test_neq():
+	assert Sectors('15MiB', std_secsz) != Sectors('1MiB', std_secsz)
+
+
+def test_gt():
+	assert Sectors('2MiB', std_secsz) > Sectors('1MiB', std_secsz)
+
+
+def test_ge():
+	assert Sectors('2MiB', std_secsz) >= Sectors('1MiB', std_secsz)
+	assert Sectors('2MiB', std_secsz) >= Sectors('2MiB', std_secsz)
+
+
+def test_eq_unit():
+	eq_(Sectors('1024MiB', std_secsz), Sectors('1GiB', std_secsz))
+
+
+def test_add():
+	eq_(Sectors('2GiB', std_secsz), Sectors('1GiB', std_secsz) + Sectors('1GiB', std_secsz))
+
+
+@raises(UnitError)
+def test_add_with_diff_secsize():
+	Sectors('1GiB', Bytes(512)) + Sectors('1GiB', Bytes(4096))
+
+
+def test_iadd():
+	s = Sectors('1GiB', std_secsz)
+	s += Sectors('1GiB', std_secsz)
+	eq_(Sectors('2GiB', std_secsz), s)
+
+
+def test_sub():
+	eq_(Sectors('1GiB', std_secsz), Sectors('2GiB', std_secsz) - Sectors('1GiB', std_secsz))
+
+
+def test_sub_int():
+	secsize = Bytes('4KiB')
+	eq_(Sectors('1MiB', secsize), Sectors('1028KiB', secsize) - 1)
+
+
+def test_isub():
+	s = Sectors('2GiB', std_secsz)
+	s -= Sectors('1GiB', std_secsz)
+	eq_(Sectors('1GiB', std_secsz), s)
+
+
+def test_mul():
+	eq_(Sectors('2GiB', std_secsz), Sectors('1GiB', std_secsz) * 2)
+
+
+@raises(UnitError)
+def test_mul_bytes():
+	Sectors('1GiB', std_secsz) * Sectors('1GiB', std_secsz)
+
+
+def test_imul():
+	s = Sectors('1GiB', std_secsz)
+	s *= 2
+	eq_(Sectors('2GiB', std_secsz), s)
+
+
+def test_div():
+	eq_(Sectors('1GiB', std_secsz), Sectors('2GiB', std_secsz) / 2)
+
+
+def test_div_bytes():
+	eq_(2, Sectors('2GiB', std_secsz) / Sectors('1GiB', std_secsz))
+
+
+def test_idiv():
+	s = Sectors('2GiB', std_secsz)
+	s /= 2
+	eq_(Sectors('1GiB', std_secsz), s)
+
+
+def test_mod():
+	eq_(Sectors('256MiB', std_secsz), Sectors('1GiB', std_secsz) % Sectors('768MiB', std_secsz))
+
+
+@raises(UnitError)
+def test_mod_int():
+	Sectors('1GiB', std_secsz) % 768
+
+
+def test_imod():
+	s = Sectors('1GiB', std_secsz)
+	s %= Sectors('768MiB', std_secsz)
+	eq_(Sectors('256MiB', std_secsz), s)
+
+
+@raises(UnitError)
+def test_imod_int():
+	s = Sectors('1GiB', std_secsz)
+	s %= 5
+
+
+def test_abs():
+	secsize = 512
+	eq_(pow(1024, 3) / secsize, abs(Sectors('1GiB', secsize)))