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