@@ -2118,12 +2118,31 @@ def method_type(self, func: FuncBase) -> FunctionLike:
21182118 return method_type_with_fallback (func , self .named_type ('builtins.function' ))
21192119
21202120
2121+ # Data structure returned by find_isinstance_check representing
2122+ # information learned from the truth or falsehood of a condition. The
2123+ # dict maps nodes representing expressions like 'a[0].x' to their
2124+ # refined types under the assumption that the condition has a
2125+ # particular truth value. A value of None means that the condition can
2126+ # never have that truth value.
2127+
2128+ # NB: The keys of this dict are nodes in the original source program,
2129+ # which are compared by reference equality--effectively, being *the
2130+ # same* expression of the program, not just two identical expressions
2131+ # (such as two references to the same variable). TODO: it would
2132+ # probably be better to have the dict keyed by the nodes' literal_hash
2133+ # field instead.
2134+
2135+ # NB: This should be `TypeMap = Optional[Dict[Node, Type]]`!
2136+ # But see https://github.com/python/mypy/issues/1637
2137+ TypeMap = Dict [Node , Type ]
2138+
2139+
21212140def conditional_type_map (expr : Node ,
21222141 current_type : Optional [Type ],
21232142 proposed_type : Optional [Type ],
21242143 * ,
21252144 weak : bool = False
2126- ) -> Tuple [Optional [ Dict [ Node , Type ]], Optional [ Dict [ Node , Type ]] ]:
2145+ ) -> Tuple [TypeMap , TypeMap ]:
21272146 """Takes in an expression, the current type of the expression, and a
21282147 proposed type of that expression.
21292148
@@ -2154,10 +2173,54 @@ def is_literal_none(n: Node) -> bool:
21542173 return isinstance (n , NameExpr ) and n .fullname == 'builtins.None'
21552174
21562175
2176+ def and_conditional_maps (m1 : TypeMap , m2 : TypeMap ) -> TypeMap :
2177+ """Calculate what information we can learn from the truth of (e1 and e2)
2178+ in terms of the information that we can learn from the truth of e1 and
2179+ the truth of e2.
2180+ """
2181+
2182+ if m1 is None or m2 is None :
2183+ # One of the conditions can never be true.
2184+ return None
2185+ # Both conditions can be true; combine the information. Anything
2186+ # we learn from either conditions's truth is valid. If the same
2187+ # expression's type is refined by both conditions, we somewhat
2188+ # arbitrarily give precedence to m2. (In the future, we could use
2189+ # an intersection type.)
2190+ result = m2 .copy ()
2191+ m2_keys = set (n2 .literal_hash for n2 in m2 )
2192+ for n1 in m1 :
2193+ if n1 .literal_hash not in m2_keys :
2194+ result [n1 ] = m1 [n1 ]
2195+ return result
2196+
2197+
2198+ def or_conditional_maps (m1 : TypeMap , m2 : TypeMap ) -> TypeMap :
2199+ """Calculate what information we can learn from the truth of (e1 or e2)
2200+ in terms of the information that we can learn from the truth of e1 and
2201+ the truth of e2.
2202+ """
2203+
2204+ if m1 is None :
2205+ return m2
2206+ if m2 is None :
2207+ return m1
2208+ # Both conditions can be true. Combine information about
2209+ # expressions whose type is refined by both conditions. (We do not
2210+ # learn anything about expressions whose type is refined by only
2211+ # one condition.)
2212+ result = {}
2213+ for n1 in m1 :
2214+ for n2 in m2 :
2215+ if n1 .literal_hash == n2 .literal_hash :
2216+ result [n1 ] = UnionType .make_simplified_union ([m1 [n1 ], m2 [n2 ]])
2217+ return result
2218+
2219+
21572220def find_isinstance_check (node : Node ,
21582221 type_map : Dict [Node , Type ],
21592222 weak : bool = False
2160- ) -> Tuple [Optional [ Dict [ Node , Type ]], Optional [ Dict [ Node , Type ]] ]:
2223+ ) -> Tuple [TypeMap , TypeMap ]:
21612224 """Find any isinstance checks (within a chain of ands). Includes
21622225 implicit and explicit checks for None.
21632226
@@ -2201,7 +2264,24 @@ def find_isinstance_check(node: Node,
22012264 _ , if_vars = conditional_type_map (node , vartype , NoneTyp (), weak = weak )
22022265 return if_vars , {}
22032266 elif isinstance (node , OpExpr ) and node .op == 'and' :
2204- left_if_vars , right_else_vars = find_isinstance_check (
2267+ left_if_vars , left_else_vars = find_isinstance_check (
2268+ node .left ,
2269+ type_map ,
2270+ weak ,
2271+ )
2272+
2273+ right_if_vars , right_else_vars = find_isinstance_check (
2274+ node .right ,
2275+ type_map ,
2276+ weak ,
2277+ )
2278+
2279+ # (e1 and e2) is true if both e1 and e2 are true,
2280+ # and false if at least one of e1 and e2 is false.
2281+ return (and_conditional_maps (left_if_vars , right_if_vars ),
2282+ or_conditional_maps (left_else_vars , right_else_vars ))
2283+ elif isinstance (node , OpExpr ) and node .op == 'or' :
2284+ left_if_vars , left_else_vars = find_isinstance_check (
22052285 node .left ,
22062286 type_map ,
22072287 weak ,
@@ -2212,16 +2292,11 @@ def find_isinstance_check(node: Node,
22122292 type_map ,
22132293 weak ,
22142294 )
2215- if left_if_vars :
2216- if right_if_vars is not None :
2217- left_if_vars .update (right_if_vars )
2218- else :
2219- left_if_vars = None
2220- else :
2221- left_if_vars = right_if_vars
22222295
2223- # Make no claim about the types in else
2224- return left_if_vars , {}
2296+ # (e1 or e2) is true if at least one of e1 or e2 is true,
2297+ # and false if both e1 and e2 are false.
2298+ return (or_conditional_maps (left_if_vars , right_if_vars ),
2299+ and_conditional_maps (left_else_vars , right_else_vars ))
22252300 elif isinstance (node , UnaryExpr ) and node .op == 'not' :
22262301 left , right = find_isinstance_check (node .expr , type_map , weak )
22272302 return right , left
0 commit comments