pyunits/units/backend/units.py

from collections import defaultdict
from .dimensional_values import Dimension, Scaling, DimensionalValue
from .errors import UnitError

class BaseUnit(object):
	"""
	The atomic unit type that cannot be expressed by other units
	(well of course they can be expressed, but they should not)

	#FIXME:
	Should I compare symbols here?
	"""
	__slots__ = "symbol", "name", "description", "is_SI"
	def __init__(self, symbol, name, description, is_SI):
		self.symbol = symbol
		self.name = name
		self.description = description
		self.is_SI = is_SI

	def __mul__(self, other):
		if(isinstance(other, CompoundUnit)):
			return other.__rmul__(self)
		if(isinstance(other, OutputUnit)):
			return other.to_compound().__rmul__(self)
		if(not isinstance(other, BaseUnit)):
			raise TypeError("Cannot __mul__ BaseUnit and {}".format(type(other)))
		return CompoundUnit([(self, 1), (other, 1)])
	def __truediv__(self, other):
		if(isinstance(other, CompoundUnit)):
			return other.__rtruediv__(self)
		if(isinstance(other, OutputUnit)):
			return other.to_compound().__rtruediv__(self)
		if(not isinstance(other, BaseUnit)):
			raise TypeError("Cannot __mul__ BaseUnit and {}".format(type(other)))
		return CompoundUnit([(self, 1), (other, -1)])

	def __pow__(self, other):
		if(not isinstance(other, (int, float, complex))):
			raise TypeError("Cannot raise BaseUnit to non-scalar")
		return CompoundUnit([(self, other)])

	def __repr__(self):
		return self.symbol
	def __str__(self):
		return self.symbol

	def __matmul__(self, other):
		return other @ Dimension(Scaling(1, 1), self)
	def __rmatmul__(self, other):
		return other @ Dimension(Scaling(1, 1), self)

	def __rrshift__(self, other):
		if(isinstance(other, DimensionalValue)):
			return other >> Dimension(Scaling(1, 1), self)
		if(isinstance(other, BaseUnit)):
			if(other == self):
				return self
			raise UnitError("Cannot shift {} to {}".format(other, self))
		if(isinstance(other, CompoundUnit)):
			if(self in [base for base, exponent in other.bases]):
				return OutputUnit([other / self, self])
			raise UnitError("Cannot shift {} to {}".format(other, self))

		if(isinstance(other, OutputUnit)):
			try_here = other.components[0]
			rest = other.components[1:]
			return OutputUnit((try_here >> self).components.extend(rest))
		raise TypeError("Cannot rightshift BaseUnit and {}".format(type(other)))

	def __rshift__(self, other):
		if(isinstance(other, (CompoundUnit, BaseUnit))):
			return other.__rrshift__(self)
		if(isinstance(other, OutputUnit)):
			return other.to_compound().__rrshift__(self)
		raise TypeError("Cannot rightshift BaseUnit and {}".format(type(other)))

	def is_scalar(self):
		return False
		
	

		
		

class CompoundUnit(object):
	"""
	A unit that can (and should) be expressed using other units.
	This unit can be used as-is and can be broke down to BaseUnits.

	Also it should be possible to automatically gather BaseUnits and
	create CompoundUnits while multiplying.

	``bases`` is a list of tuples ``[(BaseUnit, exponent)]`` where
	exponent is an integer
	"""
	__slots__ = "is_SI", "bases"
	def __check_and_split_bases(bases):
		for base, exponent in bases:
			if(isinstance(exponent, float) and exponent.is_integer()):
				exponent = int(exponent)
			if(not isinstance(exponent, int)):
				raise ValueError("Unit Exponent must be integer")
			if(isinstance(base, BaseUnit)):
				yield base, exponent
				continue
			if(isinstance(base, CompoundUnit)):
				yield from CompoundUnit.__check_and_split_bases(base.bases)
				continue
			raise TypeError("Unit bases must be BaseUnit or CompoundUnit")

	def __cleanup_bases(bases):
		dct = defaultdict(int)
		for base, exponent in bases:
			dct[base] += exponent
		for k,v in dct.items():
			# do not yield powers of 0
			if(v):
				yield (k, v)
	def __init__(self, bases):
		self.bases = list(CompoundUnit.__cleanup_bases(CompoundUnit.__check_and_split_bases(bases)))
		self.is_SI = all([base.is_SI for base,_ in self.bases])

	def split(self):
		return self.bases

	def __mul__(self, other):
		if(isinstance(other, BaseUnit)):
			return CompoundUnit(self.bases + [(other, 1)])
		if(isinstance(other, CompoundUnit)):
			return CompoundUnit(self.bases + other.bases)
		raise TypeError("Cannot __mul__ CompoundUnit and {}".format(type(other)))
	def __rmul__(self, other):
		if(isinstance(other, BaseUnit)):
			return CompoundUnit(self.bases + [(other, 1)])
		if(isinstance(other, CompoundUnit)):
			return CompoundUnit(self.bases + other.bases)
		raise TypeError("Cannot __rmul__ CompoundUnit and {}".format(type(other)))
	def __pow__(self, other):
		if(not isinstance(other, (float, int, complex))):
			raise TypeError("Cannot raise CompoundUnit to non-scalar")
		return CompoundUnit([(base, other * exponent) for base, exponent in self.bases])
	
	def __truediv__(self, other):
		if(isinstance(other, BaseUnit)):
			return CompoundUnit(self.bases + [(other, -1)])
		if(isinstance(other, CompoundUnit)):
			return CompoundUnit(self.bases + [(base, -exponent) for base, exponent in other.bases])

		raise TypeError("Cannot __truediv__ CompoundUnit and {}".format(type(other)))

	def __rtruediv__(self, other):
		if(isinstance(other, BaseUnit)):
			return CompoundUnit([(other, 1)] + [(base, -exponent) for base, exponent in self.bases])
		if(isinstance(other, CompoundUnit)):
			return CompoundUnit(other.bases + [(base, -exponent) for base, exponent in self.bases])
		raise TypeError("Cannot __rtruediv__ CompoundUnit and {}".format(type(other)))

	def __repr__(self):
		return "*".join(["{}**{}".format(base, exponent) for base, exponent in self.bases])
		
	def __matmul__(self, other):
		return other @ Dimension(Scaling(1, 1), self)
	def __rmatmul__(self, other):
		return other @ Dimension(Scaling(1, 1), self)

	def __eq__(self, other):
		if(isinstance(other, BaseUnit)):
			return self.bases == [(other, 1)]
		if(isinstance(other, CompoundUnit)):
			return self.bases == other.bases
		raise TypeError("Cannot compare CompoundUnit and {}".format(type(other)))
	def ungroup(self):
		"""
		This is just used to implement the rightshift stuff.
		"""
		return CompoundUnit(self.bases)

	def __rrshift__(self, other):
		if(isinstance(other, DimensionalValue)):
			return other >> Dimension(Scaling(1, 1), self)

		if(isinstance(other, BaseUnit)):
			chk = [base for base, exponent in other.bases]
			if(other in chk):
				return self.ungroup()
			raise UnitError("Cannot shift {} to {}".format(other, self))
		if(isinstance(other, CompoundUnit)):
			chk = [base for base, exponent in other.bases]
			for base, exponent in self.bases:
				if(not base in chk):
					raise UnitError("Cannot shift {} to {}".format(other, self))

				return OutputUnit([(other.ungroup() / self.ungroup()), self])
		if(isinstance(other, OutputUnit)):
			try_here = other.components[0]
			rest = other.components[1:]
			return OutputUnit((try_here >> self).components.extend(rest))

		raise TypeError("Cannot rightshift BaseUnit and {}".format(type(other)))
				
	def __rshift__(self, other):
		if(isinstance(other, (CompoundUnit, BaseUnit))):
			return other.__rrshift__(self)
		if(isinstance(other, OutputUnit)):
			return other.to_compound().__rrshift__(self)
		raise TypeError("Cannot rightshift CompoundUnit and {}".format(type(other)))

	def is_scalar(self):
		return len(self.bases) == 0



class NamedCompoundUnit(CompoundUnit):	
	"""
	A unit that can (and should) be expressed using other units.
	This unit can be used as-is and can be broke down to BaseUnits.

	Also it should be possible to automatically gather BaseUnits and
	create CompoundUnits while multiplying.

	``bases`` is a list of tuples ``[(BaseUnit, exponent)]`` where
	exponent is an integer
	"""
	__slots__ = "symbol", "name", "description", "is_SI", "bases"
	def __init__(self, symbol, name, description, bases):
		CompoundUnit.__init__(self, bases)
		self.symbol = symbol
		self.name = name
		self.description = description

	def __repr__(self):
		return self.symbol
	def __str__(self):
		return self.symbol



class OutputUnit(object):
	"""
	Unit type for pretty output.

	Will be created by __rrshift__ and others.
	
	It has several compound units that are listed explicitly
	for prettier output, if any operation is performed it will automatically
	cast to a CompoundUnit.
	"""
	__slots__ = "components",
	def __init__(self, components):
		self.components = components

	def to_compound(self):
		bases = []
		for compound in self.components:
			if(isinstance(compound, BaseUnit)):
				bases.append((compound, 1))
			if(isinstance(compound, CompoundUnit)):
				bases.extend(compound.bases)

		return CompoundUnit(bases)

	def __repr__(self):
		return "*".join([repr(c) for c in self.components if c])

	def __mul__(self, other):
		return self.to_compound().__mul__(other)
	def __rmul__(self, other):
		return self.to_compound().__rmul__(other)
	def __truediv__(self, other):
		return self.to_compound().__truediv__(other)
	def __rtruediv__(self, other):
		return self.to_compound().__rtruediv__(other)
	def __pow__(self, other):
		return self.to_compound().__pow__(other)
	def __matmul__(self, other):
		return self.to_compound().__matmul__(other)
	def __rmatmul__(self, other):
		return self.to_compound().__rmatmul__(other)
	def __eq__(self, other):
		return self.to_compound().__eq__(other)

	def split(self):
		return self.to_compound().split()
	def ungroup(self):
		return self.to_compound().ungroup()
	def __rrshift__(self, other):
		return self.to_compound().__rrshift__(other)

	def is_scalar(self):
		return self.to_compound().is_scalar()