Procházet zdrojové kódy

Added possibility to specify an exponent to ExpressionBase.is_power.

Taddeus Kroes před 14 roky
rodič
revize
a9209ec2d2
2 změnil soubory, kde provedl 16 přidání a 4 odebrání
  1. 10 2
      src/node.py
  2. 6 2
      tests/test_node.py

+ 10 - 2
src/node.py

@@ -33,6 +33,11 @@ OP_EXPAND = 9
 OP_COMMA = 10
 OP_SQRT = 11
 
+# Goniometry
+OP_SIN = 12
+OP_COS = 13
+OP_TAN = 14
+
 
 TYPE_MAP = {
         int: TYPE_INTEGER,
@@ -112,8 +117,11 @@ class ExpressionBase(object):
     def is_op(self, op):
         return not self.is_leaf and self.op == op
 
-    def is_power(self):
-        return not self.is_leaf and self.op == OP_POW
+    def is_power(self, exponent=None):
+        if self.is_leaf or self.op != OP_POW:
+            return False
+
+        return exponent == None or self[1] == exponent
 
     def is_nary(self):
         return not self.is_leaf and self.op in [OP_ADD, OP_SUB, OP_MUL]

+ 6 - 2
tests/test_node.py

@@ -35,8 +35,12 @@ class TestNode(RulesTestCase):
         self.assertFalse(N('+', *self.l[:2]).is_leaf)
 
     def test_is_power(self):
-        self.assertTrue(N('^', *self.l[:2]).is_power())
-        self.assertFalse(N('+', *self.l[:2]).is_power())
+        self.assertTrue(N('^', *self.l[2:]).is_power())
+        self.assertFalse(N('+', *self.l[2:]).is_power())
+
+    def test_is_power_exponent(self):
+        self.assertTrue(N('^', *self.l[2:]).is_power(5))
+        self.assertFalse(N('^', *self.l[2:]).is_power(2))
 
     def test_is_nary(self):
         self.assertTrue(N('+', *self.l[:2]).is_nary())