Repository: quora/pyanalyze
Pull request #634 (2023-05-26 23:08:29)

Fix crash on stubs that contain dict or set literals


Patch:
diff --git a/docs/changelog.md b/docs/changelog.md
index 646afc0..77e4fd7 100644
--- a/docs/changelog.md
+++ b/docs/changelog.md
@@ -2,6 +2,7 @@
 
 ## Unreleased
 
+- Fix crash on stubs that contain dict or set literals (#634)
 - Remove more old special cases and improve robustness of
   annotation parsing (#630)
 - Remove dependency on `typing_inspect` (#629)
diff --git a/pyanalyze/annotations.py b/pyanalyze/annotations.py
index b71f93b..069cfbf 100644
--- a/pyanalyze/annotations.py
+++ b/pyanalyze/annotations.py
@@ -23,6 +23,7 @@ These functions all use :class:`Context` objects to resolve names and
 show errors.
 
 """
+from _ast import Dict
 import ast
 import builtins
 import contextlib
@@ -76,6 +77,8 @@ from .signature import (
 )
 from .value import (
     _HashableValue,
+    DictIncompleteValue,
+    KVPair,
     annotate_value,
     AnnotatedValue,
     AnySource,
@@ -916,6 +919,24 @@ class _Visitor(ast.NodeVisitor):
         elts = [(False, self.visit(elt)) for elt in node.elts]
         return SequenceValue(list, elts)
 
+    def visit_Set(self, node: ast.Set) -> Value:
+        elts = [(False, self.visit(elt)) for elt in node.elts]
+        return SequenceValue(set, elts)
+
+    def visit_Dict(self, node: Dict) -> Any:
+        keys = [self.visit(key) if key is not None else None for key in node.keys]
+        values = [self.visit(value) for value in node.values]
+        kvpairs = []
+        for key, value in zip(keys, values):
+            if key is None:
+                # Just skip ** unpacking in stubs for now.
+                kvpairs.append(
+                    KVPair(AnyValue(AnySource.inference), AnyValue(AnySource.inference))
+                )
+            else:
+                kvpairs.append(KVPair(key, value))
+        return DictIncompleteValue(dict, kvpairs)
+
     def visit_Index(self, node: ast.Index) -> Value:
         # class is unused in 3.9
         return self.visit(node.value)  # static analysis: ignore[undefined_attribute]


Test patch:
diff --git a/pyanalyze/stubs/_pyanalyze_tests-stubs/defaults.pyi b/pyanalyze/stubs/_pyanalyze_tests-stubs/defaults.pyi
new file mode 100644
index 0000000..abc54ac
--- /dev/null
+++ b/pyanalyze/stubs/_pyanalyze_tests-stubs/defaults.pyi
@@ -0,0 +1,10 @@
+from typing import TypeVar
+
+T = TypeVar("T")
+U = TypeVar("U")
+V = TypeVar("V")
+W = TypeVar("W")
+
+def many_defaults(
+    a: T = {"a": 1}, b: U = [1, ()], c: V = (1, 2), d: W = {1, 2}
+) -> tuple[T, U, V, W]: ...
diff --git a/pyanalyze/test_typeshed.py b/pyanalyze/test_typeshed.py
index 8c3ca58..f6d5a91 100644
--- a/pyanalyze/test_typeshed.py
+++ b/pyanalyze/test_typeshed.py
@@ -29,6 +29,7 @@ from .typeshed import TypeshedFinder
 from .value import (
     AnySource,
     AnyValue,
+    SequenceValue,
     assert_is_value,
     CallableValue,
     DictIncompleteValue,
@@ -373,6 +374,29 @@ class TestBundledStubs(TestNameCheckVisitorBase):
 
             assert_type(x, Literal[3])
 
+    @assert_passes()
+    def test_stub_defaults(self):
+        def capybara():
+            from _pyanalyze_tests.defaults import many_defaults
+
+            a, b, c, d = many_defaults()
+            assert_is_value(
+                a, DictIncompleteValue(dict, [KVPair(KnownValue("a"), KnownValue(1))])
+            )
+            assert_is_value(
+                b,
+                SequenceValue(
+                    list, [(False, KnownValue(1)), (False, SequenceValue(tuple, []))]
+                ),
+            )
+            assert_is_value(
+                c,
+                SequenceValue(tuple, [(False, KnownValue(1)), (False, KnownValue(2))]),
+            )
+            assert_is_value(
+                d, SequenceValue(set, [(False, KnownValue(1)), (False, KnownValue(2))])
+            )
+
 
 class TestConstructors(TestNameCheckVisitorBase):
     @assert_passes()