Repository: quora/pyanalyze
Pull request #299 (2021-12-03 01:52:18)

Minor typing fixes
Found by #297


Patch:
diff --git a/docs/changelog.md b/docs/changelog.md
index 6c576e7..bca28a5 100644
--- a/docs/changelog.md
+++ b/docs/changelog.md
@@ -2,6 +2,7 @@
 
 ## Unreleased
 
+- Fix handling of `async for` (#298)
 - Add support for stub-only Protocols (#295)
 - Basic support for stub-only types (#290)
 - Require `typing_inspect>=0.7.0` (#290)
diff --git a/pyanalyze/importer.py b/pyanalyze/importer.py
index 6af58c0..05aa1ae 100644
--- a/pyanalyze/importer.py
+++ b/pyanalyze/importer.py
@@ -5,6 +5,7 @@ Module responsible for importing files.
 """
 
 from functools import lru_cache
+import importlib
 import importlib.util
 from pathlib import Path
 import sys
@@ -99,7 +100,7 @@ def load_module_from_file(
     return import_module(str(abspath), abspath), False
 
 
-def import_module(module_path: str, filename: str) -> ModuleType:
+def import_module(module_path: str, filename: Path) -> ModuleType:
     """Import a file under an arbitrary module name."""
     spec = importlib.util.spec_from_file_location(module_path, filename)
     module = importlib.util.module_from_spec(spec)
diff --git a/pyanalyze/name_check_visitor.py b/pyanalyze/name_check_visitor.py
index ca7e500..3240c20 100644
--- a/pyanalyze/name_check_visitor.py
+++ b/pyanalyze/name_check_visitor.py
@@ -15,7 +15,6 @@ from ast_decompiler import decompile
 import asyncio
 import builtins
 import collections.abc
-from collections.abc import Awaitable
 import contextlib
 from dataclasses import dataclass
 from functools import reduce
@@ -1347,7 +1346,7 @@ class NameCheckVisitor(
             is_coroutine = _is_coroutine_function(potential_function)
 
         if is_coroutine or info.is_decorated_coroutine:
-            return_value = GenericValue(Awaitable, [return_value])
+            return_value = GenericValue(collections.abc.Awaitable, [return_value])
 
         try:
             argspec = self.arg_spec_cache.get_argspec(potential_function)
@@ -2766,9 +2765,12 @@ class NameCheckVisitor(
 
     def visit_Await(self, node: ast.Await) -> Value:
         composite = self.composite_from_node(node.value)
-        tv_map = AwaitableValue.can_assign(composite[0], self)
+        return self.unpack_awaitable(composite, node.value)
+
+    def unpack_awaitable(self, composite: Composite, node: ast.AST) -> Value:
+        tv_map = AwaitableValue.can_assign(composite.value, self)
         if isinstance(tv_map, CanAssignError):
-            return self._check_dunder_call(node.value, composite, "__await__", [])
+            return self._check_dunder_call(node, composite, "__await__", [])
         else:
             return tv_map.get(T, AnyValue(AnySource.generic_argument))
 
@@ -2777,7 +2779,7 @@ class NameCheckVisitor(
         value = self.visit(node.value)
         if not TypedValue(collections.abc.Iterable).is_assignable(
             value, self
-        ) and not TypedValue(Awaitable).is_assignable(value, self):
+        ) and not AwaitableValue.is_assignable(value, self):
             self._show_error_if_checking(
                 node,
                 f"Cannot use {value} in yield from",
@@ -3076,9 +3078,10 @@ class NameCheckVisitor(
         composite = self.composite_from_node(node)
         if is_async:
             iterator = self._check_dunder_call(node, composite, "__aiter__", [])
-            return self._check_dunder_call(
+            anext = self._check_dunder_call(
                 node, Composite(iterator, None, node), "__anext__", []
             )
+            return self.unpack_awaitable(Composite(anext), node)
         iterated = composite.value
         result = concrete_values_from_iterable(iterated, self)
         if isinstance(result, CanAssignError):
@@ -3252,7 +3255,7 @@ class NameCheckVisitor(
             )
         # If the value is an awaitable or is assignable to asyncio.Future, show
         # an error about a missing await.
-        elif value.is_type(Awaitable) or (value.is_type(asyncio.Future)):
+        elif value.is_type(collections.abc.Awaitable) or value.is_type(asyncio.Future):
             if self.is_async_def:
                 new_node = ast.Expr(value=ast.Await(value=node.value))
             else:
diff --git a/pyanalyze/node_visitor.py b/pyanalyze/node_visitor.py
index 38958be..a4aec9a 100644
--- a/pyanalyze/node_visitor.py
+++ b/pyanalyze/node_visitor.py
@@ -149,7 +149,7 @@ class BaseNodeVisitor(ast.NodeVisitor):
 
     _changes_for_fixer = collections.defaultdict(list)
 
-    tree: Optional[ast.Module]
+    tree: ast.Module
     all_failures: List[Failure]
 
     def __init__(
diff --git a/pyanalyze/typeshed.py b/pyanalyze/typeshed.py
index 0e5fb09..fda010d 100644
--- a/pyanalyze/typeshed.py
+++ b/pyanalyze/typeshed.py
@@ -33,13 +33,14 @@ from collections.abc import (
     Sized,
 )
 from contextlib import AbstractContextManager
-from dataclasses import dataclass
+from dataclasses import dataclass, field
 import collections.abc
 import qcore
 import inspect
 import sys
 from types import GeneratorType
 from typing import (
+    Dict,
     Set,
     Tuple,
     cast,
@@ -93,12 +94,16 @@ _TYPING_ALIASES = {
 }
 
 
-class TypeshedFinder(object):
-    def __init__(self, verbose: bool = True) -> None:
-        self.verbose = verbose
-        self.resolver = typeshed_client.Resolver()
-        self._assignment_cache = {}
-        self._attribute_cache = {}
+@dataclass
+class TypeshedFinder:
+    verbose: bool = True
+    resolver: typeshed_client.Resolver = field(default_factory=typeshed_client.Resolver)
+    _assignment_cache: Dict[Tuple[str, ast3.AST], Value] = field(
+        default_factory=dict, repr=False, init=False
+    )
+    _attribute_cache: Dict[Tuple[str, str, bool], Value] = field(
+        default_factory=dict, repr=False, init=False
+    )
 
     def log(self, message: str, obj: object) -> None:
         if not self.verbose:


Test patch:
diff --git a/pyanalyze/test_name_check_visitor.py b/pyanalyze/test_name_check_visitor.py
index 6c79380..4ee6763 100644
--- a/pyanalyze/test_name_check_visitor.py
+++ b/pyanalyze/test_name_check_visitor.py
@@ -1849,7 +1849,7 @@ class TestAsyncAwait(TestNameCheckVisitorBase):
     @assert_passes()
     def test_async_comprehension(self):
         class ANext:
-            def __anext__(self) -> int:
+            async def __anext__(self) -> int:
                 return 42
 
         class AIter:
@@ -1860,10 +1860,28 @@ class TestAsyncAwait(TestNameCheckVisitorBase):
             x = [y async for y in AIter()]
             assert_is_value(x, make_weak(GenericValue(list, [TypedValue(int)])))
 
-    @assert_fails(ErrorCode.unsupported_operation)
+    @assert_passes()
+    def test_async_generator(self):
+        import collections.abc
+        from typing import AsyncIterator
+
+        async def f() -> AsyncIterator[int]:
+            yield 1
+            yield 2
+
+        async def capybara():
+            x = f()
+            assert_is_value(
+                x, GenericValue(collections.abc.AsyncIterator, [TypedValue(int)])
+            )
+            ints = [i async for i in x]
+            # TODO should be list[int] but we lose the type argument somewhere
+            assert_is_value(ints, TypedValue(list))
+
+    @assert_passes()
     def test_bad_async_comprehension(self):
         async def f():
-            return [x async for x in []]
+            return [x async for x in []]  # E: unsupported_operation
 
 
 class TestKeywordOnlyArguments(TestNameCheckVisitorBase):