Only allow some unary expressions

python · sobolevn · Jul 13, 2023 · Jul 13, 2023 · Jul 13, 2023 · Jul 13, 2023
commit ec71ebe662a0ce87a8b29933e6949f79a83d1080
diff --git a/mypy/stubgen.py b/mypy/stubgen.py
@@ -1397,7 +1397,7 @@ def is_private_member(self, fullname: str) -> bool:
     def get_str_type_of_node(
         self, rvalue: Expression, can_infer_optional: bool = False, can_be_any: bool = True
     ) -> str:
-        rvalue = self.unwrap_unary_expr(rvalue)
+        rvalue = self.maybe_unwrap_unary_expr(rvalue)
 
         if isinstance(rvalue, IntExpr):
             return "int"
@@ -1410,8 +1410,8 @@ def get_str_type_of_node(
         if isinstance(rvalue, ComplexExpr):  # 1j
             return "complex"
         if isinstance(rvalue, OpExpr) and rvalue.op in ("-", "+"):  # -1j + 1
-            if isinstance(self.unwrap_unary_expr(rvalue.left), ComplexExpr) or isinstance(
-                self.unwrap_unary_expr(rvalue.right), ComplexExpr
            if isinstance(self.maybe_unwrap_unary_expr(rvalue.left), ComplexExpr) or isinstance(
+                self.maybe_unwrap_unary_expr(rvalue.right), ComplexExpr
             ):
                 return "complex"
         if isinstance(rvalue, NameExpr) and rvalue.name in ("True", "False"):
@@ -1425,10 +1425,38 @@ def get_str_type_of_node(
         else:
             return ""
 
-    def unwrap_unary_expr(self, expr: Expression) -> Expression:
-        # Unwrap (possibly nested) unary expressions:
-        while isinstance(expr, UnaryExpr):
-            expr = expr.expr
+    def maybe_unwrap_unary_expr(self, expr: Expression) -> Expression:
+        """Unwrap (possibly nested) unary expressions.
+
+        But, some unary expressions can change the type of expression.
+        While we want to preserve it. For example, `~True` is `int`.
+        So, we only allow a subset of unary expressions to be unwrapped.
+        """
+        if not isinstance(expr, UnaryExpr):
+            return expr
+
+        # First, try to unwrap `[+-]+ (int|float|complex)` expr:
+        math_ops = ("+", "-")
+        if expr.op in math_ops:
+            while isinstance(expr, UnaryExpr):
+                if expr.op not in math_ops or not isinstance(
+                    expr.expr, (IntExpr, FloatExpr, ComplexExpr, UnaryExpr)
+                ):
+                    break
+                expr = expr.expr
+            return expr
+
+        # Next, try `not bool` expr:
+        if expr.op == "not":
+            while isinstance(expr, UnaryExpr):
+                if expr.op != "not" or not isinstance(expr.expr, (NameExpr, UnaryExpr)):
+                    break
+                if isinstance(expr.expr, NameExpr) and expr.expr.name not in ("True", "False"):
+                    break
+                expr = expr.expr
+            return expr
+
+        # This is some other unary expr, we cannot do anything with it (yet?).
        return expr
 
     def print_annotation(self, t: Type) -> str:

diff --git a/test-data/unit/stubgen.test b/test-data/unit/stubgen.test
@@ -147,12 +147,12 @@ bl2: bool
 bts: bytes
 
 [case testVariablesWithUnary]
-i = ++1
+i = +-1
 f = -1.5
 c1 = -1j
 c2 = -1j + 1
 bl1 = not True
-bl2 = not False
+bl2 = not not False
 [out]
 i: int
 f: float
@@ -161,6 +161,19 @@ c2: complex
 bl1: bool
 bl2: bool
 
+[case testVariablesWithUnaryWrong]
+i = not +1
+bl1 = -True
+bl2 = not -False
+bl3 = -(not False)
+[out]
+from _typeshed import Incomplete
+
+i: Incomplete
+bl1: Incomplete
+bl2: Incomplete
+bl3: Incomplete
+
 [case testAnnotatedVariable]
 x: int = 1
 [out]