add type_context to TypeChecker and fix logic in neg method callback

TH3CHARLie · TH3CHARLie · commit 32a99a2422d4 · 2019-11-06T10:18:44.000+08:00
diff --git a/mypy/checker.py b/mypy/checker.py
@@ -219,6 +219,7 @@ def __init__(self, errors: Errors, modules: Dict[str, MypyFile], options: Option
         self.msg = MessageBuilder(errors, modules)
         self.plugin = plugin
         self.expr_checker = mypy.checkexpr.ExpressionChecker(self, self.msg, self.plugin)
+        self.type_context = self.expr_checker.type_context
         self.tscope = Scope()
         self.scope = CheckerScope(tree)
         self.binder = ConditionalTypeBinder()
diff --git a/mypy/plugin.py b/mypy/plugin.py
@@ -215,6 +215,8 @@ class CheckerPluginInterface:
     msg = None  # type: MessageBuilder
     options = None  # type: Options
     path = None  # type: str
+    # Type context for type inference
+    type_context = None  # type: List[Optional[Type]]
 
     @abstractmethod
     def fail(self, msg: str, ctx: Context, *, code: Optional[ErrorCode] = None) -> None:
diff --git a/mypy/plugins/default.py b/mypy/plugins/default.py
@@ -14,6 +14,7 @@
 )
 from mypy.subtypes import is_subtype
 from mypy.typeops import make_simplified_union
+from mypy.checkexpr import is_literal_type_like
 
 
 class DefaultPlugin(Plugin):
@@ -431,11 +432,18 @@ def int_neg_callback(ctx: MethodContext) -> Type:
         value = ctx.type.last_known_value.value
         fallback = ctx.type.last_known_value.fallback
         if isinstance(value, int):
-            return Instance(ctx.type.type, ctx.type.args,
-                last_known_value=LiteralType(-value, fallback))
+            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, fallback)
+            return LiteralType(value=-value, fallback=fallback)
     return ctx.default_return_type
