Support Negative Int Literal (#7878)

TH3CHARLie · ilevkivskyi · commit 8010414c2b79 · 2019-11-06T17:07:33.000Z
This PR aims to support negative int Literal, resolves #7844 Note that this also required adding `type_context` to plugin API, but this may be helpful also for other plugins.
diff --git a/mypy/checker.py b/mypy/checker.py
@@ -251,6 +251,10 @@ def __init__(self, errors: Errors, modules: Dict[str, MypyFile], options: Option
         # argument through various `checker` and `checkmember` functions.
         self._is_final_def = False
 
+    @property
+    def type_context(self) -> List[Optional[Type]]:
+        return self.expr_checker.type_context
+
     def reset(self) -> None:
         """Cleanup stale state that might be left over from a typechecking run.
 
diff --git a/mypy/plugin.py b/mypy/plugin.py
@@ -216,6 +216,12 @@ class CheckerPluginInterface:
     options = None  # type: Options
     path = None  # type: str
 
+    # Type context for type inference
+    @abstractproperty
+    def type_context(self) -> List[Optional[Type]]:
+        """Return the type context of the plugin"""
+        raise NotImplementedError
+
     @abstractmethod
     def fail(self, msg: str, ctx: Context, *, code: Optional[ErrorCode] = None) -> None:
         """Emit an error message at given location."""
diff --git a/mypy/plugins/default.py b/mypy/plugins/default.py
@@ -10,10 +10,11 @@
 from mypy.plugins.common import try_getting_str_literals
 from mypy.types import (
     Type, Instance, AnyType, TypeOfAny, CallableType, NoneType, TypedDictType,
-    TypeVarType, TPDICT_FB_NAMES, get_proper_type
+    TypeVarType, TPDICT_FB_NAMES, get_proper_type, LiteralType
 )
 from mypy.subtypes import is_subtype
 from mypy.typeops import make_simplified_union
+from mypy.checkexpr import is_literal_type_like
 
 
 class DefaultPlugin(Plugin):
@@ -57,6 +58,8 @@ def get_method_hook(self, fullname: str
             return typed_dict_get_callback
         elif fullname == 'builtins.int.__pow__':
             return int_pow_callback
+        elif fullname == 'builtins.int.__neg__':
+            return int_neg_callback
         elif fullname in set(n + '.setdefault' for n in TPDICT_FB_NAMES):
             return typed_dict_setdefault_callback
         elif fullname in set(n + '.pop' for n in TPDICT_FB_NAMES):
@@ -417,3 +420,30 @@ def int_pow_callback(ctx: MethodContext) -> Type:
         else:
             return ctx.api.named_generic_type('builtins.float', [])
     return ctx.default_return_type
+
+
+def int_neg_callback(ctx: MethodContext) -> Type:
+    """Infer a more precise return type for int.__neg__.
+
+    This is mainly used to infer the return type as LiteralType
+    if the original underlying object is a LiteralType object
+    """
+    if isinstance(ctx.type, Instance) and ctx.type.last_known_value is not None:
+        value = ctx.type.last_known_value.value
+        fallback = ctx.type.last_known_value.fallback
+        if isinstance(value, int):
+            if is_literal_type_like(ctx.api.type_context[-1]):
+                return LiteralType(value=-value, fallback=fallback)
+            else:
+                return ctx.type.copy_modified(last_known_value=LiteralType(
+                    value=-value,
+                    fallback=ctx.type,
+                    line=ctx.type.line,
+                    column=ctx.type.column,
+                ))
+    elif isinstance(ctx.type, LiteralType):
+        value = ctx.type.value
+        fallback = ctx.type.fallback
+        if isinstance(value, int):
+            return LiteralType(value=-value, fallback=fallback)
+    return ctx.default_return_type
diff --git a/test-data/unit/check-literal.test b/test-data/unit/check-literal.test
@@ -3093,3 +3093,28 @@ def test() -> None:
     if x == 2:  # OK
         ...
 [builtins fixtures/bool.pyi]
+
+[case testNegativeIntLiteral]
+from typing_extensions import Literal
+
+a: Literal[-2] = -2
+b: Literal[-1] = -1
+c: Literal[0] = 0
+d: Literal[1] = 1
+e: Literal[2] = 2
+[out]
+[builtins fixtures/float.pyi]
+
+[case testNegativeIntLiteralWithFinal]
+from typing_extensions import Literal, Final
+
+ONE: Final = 1
+x: Literal[-1] = -ONE
+
+TWO: Final = 2
+THREE: Final = 3
+
+err_code = -TWO
+if bool():
+    err_code = -THREE
+[builtins fixtures/float.pyi]
diff --git a/test-data/unit/fixtures/float.pyi b/test-data/unit/fixtures/float.pyi
@@ -26,6 +26,7 @@ class int:
     def __float__(self) -> float: ...
     def __int__(self) -> int: ...
     def __mul__(self, x: int) -> int: ...
+    def __neg__(self) -> int: ...
     def __rmul__(self, x: int) -> int: ...
 
 class float:
