Repository: quora/pyanalyze
Pull request #436 (2022-01-23 02:59:41)

add reveal_locals()


Patch:
diff --git a/docs/changelog.md b/docs/changelog.md
index 1c6d81b..a7a23f2 100644
--- a/docs/changelog.md
+++ b/docs/changelog.md
@@ -2,6 +2,7 @@
 
 ## Unreleased
 
+- Add support for `reveal_locals()` (#436)
 - Add support for `assert_error()` (#435)
 - Add support for `assert_type()` (#434)
 - `reveal_type()` and `dump_value()` now return their argument,
diff --git a/pyanalyze/extensions.py b/pyanalyze/extensions.py
index 7c438a6..3982b3b 100644
--- a/pyanalyze/extensions.py
+++ b/pyanalyze/extensions.py
@@ -394,6 +394,18 @@ def reveal_type(value: _T) -> _T:
     return value
 
 
+def reveal_locals() -> None:
+    """Reveal the types of all local variables.
+
+    When the type checker encounters a call to this function,
+    it prints the type of all variables in the local scope.
+
+    This does nothing at runtime.
+
+    """
+    pass
+
+
 def assert_type(val: _T, typ: Any) -> _T:
     """Assert the inferred static type of an expression.
 
diff --git a/pyanalyze/implementation.py b/pyanalyze/implementation.py
index 6661be2..b754a50 100644
--- a/pyanalyze/implementation.py
+++ b/pyanalyze/implementation.py
@@ -1,6 +1,6 @@
 from .annotations import type_from_value
 from .error_code import ErrorCode
-from .extensions import assert_type, reveal_type
+from .extensions import assert_type, reveal_locals, reveal_type
 from .format_strings import parse_format_string
 from .predicates import IsAssignablePredicate
 from .safe import safe_hasattr, safe_isinstance, safe_issubclass
@@ -948,6 +948,24 @@ def _reveal_type_impl(ctx: CallContext) -> Value:
     return value
 
 
+def _reveal_locals_impl(ctx: CallContext) -> Value:
+    scope = ctx.visitor.scopes.current_scope()
+    if ctx.visitor._is_collecting():
+        for varname in scope.all_variables():
+            scope.get(varname, ctx.node, ctx.visitor.state)
+    else:
+        details = []
+        for varname in scope.all_variables():
+            val, _, _ = scope.get(varname, ctx.node, ctx.visitor.state)
+            details.append(CanAssignError(f"{varname}: {val}"))
+        ctx.show_error(
+            "Revealed local types are:",
+            ErrorCode.inference_failure,
+            detail=str(CanAssignError(children=details)),
+        )
+    return KnownValue(None)
+
+
 def _dump_value_impl(ctx: CallContext) -> Value:
     value = ctx.vars["value"]
     if ctx.visitor._is_checking():
@@ -1155,6 +1173,9 @@ def get_default_argspecs() -> Dict[object, Signature]:
             impl=_reveal_type_impl,
             callable=reveal_type,
         ),
+        Signature.make(
+            [], KnownValue(None), impl=_reveal_locals_impl, callable=reveal_locals
+        ),
         Signature.make(
             [SigParameter("value", _POS_ONLY, annotation=TypeVarValue(T))],
             TypeVarValue(T),
@@ -1608,4 +1629,17 @@ def get_default_argspecs() -> Dict[object, Signature]:
                 impl=_assert_type_impl,
             )
             signatures.append(sig)
+        # Anticipating that this will be added to the stdlib
+        try:
+            reveal_locals_func = getattr(mod, "reveal_locals")
+        except AttributeError:
+            pass
+        else:
+            sig = Signature.make(
+                [],
+                KnownValue(None),
+                callable=reveal_locals_func,
+                impl=_reveal_locals_impl,
+            )
+            signatures.append(sig)
     return {sig.callable: sig for sig in signatures}
diff --git a/pyanalyze/stacked_scopes.py b/pyanalyze/stacked_scopes.py
index 3a5a652..295c4f8 100644
--- a/pyanalyze/stacked_scopes.py
+++ b/pyanalyze/stacked_scopes.py
@@ -121,6 +121,15 @@ class CompositeVariable:
     def extend_with(self, index: CompositeIndex) -> "CompositeVariable":
         return CompositeVariable(self.varname, (*self.attributes, index))
 
+    def __str__(self) -> str:
+        pieces = [self.varname]
+        for attr in self.attributes:
+            if isinstance(attr, str):
+                pieces.append(f".{attr}")
+            else:
+                pieces.append(f"[{attr.val!r}]")
+        return "".join(pieces)
+
 
 Varname = Union[str, CompositeVariable]
 
@@ -746,6 +755,9 @@ class Scope:
     def items(self) -> Iterable[Tuple[Varname, Value]]:
         return self.variables.items()
 
+    def all_variables(self) -> Iterable[Varname]:
+        return self.variables
+
     def __contains__(self, varname: Varname) -> bool:
         return varname in self.variables
 
@@ -1235,6 +1247,9 @@ class FunctionScope(Scope):
     def items(self) -> Iterable[Tuple[Varname, Value]]:
         raise NotImplementedError
 
+    def all_variables(self) -> Iterable[Varname]:
+        yield from self.name_to_current_definition_nodes
+
     def __contains__(self, varname: Varname) -> bool:
         return varname in self.name_to_all_definition_nodes
 


Test patch:
diff --git a/pyanalyze/test_inference_helpers.py b/pyanalyze/test_inference_helpers.py
index ad11c16..52cfd79 100644
--- a/pyanalyze/test_inference_helpers.py
+++ b/pyanalyze/test_inference_helpers.py
@@ -53,3 +53,15 @@ class TestAssertError(TestNameCheckVisitorBase):
 
             with assert_error():  # E: inference_failure
                 f(1)
+
+
+class TestRevealLocals(TestNameCheckVisitorBase):
+    @assert_passes()
+    def test(self) -> None:
+        from pyanalyze.extensions import reveal_locals
+
+        def capybara(a: object, b: str) -> None:
+            c = 3
+            if b == "x":
+                reveal_locals()  # E: inference_failure
+            print(a, b, c)