Merge pull request #206 from mahmoud/omni-t

Omni-T

glom/core.py

@@ -1431,6 +1431,44 @@ def __call__(self, *args, **kwargs):
             # TODO: typecheck kwarg vals?
         return _t_child(self, '(', (args, kwargs))
 
+    def __add__(self, arg):
+        return _t_child(self, '+', arg)
+
+    def __sub__(self, arg):
+        return _t_child(self, '-', arg)
+
+    def __mul__(self, arg):
+        return _t_child(self, '*', arg)
+
+    def __floordiv__(self, arg):
+        return _t_child(self, '#', arg)
+
+    def __truediv__(self, arg):
+        return _t_child(self, '/', arg)
+
+    __div__ = __truediv__
+
+    def __mod__(self, arg):
+        return _t_child(self, '%', arg)
+
+    def __pow__(self, arg):
+        return _t_child(self, ':', arg)
+
+    def __and__(self, arg):
+        return _t_child(self, '&', arg)
+
+    def __or__(self, arg):
+        return _t_child(self, '|', arg)
+
+    def __xor__(self, arg):
+        return _t_child(self, '^', arg)
+
+    def __invert__(self):
+        return _t_child(self, '~', None)
+
+    def __neg__(self):
+        return _t_child(self, '_', None)
+
     def __(self, name):
         return _t_child(self, '.', '__' + name)
 
@@ -1533,6 +1571,34 @@ def _t_eval(target, _t, scope):
             # if args to the call "reset" their path
             # e.g. "T.a" should mean the same thing
             # in both of these specs: T.a and T.b(T.a)
+        else:  # arithmetic operators
+            try:
+                if op == '+':
+                    cur = cur + arg
+                elif op == '-':
+                    cur = cur - arg
+                elif op == '*':
+                    cur = cur * arg
+                #elif op == '#':
+                #    cur = cur // arg  # TODO: python 2 friendly approach?
+                elif op == '/':
+                    cur = cur / arg
+                elif op == '%':
+                    cur = cur % arg
+                elif op == ':':
+                    cur = cur ** arg
+                elif op == '&':
+                    cur = cur & arg
+                elif op == '|':
+                    cur = cur | arg
+                elif op == '^':
+                    cur = cur ^ arg
+                elif op == '~':
+                    cur = ~cur
+                elif op == '_':
+                    cur = -cur
+            except (TypeError, ZeroDivisionError) as e:
+                pae = PathAccessError(e, Path(_t), i // 2)
         if pae:
             raise pae
         i += 2
@@ -1593,6 +1659,18 @@ def _format_t(path, root=T):
             prepr.append(format_invocation(args=args, kwargs=kwargs, repr=bbrepr))
         elif op == 'P':
             return _format_path(path)
+        elif op in ('_', '~'):  # unary arithmetic operators
+            if any([o in path[:i] for o in '+-/%:&|^~_']):
+                prepr = ['('] + prepr + [')']
+            prepr = ['-' if op == '_' else op] + prepr
+        else:  # binary arithmetic operators
+            formatted_arg = bbrepr(arg)
+            if type(arg) is TType:
+                arg_path = _T_PATHS[arg]
+                if any([o in arg_path for o in '+-/%:&|^~_']):
+                    formatted_arg = '(' + formatted_arg + ')'
+            prepr.append(' ' + ('**' if op == ':' else op) + ' ')
+            prepr.append(formatted_arg)
         i += 2
     return "".join(prepr)
 

glom/test/test_path_and_t.py

@@ -229,6 +229,38 @@ def test_t_dict_key():
     assert glom(target, {T['a']: 'a'}) == {'A': 'A'}
 
 
+def test_t_arithmetic():
+    t = 2
+    assert glom(t, T + T) == 4
+    assert glom(t, T * T) == 4
+    assert glom(t, T ** T) == 4
+    assert glom(t, T / 1) == 2
+    assert glom(t, T % 1) == 0
+    assert glom(t, T - 1) == 1
+    assert glom(t, T & T) == 2
+    assert glom(t, T | 1) == 3
+    assert glom(t, T ^ T) == 0
+    assert glom(2, ~T) == -3
+    assert glom(t, -T) == -2
+
+
+def test_t_arithmetic_reprs():
+    assert repr(T + T) == "T + T"
+    assert repr(T + (T / 2 * (T - 5) % 4)) == "T + (T / 2 * (T - 5) % 4)"
+    assert repr(T & 7 | (T ^ 6)) == "T & 7 | (T ^ 6)"
+    assert repr(-(~T)) == "-(~T)"
+
+
+def test_t_arithmetic_errors():
+    with raises(PathAccessError, match='zero'):
+        glom(0, T / 0)
+
+    with raises(PathAccessError, match='unsupported operand type'):
+        glom(None, T / 2)
+
+    return
+
+
 def test_t_dunders():
     with raises(AttributeError) as exc_info:
         T.__name__