Repository: quora/pyanalyze
Pull request #347 (2021-12-19 03:37:07)

huge ranges are still always entered


Patch:
diff --git a/pyanalyze/name_check_visitor.py b/pyanalyze/name_check_visitor.py
index c51cc1b..f4b002c 100644
--- a/pyanalyze/name_check_visitor.py
+++ b/pyanalyze/name_check_visitor.py
@@ -113,6 +113,7 @@ from .asynq_checker import AsyncFunctionKind, AsynqChecker, FunctionInfo
 from .yield_checker import YieldChecker
 from .type_object import TypeObject, get_mro
 from .value import (
+    AlwaysPresentExtension,
     AnnotatedValue,
     AnySource,
     AnyValue,
@@ -125,6 +126,7 @@ from .value import (
     NO_RETURN_VALUE,
     kv_pairs_from_mapping,
     make_weak,
+    unannotate_value,
     unite_and_simplify,
     unite_values,
     KnownValue,
@@ -3051,7 +3053,10 @@ class NameCheckVisitor(
         if self.config.FOR_LOOP_ALWAYS_ENTERED:
             always_entered = True
         elif isinstance(iterated_value, Value):
-            always_entered = False
+            iterated_value, present = unannotate_value(
+                iterated_value, AlwaysPresentExtension
+            )
+            always_entered = bool(present)
         else:
             always_entered = len(iterated_value) > 0
         if not isinstance(iterated_value, Value):
diff --git a/pyanalyze/value.py b/pyanalyze/value.py
index 448434f..3a4c411 100644
--- a/pyanalyze/value.py
+++ b/pyanalyze/value.py
@@ -48,6 +48,7 @@ from typing import (
 from typing_extensions import Literal, Protocol
 
 import pyanalyze
+from pyanalyze import extensions
 from pyanalyze.extensions import CustomCheck
 
 from .safe import all_of_type, safe_issubclass, safe_isinstance
@@ -1686,6 +1687,15 @@ class WeakExtension(Extension):
     """
 
 
+@dataclass(frozen=True)
+class AlwaysPresentExtension(Extension):
+    """Extension that indicates that an iterable value is nonempty.
+
+    Currently cannot be used from user code.
+
+    """
+
+
 @dataclass(frozen=True)
 class AnnotatedValue(Value):
     """Value representing a `PEP 593 <https://www.python.org/dev/peps/pep-0593/>`_ Annotated object.
@@ -1874,6 +1884,26 @@ def annotate_value(origin: Value, metadata: Sequence[Union[Value, Extension]]) -
     return AnnotatedValue(origin, metadata)
 
 
+def unannotate_value(
+    origin: Value, extension: Type[Extension]
+) -> Tuple[Value, Sequence[Extension]]:
+    if not isinstance(origin, AnnotatedValue):
+        return origin, []
+    matches = [
+        metadata for metadata in origin.metadata if isinstance(metadata, extension)
+    ]
+    # the all_of_type call is redundant but necessary for pyanalyze's narrower for now
+    # TODO remove it
+    if matches and all_of_type(matches, Extension):
+        remaining = [
+            metadata
+            for metadata in origin.metadata
+            if not isinstance(metadata, extension)
+        ]
+        return annotate_value(origin.value, remaining), matches
+    return origin, []
+
+
 def unite_and_simplify(*values: Value, limit: int) -> Value:
     united = unite_values(*values)
     if not isinstance(united, MultiValuedValue) or len(united.vals) < limit:
@@ -1955,6 +1985,7 @@ def concrete_values_from_iterable(
 
     """
     value = replace_known_sequence_value(value)
+    is_nonempty = False
     if isinstance(value, MultiValuedValue):
         subvals = [concrete_values_from_iterable(val, ctx) for val in value.vals]
         errors = [subval for subval in subvals if isinstance(subval, CanAssignError)]
@@ -1981,20 +2012,26 @@ def concrete_values_from_iterable(
         if all(pair.is_required and not pair.is_many for pair in value.kv_pairs):
             return [pair.key for pair in value.kv_pairs]
     elif isinstance(value, KnownValue):
-        if (
-            isinstance(value.val, (str, bytes, range))
-            and len(value.val) < ITERATION_LIMIT
-        ):
-            return [KnownValue(c) for c in value.val]
+        if isinstance(value.val, (str, bytes, range)):
+            if len(value.val) < ITERATION_LIMIT:
+                return [KnownValue(c) for c in value.val]
+            is_nonempty = True
     elif value is NO_RETURN_VALUE:
         return NO_RETURN_VALUE
     iter_tv_map = IterableValue.can_assign(value, ctx)
     if not isinstance(iter_tv_map, CanAssignError):
-        return iter_tv_map.get(T, AnyValue(AnySource.generic_argument))
-    getitem_tv_map = GetItemProtoValue.can_assign(value, ctx)
-    if not isinstance(getitem_tv_map, CanAssignError):
-        return getitem_tv_map.get(T, AnyValue(AnySource.generic_argument))
-    return iter_tv_map
+        val = iter_tv_map.get(T, AnyValue(AnySource.generic_argument))
+    else:
+        getitem_tv_map = GetItemProtoValue.can_assign(value, ctx)
+        if not isinstance(getitem_tv_map, CanAssignError):
+            val = getitem_tv_map.get(T, AnyValue(AnySource.generic_argument))
+        else:
+            # We return the error from the __iter__ check because the __getitem__
+            # check is more arcane.
+            return iter_tv_map
+    if is_nonempty:
+        return annotate_value(val, [AlwaysPresentExtension()])
+    return val
 
 
 K = TypeVar("K")


Test patch:
diff --git a/pyanalyze/test_stacked_scopes.py b/pyanalyze/test_stacked_scopes.py
index e34c3ad..2be69db 100644
--- a/pyanalyze/test_stacked_scopes.py
+++ b/pyanalyze/test_stacked_scopes.py
@@ -427,6 +427,13 @@ class TestLoops(TestNameCheckVisitorBase):
                 assert_is_value(x, KnownValue(4))
             assert_is_value(x, KnownValue(4))
 
+        def huge_range():
+            for i in range(1000000000):
+                # We don't create the whole Union
+                assert_is_value(i, TypedValue(int))
+            # But we do recognize that the iterable is nonempty
+            assert_is_value(i, TypedValue(int))
+
     @assert_passes()
     def test_range_always_entered(self):
         def capybara():