291 lines
9.3 KiB
Python
291 lines
9.3 KiB
Python
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()
|