Skip to content

Commit

Permalink
Imporves __div__ with __future__ import on py2 (python#11276)
Browse files Browse the repository at this point in the history
  • Loading branch information
sobolevn authored Oct 6, 2021
1 parent d0a41a0 commit e6b91bd
Show file tree
Hide file tree
Showing 4 changed files with 39 additions and 5 deletions.
7 changes: 5 additions & 2 deletions mypy/build.py
Original file line number Diff line number Diff line change
Expand Up @@ -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]:
Expand Down
8 changes: 7 additions & 1 deletion mypy/checker.py
Original file line number Diff line number Diff line change
Expand Up @@ -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.
Expand Down Expand Up @@ -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
Expand Down
3 changes: 1 addition & 2 deletions mypy/checkexpr.py
Original file line number Diff line number Diff line change
Expand Up @@ -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]

Expand Down
26 changes: 26 additions & 0 deletions test-data/unit/check-expressions.test
Original file line number Diff line number Diff line change
Expand Up @@ -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():
Expand Down

0 comments on commit e6b91bd

Please sign in to comment.