From fcef037323b7e9833c0166446e0a08ea8bca23aa Mon Sep 17 00:00:00 2001
From: Tim Rid <6593626+timrid@users.noreply.github.com>
Date: Tue, 29 Dec 2020 13:00:24 +0100
Subject: [PATCH] added missing operator overloads

---
 construct-stubs/expr.pyi        | 31 +++++++++++++++++++++++
 scripts/expr_mixin_generator.py | 44 +++++++++++++++++++++++++--------
 2 files changed, 65 insertions(+), 10 deletions(-)

diff --git a/construct-stubs/expr.pyi b/construct-stubs/expr.pyi
index 46ed44e..603efbe 100644
--- a/construct-stubs/expr.pyi
+++ b/construct-stubs/expr.pyi
@@ -489,6 +489,37 @@ class ExprMixin(t.Generic[ReturnType], object):
     @t.overload
     def __ne__(self, other: t.Any) -> BinExpr[t.Any]: ...
 
+    # __neg__ ##########################################################################################################
+    @t.overload
+    def __neg__(self: ExprMixin[int]) -> BinExpr[int]: ...
+    @t.overload
+    def __neg__(self: ExprMixin[bool]) -> BinExpr[int]: ...
+    @t.overload
+    def __neg__(self: ExprMixin[float]) -> BinExpr[float]: ...
+    @t.overload
+    def __neg__(self) -> UniExpr[t.Any]: ...
+
+    # __pos__ ##########################################################################################################
+    @t.overload
+    def __pos__(self: ExprMixin[int]) -> BinExpr[int]: ...
+    @t.overload
+    def __pos__(self: ExprMixin[bool]) -> BinExpr[int]: ...
+    @t.overload
+    def __pos__(self: ExprMixin[float]) -> BinExpr[float]: ...
+    @t.overload
+    def __pos__(self) -> UniExpr[t.Any]: ...
+
+    # __invert__ #######################################################################################################
+    @t.overload
+    def __invert__(self: ExprMixin[int]) -> BinExpr[int]: ...
+    @t.overload
+    def __invert__(self: ExprMixin[bool]) -> BinExpr[int]: ...
+    @t.overload
+    def __invert__(self) -> UniExpr[t.Any]: ...
+
+    # __inv__ ##########################################################################################################
+    def __inv__(self) -> UniExpr[t.Any]: ...
+
 class UniExpr(ExprMixin[ReturnType]):
     def __init__(self, op: UniOperator, operand: t.Any) -> None: ...
     def __call__(self, obj: t.Union[Context, dict[str, t.Any], t.Any], *args: t.Any) -> ReturnType: ...
diff --git a/scripts/expr_mixin_generator.py b/scripts/expr_mixin_generator.py
index 91a2613..45b7c8e 100644
--- a/scripts/expr_mixin_generator.py
+++ b/scripts/expr_mixin_generator.py
@@ -6,7 +6,7 @@ testobjs = [
     float(10.0),
 ]
 
-operators = [
+operators_bin = [
     "__add__",
     "__sub__",
     "__mul__",
@@ -35,11 +35,6 @@ operators = [
     "__rand__",
     "__ror__",
 
-    # "__neg__",
-    # "__pos__",
-    # "__invert__",
-    # "__inv__",
-
     "__contains__",
     "__gt__",
     "__ge__",
@@ -49,8 +44,15 @@ operators = [
     "__ne__",
 ]
 
+operators_uni = [
+    "__neg__",
+    "__pos__",
+    "__invert__",
+    "__inv__",
+]
 
-def create_overload(op, lhs=None, rhs=None):
+
+def create_overload_bin(op, lhs=None, rhs=None):
     if lhs is None and rhs is None:
         print("    @t.overload")
         print(f"    def {op}(self, other: t.Any) -> BinExpr[t.Any]: ...")
@@ -66,12 +68,34 @@ def create_overload(op, lhs=None, rhs=None):
         except AttributeError:
             pass
 
+def create_overload_uni(op, obj=None):
+    if obj is None:
+        print("    @t.overload")
+        print(f"    def {op}(self) -> UniExpr[t.Any]: ...")
+    else:
+        try:
+            result = getattr(obj, op)()
+            obj_type = type(obj).__name__
+            result_type = type(result).__name__
+            if result_type != "NotImplementedType":
+                print("    @t.overload")
+                print(f"    def {op}(self: ExprMixin[{obj_type}]) -> BinExpr[{result_type}]: ...")
+        except AttributeError:
+            pass
+
 
 print("class ExprMixin(t.Generic[ReturnType], object):")
-for op in operators:
+for op in operators_bin:
     print(f"    # {op} ".ljust(120, "#"))
     for lhs in testobjs:
         for rhs in testobjs:
-            create_overload(op, lhs, rhs)
-    create_overload(op)
+            create_overload_bin(op, lhs, rhs)
+    create_overload_bin(op)
     print("")
+
+for op in operators_uni:
+    print(f"    # {op} ".ljust(120, "#"))
+    for obj in testobjs:
+        create_overload_uni(op, obj)
+    create_overload_uni(op)
+    print("")
\ No newline at end of file