Return generator annotation when it exists

Summary:
If a function already has a return annotation that's a generator, we return it.

The next step to consider is:
what if we have a generator annotation that is different than the one we inferred? this is a todo. For now, we return the annotation directly without considering the type we inferred.

Other todos include
- send type
- next keyword
- async generators

Reviewed By: ndmitchell

Differential Revision: D67732875

fbshipit-source-id: c3f0df46acb3cdcdbee1a58f159e71d385450a6d

pyre2/pyre2/bin/alt/answers.rs

@@ -1257,7 +1257,8 @@ impl<'a, Ans: LookupAnswer> AnswersSolver<'a, Ans> {
                 let yield_expr_type = self
                     .get(&Key::YieldType(ShortIdentifier::new(&x.name)))
                     .arc_clone();
-                let ret = if yield_expr_type.is_never() {
+                // # TODO zeina: raise a type error if return type is inconsistent with inferred generator type
+                let ret = if yield_expr_type.is_never() || ret.is_generator() {
                     ret
                 } else {
                     self.stdlib

pyre2/pyre2/bin/test/yield.rs

@@ -7,6 +7,8 @@
 
 use crate::testcase_with_bug;
 
+// TODO zeina: use assert_type instead of reveal_type after I support most of these cases.
+
 // TODO zeina: 1- We need a generator type; 2- next keyword currently unsupported
 testcase_with_bug!(
     test_generator,
@@ -54,7 +56,7 @@ def accumulate(x: int) -> Generator[int, int, None]:
     yield x # E: TODO: ExprYield - Answers::expr_infer # E:  EXPECTED None <: Generator[int, int, None]
 
 gen = accumulate(10)
-reveal_type(gen) # E: revealed type: Generator[int, Unknown, Generator[int, int, None]]
+reveal_type(gen) # E: revealed type: Generator[int, int, None]
 gen.send(5) # E:  Object of class `Generator` has no attribute `send`
 
 "#,
@@ -108,7 +110,7 @@ def f(value) -> Generator[int, None, None]:
     while True: # E: EXPECTED None <: Generator[int, None, None]
         yield value # E: TODO: ExprYield - Answers::expr_infer
 
-reveal_type(f(3)) # E: revealed type: Generator[Unknown, Unknown, Generator[int, None, None]]
+reveal_type(f(3)) # E: revealed type: Generator[int, None, None]
 
 "#,
 );
@@ -125,12 +127,12 @@ def f(value: T) -> Generator[T, None, None]:
     while True: # E: EXPECTED None <: Generator[?_, None, None]
         yield value # E: TODO: ExprYield - Answers::expr_infer
 
-reveal_type(f(3)) # E: revealed type: Generator[int, Unknown, Generator[int, None, None]]
+reveal_type(f(3)) # E: revealed type: Generator[int, None, None]
 
 "#,
 );
 
-// TODO zeina: This should typecheck
+// TODO zeina: This should typecheck; we should first support async generators.
 testcase_with_bug!(
     test_async_generator_basic_type,
     r#"
@@ -139,12 +141,12 @@ from typing import AsyncGenerator, reveal_type # E: Could not import `AsyncGener
 async def async_count_up_to() -> AsyncGenerator[int, None]:
     yield 2 # E: TODO: ExprYield - Answers::expr_infer
 
-reveal_type(async_count_up_to()) # E: Coroutine[Unknown, Unknown, Generator[Literal[2], Unknown, Error]]
+reveal_type(async_count_up_to()) # E: Generator[Literal[2], Unknown, Error]
 
 "#,
 );
 
-// TODO zeina: This should typecheck
+// TODO zeina: This should typecheck; we should first support async generators.
 testcase_with_bug!(
     test_async_generator_basic_inference,
     r#"
@@ -153,7 +155,7 @@ from typing import reveal_type
 async def async_count_up_to():
     yield 2 # E: TODO: ExprYield - Answers::expr_infer
 
-reveal_type(async_count_up_to()) # E: Coroutine[Unknown, Unknown, Generator[Literal[2], Unknown, None]]
+reveal_type(async_count_up_to()) # E: Generator[Literal[2], Unknown, None]
 
 "#,
 );

pyre2/pyre2/bin/types/types.rs

@@ -399,6 +399,20 @@ impl Type {
         }
     }
 
+    pub fn is_generator(&self) -> bool {
+        match self {
+            Type::ClassType(cls) => {
+                let typ = cls.class_object();
+                let name = typ.name().id().as_str();
+                if name == "Generator" {
+                    return true;
+                }
+                false
+            }
+            _ => false,
+        }
+    }
+
     pub fn is_var(&self) -> bool {
         matches!(self, Type::Var(_))
     }