Repository: quora/pyanalyze
Pull request #680 (2023-09-12 04:03:00)

Infer len constraints


Patch:
diff --git a/docs/changelog.md b/docs/changelog.md
index 12a9e23..cd4b471 100644
--- a/docs/changelog.md
+++ b/docs/changelog.md
@@ -2,6 +2,10 @@
 
 ## Unreleased
 
+- Fix usage of `assert_type()` with `Any` and with unions of
+  `Annotated` objects (#680)
+- Support inferring `MinLen` and `MaxLen` annotations based
+  on `len()` checks (#680)
 - Expose a convenience API for runtime type checking in the
   `pyanalyze.runtime` module (#674)
 - Support for annotations from the `annotated-types` library (#673)
diff --git a/pyanalyze/implementation.py b/pyanalyze/implementation.py
index c744ba5..316042a 100644
--- a/pyanalyze/implementation.py
+++ b/pyanalyze/implementation.py
@@ -1,9 +1,23 @@
+import ast
 import collections
 import collections.abc
 import inspect
 import typing
 from itertools import product
-from typing import Callable, cast, Dict, NewType, Optional, Sequence, TypeVar, Union
+from typing import (
+    Callable,
+    Type,
+    cast,
+    Dict,
+    NewType,
+    Optional,
+    Sequence,
+    TypeVar,
+    Union,
+)
+
+from .annotated_types import MaxLen
+from .annotated_types import MinLen
 
 import qcore
 import typing_extensions
@@ -37,6 +51,8 @@ from .value import (
     AnnotatedValue,
     AnySource,
     AnyValue,
+    CustomCheckExtension,
+    annotate_value,
     assert_is_value,
     CallableValue,
     CanAssignContext,
@@ -61,7 +77,6 @@ from .value import (
     TypedDictValue,
     TypedValue,
     TypeVarValue,
-    unannotate,
     UNINITIALIZED_VALUE,
     unite_values,
     unpack_values,
@@ -1266,11 +1281,27 @@ def _cast_impl(ctx: CallContext) -> Value:
     return type_from_value(typ, visitor=ctx.visitor, node=ctx.node)
 
 
+def _recursive_unanotate(val: Value) -> Value:
+    # Maybe we should also recurse into type parameters?
+    if isinstance(val, AnnotatedValue):
+        return _recursive_unanotate(val.value)
+    elif isinstance(val, MultiValuedValue):
+        return unite_values(*[_recursive_unanotate(subval) for subval in val.vals])
+    else:
+        return val
+
+
 def _assert_type_impl(ctx: CallContext) -> Value:
-    # TODO maybe we should walk over the whole value and remove Annotated.
-    val = unannotate(ctx.vars["val"])
+    val = _recursive_unanotate(ctx.vars["val"])
     typ = ctx.vars["typ"]
     expected_type = type_from_value(typ, visitor=ctx.visitor, node=ctx.node)
+    # Can't distinguish between different kinds of Any here
+    if (
+        isinstance(val, AnyValue)
+        and isinstance(expected_type, AnyValue)
+        and expected_type.source is not AnySource.error
+    ):
+        return val
     if val != expected_type:
         ctx.show_error(
             f"Type is {val} (expected {expected_type})",
@@ -1332,12 +1363,37 @@ def len_of_value(val: Value) -> Value:
     return TypedValue(int)
 
 
+def len_transformer(val: Value, op: Type[ast.AST], comparator: object) -> Value:
+    if not isinstance(comparator, int):
+        return val
+    if isinstance(len_of_value(val), KnownValue):
+        return val  # no need to specify
+    if op is ast.Eq:
+        return annotate_value(
+            val,
+            [
+                CustomCheckExtension(MinLen(comparator)),
+                CustomCheckExtension(MaxLen(comparator)),
+            ],
+        )
+    elif op is ast.Lt:
+        return annotate_value(val, [CustomCheckExtension(MaxLen(comparator - 1))])
+    elif op is ast.LtE:
+        return annotate_value(val, [CustomCheckExtension(MaxLen(comparator))])
+    elif op is ast.Gt:
+        return annotate_value(val, [CustomCheckExtension(MinLen(comparator + 1))])
+    elif op is ast.GtE:
+        return annotate_value(val, [CustomCheckExtension(MinLen(comparator))])
+    else:
+        return val
+
+
 def _len_impl(ctx: CallContext) -> ImplReturn:
     varname = ctx.varname_for_arg("obj")
     if varname is None:
         constraint = NULL_CONSTRAINT
     else:
-        constraint = PredicateProvider(varname, len_of_value)
+        constraint = PredicateProvider(varname, len_of_value, len_transformer)
     return ImplReturn(len_of_value(ctx.vars["obj"]), constraint)
 
 
diff --git a/pyanalyze/name_check_visitor.py b/pyanalyze/name_check_visitor.py
index bb45b51..cfc55a9 100644
--- a/pyanalyze/name_check_visitor.py
+++ b/pyanalyze/name_check_visitor.py
@@ -52,7 +52,7 @@ import qcore
 import typeshed_client
 from typing_extensions import Annotated, Protocol, get_args, get_origin
 
-from pyanalyze.annotated_types import Gt, Ge, Le, Lt
+from .annotated_types import Gt, Ge, Le, Lt
 
 from . import attributes, format_strings, importer, node_visitor, type_evaluation
 from .analysis_lib import get_attribute_path
@@ -307,6 +307,13 @@ COMPARATOR_TO_OPERATOR = {
     ast.In: (_in, _not_in, None),
     ast.NotIn: (_not_in, _in, None),
 }
+_NEG_OPERATOR_TO_AST = {
+    neg_op: node_cls for node_cls, (_, neg_op, _) in COMPARATOR_TO_OPERATOR.items()
+}
+AST_TO_REVERSE = {
+    node_cls: _NEG_OPERATOR_TO_AST[op]
+    for node_cls, (op, _, _) in COMPARATOR_TO_OPERATOR.items()
+}
 
 SAFE_DECORATORS_FOR_ARGSPEC_TO_RETVAL = [KnownValue(asynq.asynq), KnownValue(property)]
 if sys.version_info < (3, 11):
@@ -3365,14 +3372,22 @@ class NameCheckVisitor(node_visitor.ReplacingNodeVisitor):
         def predicate_func(value: Value, positive: bool) -> Optional[Value]:
             predicate_value = pred.provider(value)
             if isinstance(predicate_value, KnownValue):
-                op = positive_operator if positive else negative_operator
+                operator = positive_operator if positive else negative_operator
                 try:
-                    result = op(predicate_value.val, other_val)
+                    result = operator(predicate_value.val, other_val)
                 except Exception:
                     pass
                 else:
                     if not result:
                         return None
+            if pred.value_transformer is not None:
+                op_cls = type(op)
+                if not positive:
+                    if op_cls in AST_TO_REVERSE:
+                        op_cls = AST_TO_REVERSE[op_cls]
+                    else:
+                        return value
+                return pred.value_transformer(value, op_cls, other_val)
             return value
 
         return Constraint(pred.varname, ConstraintType.predicate, True, predicate_func)
diff --git a/pyanalyze/stacked_scopes.py b/pyanalyze/stacked_scopes.py
index 5808a47..b461db9 100644
--- a/pyanalyze/stacked_scopes.py
+++ b/pyanalyze/stacked_scopes.py
@@ -41,6 +41,7 @@ from typing import (
     Sequence,
     Set,
     Tuple,
+    Type,
     TypeVar,
     Union,
 )
@@ -510,6 +511,7 @@ class PredicateProvider(AbstractConstraint):
 
     varname: VarnameWithOrigin
     provider: Callable[[Value], Value]
+    value_transformer: Optional[Callable[[Value, Type[AST], object], Value]] = None
 
     def apply(self) -> Iterable[Constraint]:
         return []


Test patch:
diff --git a/pyanalyze/test_annotated_types.py b/pyanalyze/test_annotated_types.py
index cb14b31..54485f5 100644
--- a/pyanalyze/test_annotated_types.py
+++ b/pyanalyze/test_annotated_types.py
@@ -279,3 +279,40 @@ class TestInferAnnotations(TestNameCheckVisitorBase):
 
             if i > 6:
                 takes_gt_5(i)
+
+    @assert_passes()
+    def test_len(self):
+        from typing_extensions import Annotated
+        from annotated_types import MinLen, MaxLen, Len
+
+        def takes_len_5(x: Annotated[str, Len(5)]) -> None:
+            pass
+
+        def takes_min_len_5(x: Annotated[str, MinLen(5)]) -> None:
+            pass
+
+        def takes_max_len_5(x: Annotated[str, MaxLen(5)]) -> None:
+            pass
+
+        def capybara(s: str) -> None:
+            if len(s) == 5:
+                takes_len_5(s)
+                takes_min_len_5(s)
+                takes_max_len_5(s)
+
+            if len(s) >= 6:
+                takes_min_len_5(s)
+                takes_max_len_5(s)  # E: incompatible_argument
+
+            if len(s) <= 4:
+                takes_max_len_5(s)
+                takes_min_len_5(s)  # E: incompatible_argument
+
+            if len(s) > 5:
+                takes_min_len_5(s)
+                takes_max_len_5(s)  # E: incompatible_argument
+
+            if len(s) < 5:
+                takes_max_len_5(s)
+                takes_min_len_5(s)  # E: incompatible_argument
+                takes_len_5(s)  # E: incompatible_argument
diff --git a/pyanalyze/test_implementation.py b/pyanalyze/test_implementation.py
index 4bcce17..8de90a7 100644
--- a/pyanalyze/test_implementation.py
+++ b/pyanalyze/test_implementation.py
@@ -1327,3 +1327,21 @@ class TestCallableGuards(TestNameCheckVisitorBase):
             assert_is_value(o, TypedValue(object))
             if inspect.isfunction(o):
                 assert_is_value(o, TypedValue(FunctionType))
+
+
+class TestAssertType(TestNameCheckVisitorBase):
+    @assert_passes()
+    def test_assert_type(self):
+        from typing_extensions import assert_type, Literal
+        from typing import Any
+
+        def capybara(x: object, unannotated, explicit_any: Any) -> None:
+            assert_type(x, object)
+
+            val = assert_type(1, Literal[1])
+            assert_type(val, Literal[1])
+
+            assert_type(unannotated, Any)
+            assert_type(explicit_any, Any)
+
+            assert_type(x, int)  # E: inference_failure
diff --git a/pyanalyze/test_stacked_scopes.py b/pyanalyze/test_stacked_scopes.py
index 5e4aaa1..df99da1 100644
--- a/pyanalyze/test_stacked_scopes.py
+++ b/pyanalyze/test_stacked_scopes.py
@@ -1786,13 +1786,15 @@ class TestInvalidation(TestNameCheckVisitorBase):
 
     @assert_passes()
     def test_len_condition(self) -> None:
+        from typing_extensions import Any, assert_type
+
         def capybara(file_list, key, ids):
             has_bias = len(key) > 0
             data = []
             for _ in file_list:
-                assert_is_value(key, AnyValue(AnySource.unannotated))
+                assert_type(key, Any)
                 if has_bias:
-                    assert_is_value(key, AnyValue(AnySource.unannotated))
+                    assert_type(key, Any)
                     data = [ids, data[key]]
                 else:
                     data = [ids]