Make other comparison operators work

Author: hatal175

mypy/checker.py

@@ -4256,8 +4256,9 @@ def has_no_custom_eq_checks(t: Type) -> bool:
                         else_map = {}
                 else:
                     # comparison expression with len
-                    if operator in {'==', '!='}:
+                    if operator in {'==', '!=', '>=', '<=', '<', '>'}:
                         if_map, else_map = self.refine_len_comparison_expression(
+                            operator,
                             operands,
                             operand_types,
                             expr_indices,
@@ -4267,7 +4268,7 @@ def has_no_custom_eq_checks(t: Type) -> bool:
                         if_map = {}
                         else_map = {}
 
-                if operator in {'is not', '!=', 'not in'}:
+                if operator in {'is not', '!=', 'not in', '<', '>'}:
                     if_map, else_map = else_map, if_map
 
                 partial_type_maps.append((if_map, else_map))
@@ -4608,6 +4609,7 @@ def refine_identity_comparison_expression(self,
         return reduce_conditional_maps(partial_type_maps)
 
     def refine_len_comparison_expression(self,
+                                         operator: str,
                                          operands: List[Expression],
                                          operand_types: List[Type],
                                          chain_indices: List[int],
@@ -4641,17 +4643,24 @@ def refine_len_comparison_expression(self,
         """
 
         target = None  # type: Optional[int]
+        target_index = None  # type: Optional[int]
         possible_target_indices = []
         for i in chain_indices:
             expr_type = operand_types[i]
             expr_type = coerce_to_literal(expr_type)
             if not isinstance(get_proper_type(expr_type), LiteralType):
                 continue
             if target and target != expr_type.value:
-                # We have multiple different target values. So the 'if' branch
-                # must be unreachable.
-                return None, {}
+                if operator in {'==', '!='}:
+                    # We have multiple different target values. So the 'if' branch
+                    # must be unreachable.
+                    return None, {}
+                else:
+                    # Other operators can go either way
+                    return {}, {}
+
             target = expr_type.value
+            target_index = i
             possible_target_indices.append(i)
 
         # There's nothing we can currently infer if none of the operands are valid targets,
@@ -4671,19 +4680,25 @@ def refine_len_comparison_expression(self,
             # We intentionally use 'conditional_type_map' directly here instead of
             # 'self.conditional_type_map_with_intersection': we only compute ad-hoc
             # intersections when working with pure instances.
-            partial_type_maps.append(self.conditional_len_map(expr, expr_type, target))
+            partial_type_maps.append(
+                self.conditional_len_map(operator, expr, expr_type, i, target, target_index))
 
         return reduce_conditional_maps(partial_type_maps)
 
-    def narrow_type_by_length(self, typ: Type, length: int) -> Type:
+    def narrow_type_by_length(self, operator: str, typ: Type, length: int) -> Type:
+        if operator not in {"==", "!="}:
+            return typ
         if (isinstance(typ, Instance) and typ.type.fullname == "builtins.tuple" and length >= 0):
             return TupleType(typ.args[0:1] * length, self.named_type('builtins.tuple'))
         return typ
 
     def conditional_len_map(self,
+                            operator: str,
                             expr: Expression,
                             current_type: Optional[Type],
+                            expr_index: int,
                             length: Optional[int],
+                            target_index: int,
                             ) -> Tuple[TypeMap, TypeMap]:
         """Takes in an expression, the current type of the expression, and a
         proposed length of that expression.
@@ -4702,13 +4717,36 @@ def conditional_len_map(self,
                 possible_types = union_items(current_type)
                 len_of_types = [len_of_type(typ) for typ in possible_types]
 
+                if operator in {'>=', '<=', '<', '>'} and target_index < expr_index:
+                    if operator == '>=':
+                        operator = '<='
+                    elif operator == '>':
+                        operator = '<'
+                    elif operator == '<=':
+                        operator = '>='
+                    else:
+                        operator = '>'
+
+                # We reverse the map for some operator outside this function
+                length_op_translator = {
+                    '==': int.__eq__,
+                    '!=': int.__eq__,
+                    '>=': int.__ge__,
+                    '<': int.__ge__,
+                    '<=': int.__le__,
+                    '>': int.__le__,
+                }
+
+                assert operator in length_op_translator
+                length_op = length_op_translator[operator]
+
                 proposed_type = make_simplified_union([
-                    self.narrow_type_by_length(typ, length)
+                    self.narrow_type_by_length(operator, typ, length)
                     for typ, l in zip(possible_types, len_of_types)
-                    if l is None or l == length])
+                    if l is None or length_op(l, length)])
                 remaining_type = make_simplified_union([
                     typ for typ, l in zip(possible_types, len_of_types)
-                    if l is None or l != length])
+                    if l is None or not length_op(l, length)])
                 if_map = (
                     {} if is_same_type(proposed_type, current_type)
                     else {expr: proposed_type})

test-data/unit/check-narrowing.test

@@ -1090,7 +1090,10 @@ if len(x) == 3:
     reveal_type(x) # N: Revealed type is "Union[builtins.str, builtins.list[builtins.int]]"
 else:
     reveal_type(x) # N: Revealed type is "Union[builtins.str, builtins.list[builtins.int]]"
+[builtins fixtures/len.pyi]
 
+[case testNarrowingLenAnyListElseNotAffected]
+from typing import Any
 def f(self, value: Any) -> Any:
     if isinstance(value, list) and len(value) == 0:
         reveal_type(value) # N: Revealed type is "builtins.list[Any]"
@@ -1143,3 +1146,48 @@ fin: Final = 3
 if len(x) == fin:
     reveal_type(x) # N: Revealed type is "Tuple[builtins.int, builtins.int, builtins.int]"
 [builtins fixtures/len.pyi]
+
+[case testNarrowingLenBiggerThan]
+from typing import Tuple, Union
+
+VarTuple = Union[Tuple[int], Tuple[int, int], Tuple[int, int, int]]
+
+def make_tuple() -> VarTuple:
+    return (1, 1)
+
+x = make_tuple()
+if len(x) > 1:
+    reveal_type(x) # N: Revealed type is "Union[Tuple[builtins.int, builtins.int], Tuple[builtins.int, builtins.int, builtins.int]]"
+else:
+    reveal_type(x) # N: Revealed type is "Tuple[builtins.int]"
+
+if len(x) < 2:
+    reveal_type(x) # N: Revealed type is "Tuple[builtins.int]"
+else:
+    reveal_type(x) # N: Revealed type is "Union[Tuple[builtins.int, builtins.int], Tuple[builtins.int, builtins.int, builtins.int]]"
+
+if len(x) >= 2:
+    reveal_type(x) # N: Revealed type is "Union[Tuple[builtins.int, builtins.int], Tuple[builtins.int, builtins.int, builtins.int]]"
+else:
+    reveal_type(x) # N: Revealed type is "Tuple[builtins.int]"
+
+if len(x) <= 2:
+    reveal_type(x) # N: Revealed type is "Union[Tuple[builtins.int], Tuple[builtins.int, builtins.int]]"
+else:
+    reveal_type(x) # N: Revealed type is "Tuple[builtins.int, builtins.int, builtins.int]"
+[builtins fixtures/len.pyi]
+
+[case testNarrowingLenBiggerThanVariantTuple]
+from typing import Tuple
+
+VarTuple = Tuple[int, ...]
+
+def make_tuple() -> VarTuple:
+    return (1, 1)
+
+x = make_tuple()
+if len(x) < 3:
+    reveal_type(x) # N: Revealed type is "builtins.tuple[builtins.int]"
+else:
+    reveal_type(x) # N: Revealed type is "builtins.tuple[builtins.int]"
+[builtins fixtures/len.pyi]

test-data/unit/fixtures/len.pyi

@@ -26,6 +26,10 @@ class int:
     def __add__(self, other: 'int') -> 'int': pass
     def __eq__(self, other: 'int') -> 'bool': pass
     def __ne__(self, other: 'int') -> 'bool': pass
+    def __lt__(self, n: 'int') -> 'bool': pass
+    def __gt__(self, n: 'int') -> 'bool': pass
+    def __le__(self, n: 'int') -> 'bool': pass
+    def __ge__(self, n: 'int') -> 'bool': pass
 class float: pass
 class bool(int): pass
 class str: