Added proper functionality for unary minus and plus

Author: charon25

setup.py

@@ -5,7 +5,7 @@
 
 setuptools.setup(
 	name="shunting-yard",
-	version="1.0.2",
+	version="1.0.3",
 	author="Paul 'charon25' Kern",
 	description="Compute any math expression",
 	long_description=long_description,

shunting_yard/constants.py

@@ -1,6 +1,7 @@
 from string import ascii_lowercase
 
 
+UNARY_OPERATORS = '+-'
 BASE_OPERATORS = '+-*/^'
 NUMBER_CHARS = '0123456789.'
 FUNCTION_CHARS = ascii_lowercase + '_'

shunting_yard/rpn.py

@@ -1,5 +1,5 @@
 import math
-from operator import add, mul, sub, truediv
+from operator import add, mul, neg, pos, sub, truediv
 from typing import Any, Callable, Optional, Union
 
 from shunting_yard.constants import NUMBER_CHARS
@@ -14,7 +14,9 @@ class WrongExpressionError(Exception):
 
 FUNCTIONS: FunctionDictionary = {
     '+': (2, add),
+    '+u': (1, pos),
     '-': (2, sub),
+    '-u': (1, neg),
     '*': (2, mul),
     '/': (2, truediv),
     '^': (2, pow),

shunting_yard/shunting_yard.py

@@ -15,7 +15,9 @@ class MismatchedBracketsError(Exception):
     '-': 10,
     '*': 20,
     '/': 20,
-    '^': 30
+    '^': 30,
+    '-u': 40,
+    '+u': 40,
 }
 
 # Returns the precendence if the operator is one of +-*/^, or infinity if it is a function
@@ -33,6 +35,8 @@ class Associativity(Enum):
     '*': Associativity.LEFT,
     '/': Associativity.LEFT,
     '^': Associativity.RIGHT,
+    '-u': Associativity.RIGHT,
+    '+u': Associativity.LEFT,
 }
 
 

shunting_yard/tokenize.py

@@ -1,6 +1,6 @@
 from typing import Iterator
 
-from shunting_yard.constants import BASE_OPERATORS, NUMBER_CHARS, FUNCTION_CHARS
+from shunting_yard.constants import BASE_OPERATORS, FUNCTION_CHARS, NUMBER_CHARS, UNARY_OPERATORS
 
 
 def tokenize(string: str) -> Iterator[str]:
@@ -9,11 +9,20 @@ def tokenize(string: str) -> Iterator[str]:
 
     string = string.lower()
     cursor = 0
+    is_infix = False
+
     while cursor < len(string):
         char = string[cursor]
-        if char in BASE_OPERATORS or char in '()':
+
+        if not is_infix and char in UNARY_OPERATORS:
+            yield f'{char}u'
+            cursor += 1
+
+        elif char in BASE_OPERATORS or char in '()':
             yield char
             cursor += 1
+            is_infix = (char == ')')
+
         elif char in NUMBER_CHARS:
             # Go through until not a number anymore
             cursor_end = cursor + 1
@@ -22,6 +31,8 @@ def tokenize(string: str) -> Iterator[str]:
 
             yield string[cursor:cursor_end]
             cursor += (cursor_end - cursor)
+            is_infix = True
+
         elif char in FUNCTION_CHARS:
             # Go through until not a number anymore
             cursor_end = cursor + 1
@@ -30,5 +41,7 @@ def tokenize(string: str) -> Iterator[str]:
 
             yield string[cursor:cursor_end]
             cursor += (cursor_end - cursor)
+            is_infix = True
+
         else:
             cursor += 1

tests/test_rpn.py

@@ -13,6 +13,16 @@ def test_base_operators(self):
         self.assertAlmostEqual(compute_rpn('1 2 /'), 0.5)
         self.assertEqual(compute_rpn('1 2 ^'), 1)
 
+    def test_unary_minus(self):
+        self.assertEqual(compute_rpn('1 -u'), -1)
+        self.assertEqual(compute_rpn('1 2 + -u'), -3)
+        self.assertEqual(compute_rpn('1 -u 2 *'), -2)
+
+    def test_unary_plus(self):
+        self.assertEqual(compute_rpn('1 +u'), 1)
+        self.assertEqual(compute_rpn('1 2 + +u'), 3)
+        self.assertEqual(compute_rpn('1 +u 2 *'), 2)
+
     def test_priority_brackets(self):
         self.assertEqual(compute_rpn('2 1 * 3 +'), 5)
         self.assertEqual(compute_rpn('2 1 3 + *'), 8)

tests/test_tokenize.py

@@ -37,7 +37,7 @@ def test_function_one_argument(self):
         self.assertListEqual(list(tokenize('sin(1 + 4)')), ['sin', '(', '1', '+', '4', ')'])
 
     def test_function_no_argument(self):
-        self.assertListEqual(list(tokenize('-3/2*pi')), ['-', '3', '/', '2', '*', 'pi'])
+        self.assertListEqual(list(tokenize('-3/2*pi')), ['-u', '3', '/', '2', '*', 'pi'])
 
     def test_function_2_arguments(self):
         self.assertListEqual(list(tokenize('max(1, 4)')), ['max', '(', '1', '4', ')'])
@@ -48,6 +48,19 @@ def test_function_upper_case(self):
     def test_function_underscore(self):
         self.assertListEqual(list(tokenize('arc_cos(0)')), ['arc_cos', '(', '0', ')'])
 
+    def test_function_unary_minus(self):
+        self.assertListEqual(list(tokenize('-1')), ['-u', '1'])
+        self.assertListEqual(list(tokenize('-(1+2)')), ['-u', '(', '1', '+', '2', ')'])
+        self.assertListEqual(list(tokenize('2*-1')), ['2', '*', '-u', '1'])
+        self.assertListEqual(list(tokenize('-(-1)')), ['-u', '(', '-u', '1', ')'])
+
+    def test_function_unary_plus(self):
+        self.assertListEqual(list(tokenize('+1')), ['+u', '1'])
+        self.assertListEqual(list(tokenize('+(1+2)')), ['+u', '(', '1', '+', '2', ')'])
+        self.assertListEqual(list(tokenize('2*+1')), ['2', '*', '+u', '1'])
+        self.assertListEqual(list(tokenize('+(+1)')), ['+u', '(', '+u', '1', ')'])
+
+
 
 if __name__ == '__main__':
     unittest.main()