Allow using integers for more operations with Sectors

This commit is contained in:
Anders Ingemann 2015-01-19 23:11:00 +01:00
parent d105d10c76
commit 05006f2d83

View file

@ -68,87 +68,104 @@ class Sectors(object):
return self.bytes > other.bytes return self.bytes > other.bytes
def __add__(self, other): def __add__(self, other):
if not isinstance(other, Sectors): if isinstance(other, (int, long)):
raise UnitError('Can only add sectors to sectors') return Sectors(self.bytes + self.sector_size * other, self.sector_size)
if isinstance(other, Bytes):
return Sectors(self.bytes + other, self.sector_size)
if isinstance(other, Sectors):
if self.sector_size != other.sector_size: if self.sector_size != other.sector_size:
raise UnitError('Cannot sum sectors with different sector sizes') raise UnitError('Cannot sum sectors with different sector sizes')
return Sectors(self.bytes + other.bytes, self.sector_size) return Sectors(self.bytes + other.bytes, self.sector_size)
raise UnitError('Can only add sectors, bytes or integers to sectors')
def __iadd__(self, other): def __iadd__(self, other):
if not isinstance(other, (Bytes, Sectors)): if isinstance(other, (int, long)):
raise UnitError('Can only add Bytes or sectors to sectors') self.bytes += self.sector_size * other
return self
if isinstance(other, Bytes): if isinstance(other, Bytes):
self.bytes += other self.bytes += other
return self
if isinstance(other, Sectors): if isinstance(other, Sectors):
if self.sector_size != other.sector_size: if self.sector_size != other.sector_size:
raise UnitError('Cannot sum sectors with different sector sizes') raise UnitError('Cannot sum sectors with different sector sizes')
self.bytes += other.bytes self.bytes += other.bytes
return self return self
raise UnitError('Can only add sectors, bytes or integers to sectors')
def __sub__(self, other): def __sub__(self, other):
if not isinstance(other, (Sectors, int, long)): if isinstance(other, (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) return Sectors(self.bytes - self.sector_size * other, self.sector_size)
else: if isinstance(other, Bytes):
return Sectors(self.bytes - other, self.sector_size)
if isinstance(other, Sectors):
if self.sector_size != other.sector_size: if self.sector_size != other.sector_size:
raise UnitError('Cannot subtract sectors with different sector sizes') raise UnitError('Cannot subtract sectors with different sector sizes')
return Sectors(self.bytes - other.bytes, self.sector_size) return Sectors(self.bytes - other.bytes, self.sector_size)
raise UnitError('Can only subtract sectors, bytes or integers from sectors')
def __isub__(self, other): def __isub__(self, other):
if not isinstance(other, (Sectors, int, long)): if isinstance(other, (int, long)):
raise UnitError('Can only subtract sectors or integers from sectors')
if isinstance(other, int):
self.bytes -= self.sector_size * other self.bytes -= self.sector_size * other
else: return self
if isinstance(other, Bytes):
self.bytes -= other
return self
if isinstance(other, Sectors):
if self.sector_size != other.sector_size: if self.sector_size != other.sector_size:
raise UnitError('Cannot subtract sectors with different sector sizes') raise UnitError('Cannot subtract sectors with different sector sizes')
self.bytes -= other.bytes self.bytes -= other.bytes
return self return self
raise UnitError('Can only subtract sectors, bytes or integers from sectors')
def __mul__(self, other): def __mul__(self, other):
if not isinstance(other, (int, long)): if isinstance(other, (int, long)):
raise UnitError('Can only multiply sectors with integers')
return Sectors(self.bytes * other, self.sector_size) return Sectors(self.bytes * other, self.sector_size)
else:
raise UnitError('Can only multiply sectors with integers')
def __imul__(self, other): def __imul__(self, other):
if not isinstance(other, (int, long)): if isinstance(other, (int, long)):
raise UnitError('Can only multiply sectors with integers')
self.bytes *= other self.bytes *= other
return self return self
else:
raise UnitError('Can only multiply sectors with integers')
def __div__(self, other): def __div__(self, other):
if isinstance(other, (int, long)):
return Sectors(self.bytes / other, self.sector_size)
if isinstance(other, Sectors): if isinstance(other, Sectors):
if self.sector_size != other.sector_size: if self.sector_size != other.sector_size:
raise UnitError('Cannot divide sectors with different sector sizes')
return self.bytes / other.bytes return self.bytes / other.bytes
if not isinstance(other, (int, long)): else:
raise UnitError('Cannot divide sectors with different sector sizes')
raise UnitError('Can only divide sectors with integers or sectors') raise UnitError('Can only divide sectors with integers or sectors')
return Sectors(self.bytes / other, self.sector_size)
def __idiv__(self, other): def __idiv__(self, other):
if isinstance(other, Sectors): if isinstance(other, (int, long)):
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 self.bytes /= other
return self return self
if isinstance(other, Sectors):
if self.sector_size == other.sector_size:
self.bytes /= other.bytes
return self
else:
raise UnitError('Cannot divide sectors with different sector sizes')
raise UnitError('Can only divide sectors with integers or sectors')
@onlysectors('Can only take modulus of sectors with sectors') @onlysectors('Can only take modulus of sectors with sectors')
def __mod__(self, other): def __mod__(self, other):
if self.sector_size != other.sector_size: 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) return Sectors(self.bytes % other.bytes, self.sector_size)
else:
raise UnitError('Cannot take modulus of sectors with different sector sizes')
@onlysectors('Can only take modulus of sectors with sectors') @onlysectors('Can only take modulus of sectors with sectors')
def __imod__(self, other): def __imod__(self, other):
if self.sector_size != other.sector_size: if self.sector_size == other.sector_size:
raise UnitError('Cannot take modulus of sectors with different sector sizes')
self.bytes %= other.bytes self.bytes %= other.bytes
return self return self
else:
raise UnitError('Cannot take modulus of sectors with different sector sizes')
def __getstate__(self): def __getstate__(self):
return {'__class__': self.__module__ + '.' + self.__class__.__name__, return {'__class__': self.__module__ + '.' + self.__class__.__name__,