Imporves __div__ with __future__ import on py2 (#11276)

Author: sobolevn

mypy/build.py

@@ -2178,8 +2178,11 @@ def type_checker(self) -> TypeChecker:
         if not self._type_checker:
             assert self.tree is not None, "Internal error: must be called on parsed file only"
             manager = self.manager
-            self._type_checker = TypeChecker(manager.errors, manager.modules, self.options,
-                                             self.tree, self.xpath, manager.plugin)
+            self._type_checker = TypeChecker(
+                manager.errors, manager.modules, self.options,
+                self.tree, self.xpath, manager.plugin,
+                self.manager.semantic_analyzer.future_import_flags,
+            )
         return self._type_checker
 
     def type_map(self) -> Dict[Expression, Type]:

mypy/checker.py

@@ -216,8 +216,12 @@ class TypeChecker(NodeVisitor[None], CheckerPluginInterface):
     # functions such as open(), etc.
     plugin: Plugin
 
+    # Future flags that we get from semantic analyzer.
+    future_import_flags: Set[str]
+
     def __init__(self, errors: Errors, modules: Dict[str, MypyFile], options: Options,
-                 tree: MypyFile, path: str, plugin: Plugin) -> None:
+                 tree: MypyFile, path: str, plugin: Plugin,
+                 future_import_flags: Set[str]) -> None:
         """Construct a type checker.
 
         Use errors to report type check errors.
@@ -263,6 +267,8 @@ def __init__(self, errors: Errors, modules: Dict[str, MypyFile], options: Option
         # argument through various `checker` and `checkmember` functions.
         self._is_final_def = False
 
+        self.future_import_flags = future_import_flags
+
     @property
     def type_context(self) -> List[Optional[Type]]:
         return self.expr_checker.type_context

mypy/checkexpr.py

@@ -2373,8 +2373,7 @@ def dangerous_comparison(self, left: Type, right: Type,
 
     def get_operator_method(self, op: str) -> str:
         if op == '/' and self.chk.options.python_version[0] == 2:
-            # TODO also check for "from __future__ import division"
-            return '__div__'
+            return '__truediv__' if 'division' in self.chk.future_import_flags else '__div__'
         else:
             return operators.op_methods[op]
 

test-data/unit/check-expressions.test

@@ -221,6 +221,32 @@ class C:
     pass
 [builtins fixtures/tuple.pyi]
 
+[case testDivPython2]
+# flags: --python-version 2.7
+class A(object):
+    def __div__(self, other):
+        # type: (A, str) -> str
+        return 'a'
+
+a = A()
+reveal_type(a / 'b')  # N: Revealed type is "builtins.str"
+a / 1  # E: Unsupported operand types for / ("A" and "int")
+[builtins fixtures/bool.pyi]
+
+[case testDivPython2FutureImport]
+# flags: --python-version 2.7
+from __future__ import division
+
+class A(object):
+    def __truediv__(self, other):
+        # type: (A, str) -> str
+        return 'a'
+
+a = A()
+reveal_type(a / 'b')  # N: Revealed type is "builtins.str"
+a / 1  # E: Unsupported operand types for / ("A" and "int")
+[builtins fixtures/bool.pyi]
+
 [case testIntDiv]
 a, b, c = None, None, None # type: (A, B, C)
 if int():