Fix(core): Add doctest and simplify tree build

This commit is contained in:
Bertrand Benjamin 2018-12-07 11:21:40 +01:00
parent 24b421209f
commit 21cc3b68a3
1 changed files with 17 additions and 16 deletions

View File

@ -103,6 +103,15 @@ class MOMonomial(MO):
>>> x = MOstrPower('x', 2)
>>> MOMonomial(4, x)
<MOMonomial 4x^2>
>>> MOMonomial(4, 'x')
<MOMonomial 4x>
>>> MOMonomial(4, 'x', 1)
<MOMonomial 4x>
>>> MOMonomial(4, 'x', 2)
<MOMonomial 4x^2>
>>> x = MOstrPower('x', 2)
>>> MOMonomial(4, x, 3)
<MOMonomial 4x^6>
>>> MOMonomial(0, x)
Traceback (most recent call last):
...
@ -118,30 +127,22 @@ class MOMonomial(MO):
_power = MO.factory(_variable.power.value * power)
_variable = _variable.variable
elif isinstance(_variable, MOstr):
_variable = variable
_power = MO.factory(power)
_variable = variable
else:
raise MOError(f"variable need to be a MOstrPower or a MOstr. Got {type(variable)}.")
self._variable = _variable
self._power = _power
if self._power == 1:
if self._coefficient == 1:
_tree = self._variable
try:
if self._coefficient.value != 1:
_tree = Tree("*", self._coefficient, self.strpower)
else:
_tree = Tree("*",
self._coefficient,
self._variable
)
else:
if self._coefficient == 1:
_tree = MOstrPower(self._variable, self._power)
else:
_tree = Tree("*",
self._coefficient,
MOstrPower(self._variable, self._power)
)
_tree = self.strpower
except AttributeError:
_tree = Tree("*", self._coefficient, self.strpower)
MO.__init__(self, _tree)